vortex-array 0.78.0

Vortex in memory columnar data format
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

//! Native comparison of boolean arrays using word-wise bit operations.

use vortex_buffer::BitBuffer;
use vortex_error::VortexResult;
use vortex_error::vortex_err;

use crate::ArrayRef;
use crate::ExecutionCtx;
use crate::IntoArray;
use crate::arrays::BoolArray;
use crate::arrays::Constant;
use crate::dtype::Nullability;
use crate::scalar_fn::fns::binary::compare::compare_validity;
use crate::scalar_fn::fns::operators::CompareOperator;
use crate::validity::Validity;

enum BoolOperand {
    Array { bits: BitBuffer, validity: Validity },
    Constant { value: bool, validity: Validity },
}

impl BoolOperand {
    fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
        if let Some(constant) = array.as_opt::<Constant>() {
            let value = constant
                .scalar()
                .as_bool_opt()
                .ok_or_else(|| vortex_err!("expected boolean scalar"))?
                .value()
                .ok_or_else(|| vortex_err!("null constant handled by execute_compare"))?;
            return Ok(Self::Constant {
                value,
                validity: if constant.scalar().dtype().is_nullable() {
                    Validity::AllValid
                } else {
                    Validity::NonNullable
                },
            });
        }

        let array = array.clone().execute::<BoolArray>(ctx)?;
        let validity = array.validity()?;
        Ok(Self::Array {
            bits: array.into_bit_buffer(),
            validity,
        })
    }

    fn validity(&self) -> Validity {
        match self {
            Self::Array { validity, .. } | Self::Constant { validity, .. } => validity.clone(),
        }
    }
}

/// Compare two boolean arrays.
///
/// Values compare as `false < true`; every operator reduces to at most two word-wise passes over
/// the value bit buffers.
pub(super) fn compare_bool(
    lhs: &ArrayRef,
    rhs: &ArrayRef,
    op: CompareOperator,
    nullability: Nullability,
    ctx: &mut ExecutionCtx,
) -> VortexResult<ArrayRef> {
    let len = lhs.len();
    let lhs = BoolOperand::try_new(lhs, ctx)?;
    let rhs = BoolOperand::try_new(rhs, ctx)?;
    let validity = compare_validity(lhs.validity(), rhs.validity(), nullability)?;

    let bits = match (lhs, rhs) {
        (BoolOperand::Array { bits: l, .. }, BoolOperand::Array { bits: r, .. }) => {
            compare_bits(l, r, op)
        }
        (BoolOperand::Array { bits, .. }, BoolOperand::Constant { value, .. }) => {
            compare_bits_constant(bits, value, op)
        }
        (BoolOperand::Constant { value, .. }, BoolOperand::Array { bits, .. }) => {
            compare_bits_constant(bits, value, op.swap())
        }
        (BoolOperand::Constant { value: l, .. }, BoolOperand::Constant { value: r, .. }) => {
            // Unreachable through `execute_compare` (constant-constant is folded there), but
            // cheap to answer anyway.
            let result = super::ordering_predicate(op)(l.cmp(&r));
            BitBuffer::full(result, len)
        }
    };

    Ok(BoolArray::try_new(bits, validity)?.into_array())
}

fn compare_bits(lhs: BitBuffer, rhs: BitBuffer, op: CompareOperator) -> BitBuffer {
    match op {
        CompareOperator::Eq => !(lhs ^ &rhs),
        CompareOperator::NotEq => lhs ^ &rhs,
        // a < b  ⟺  !a & b
        CompareOperator::Lt => rhs.into_bitand_not(&lhs),
        // a <= b  ⟺  !(a & !b)
        CompareOperator::Lte => !lhs.into_bitand_not(&rhs),
        // a > b  ⟺  a & !b
        CompareOperator::Gt => lhs.into_bitand_not(&rhs),
        // a >= b  ⟺  !(!a & b)
        CompareOperator::Gte => !rhs.into_bitand_not(&lhs),
    }
}

/// Compare array bits against a non-null constant: `bits <op> value`.
fn compare_bits_constant(bits: BitBuffer, value: bool, op: CompareOperator) -> BitBuffer {
    let len = bits.len();
    match (op, value) {
        (CompareOperator::Eq, true)
        | (CompareOperator::NotEq, false)
        | (CompareOperator::Gt, false)
        | (CompareOperator::Gte, true) => bits,
        (CompareOperator::Eq, false)
        | (CompareOperator::NotEq, true)
        | (CompareOperator::Lt, true)
        | (CompareOperator::Lte, false) => !bits,
        (CompareOperator::Lt, false) | (CompareOperator::Gt, true) => BitBuffer::new_unset(len),
        (CompareOperator::Lte, true) | (CompareOperator::Gte, false) => BitBuffer::new_set(len),
    }
}