vortex_sequence/compute/
compare.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_array::ArrayRef;
5use vortex_array::ArrayView;
6use vortex_array::ExecutionCtx;
7use vortex_array::IntoArray;
8use vortex_array::arrays::BoolArray;
9use vortex_array::arrays::ConstantArray;
10use vortex_array::dtype::NativePType;
11use vortex_array::dtype::Nullability;
12use vortex_array::match_each_integer_ptype;
13use vortex_array::scalar::PValue;
14use vortex_array::scalar::Scalar;
15use vortex_array::scalar_fn::fns::binary::CompareKernel;
16use vortex_array::scalar_fn::fns::operators::CompareOperator;
17use vortex_buffer::BitBufferMut;
18use vortex_error::VortexExpect;
19use vortex_error::VortexResult;
20use vortex_error::vortex_bail;
21use vortex_error::vortex_err;
22
23use crate::array::Sequence;
24
25impl CompareKernel for Sequence {
26    fn compare(
27        lhs: ArrayView<'_, Self>,
28        rhs: &ArrayRef,
29        operator: CompareOperator,
30        _ctx: &mut ExecutionCtx,
31    ) -> VortexResult<Option<ArrayRef>> {
32        // TODO(joe): support other operators (NotEq, Lt, Lte, Gt, Gte) in encoded space.
33        if operator != CompareOperator::Eq {
34            return Ok(None);
35        }
36
37        let Some(constant) = rhs.as_constant() else {
38            return Ok(None);
39        };
40
41        // Check if there exists an integer solution to const = base + (0..len) * multiplier.
42        let set_idx = find_intersection_scalar(
43            lhs.base(),
44            lhs.multiplier(),
45            lhs.len(),
46            constant
47                .as_primitive()
48                .pvalue()
49                .vortex_expect("null constant handled in adaptor"),
50        );
51
52        let nullability = lhs.dtype().nullability() | rhs.dtype().nullability();
53        let validity = match nullability {
54            Nullability::NonNullable => vortex_array::validity::Validity::NonNullable,
55            Nullability::Nullable => vortex_array::validity::Validity::AllValid,
56        };
57
58        if let Ok(set_idx) = set_idx {
59            let mut buffer = BitBufferMut::new_unset(lhs.len());
60            buffer.set(set_idx);
61            let buffer = buffer.freeze();
62            Ok(Some(BoolArray::new(buffer, validity).into_array()))
63        } else {
64            Ok(Some(
65                ConstantArray::new(Scalar::bool(false, nullability), lhs.len()).into_array(),
66            ))
67        }
68    }
69}
70
71/// Find the index where `base + idx * multiplier == intercept`, if one exists.
72///
73/// # Errors
74/// Return `VortexError` if:
75/// - `len` is 0
76/// - `intercept` or `multiplier` can't be cast to `base`'s PType
77/// - `intercept` is outside the range of the sequence
78/// - `intercept` doesn't fall exactly on a sequence value
79pub(crate) fn find_intersection_scalar(
80    base: PValue,
81    multiplier: PValue,
82    len: usize,
83    intercept: PValue,
84) -> VortexResult<usize> {
85    match_each_integer_ptype!(base.ptype(), |P| {
86        let intercept = intercept.cast::<P>()?;
87        let base = base.cast::<P>()?;
88        let multiplier = multiplier.cast::<P>()?;
89        find_intersection(base, multiplier, len, intercept)
90    })
91}
92
93fn find_intersection<P: NativePType>(
94    base: P,
95    multiplier: P,
96    len: usize,
97    intercept: P,
98) -> VortexResult<usize> {
99    if len == 0 {
100        vortex_bail!("len == 0")
101    }
102
103    let count = P::from_usize(len - 1).vortex_expect("idx must fit into type");
104    let end_element = base + (multiplier * count);
105
106    // Handle ascending vs descending sequences
107    let (min_val, max_val) = if multiplier.is_ge(P::zero()) {
108        (base, end_element)
109    } else {
110        (end_element, base)
111    };
112
113    // Check if intercept is in range
114    if !intercept.is_ge(min_val) || !intercept.is_le(max_val) {
115        vortex_bail!("{intercept} is outside of ({min_val}, {max_val}) range")
116    }
117
118    // Handle zero multiplier (constant sequence)
119    if multiplier == P::zero() {
120        if intercept == base {
121            return Ok(0);
122        } else {
123            vortex_bail!("{intercept} != {base} with zero multiplier")
124        }
125    }
126
127    // Check if (intercept - base) is evenly divisible by multiplier
128    let diff = intercept - base;
129    if diff % multiplier != P::zero() {
130        vortex_bail!("{diff} % {multiplier} != 0")
131    }
132
133    let idx = diff / multiplier;
134    idx.to_usize()
135        .ok_or_else(|| vortex_err!("Cannot represent {idx} as usize"))
136}
137
138#[cfg(test)]
139mod tests {
140    use std::sync::LazyLock;
141
142    use vortex_array::IntoArray;
143    use vortex_array::VortexSessionExecute;
144    use vortex_array::arrays::BoolArray;
145    use vortex_array::arrays::ConstantArray;
146    use vortex_array::assert_arrays_eq;
147    use vortex_array::builtins::ArrayBuiltins;
148    use vortex_array::dtype::Nullability::NonNullable;
149    use vortex_array::dtype::Nullability::Nullable;
150    use vortex_array::scalar_fn::fns::operators::Operator;
151    use vortex_session::VortexSession;
152
153    use crate::Sequence;
154
155    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
156        let session = vortex_array::array_session();
157        crate::initialize(&session);
158        session
159    });
160
161    #[test]
162    fn test_compare_match() {
163        let lhs = Sequence::try_new_typed(2i64, 1, NonNullable, 4).unwrap();
164        let rhs = ConstantArray::new(4i64, lhs.len());
165        let result = lhs
166            .into_array()
167            .binary(rhs.into_array(), Operator::Eq)
168            .unwrap();
169        let expected = BoolArray::from_iter([false, false, true, false]);
170        assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx());
171    }
172
173    #[test]
174    fn test_compare_match_scale() {
175        let lhs = Sequence::try_new_typed(2i64, 3, Nullable, 4).unwrap();
176        let rhs = ConstantArray::new(8i64, lhs.len());
177        let result = lhs
178            .into_array()
179            .binary(rhs.into_array(), Operator::Eq)
180            .unwrap();
181        let expected = BoolArray::from_iter([Some(false), Some(false), Some(true), Some(false)]);
182        assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx());
183    }
184
185    #[test]
186    fn test_compare_no_match() {
187        let lhs = Sequence::try_new_typed(2i64, 1, NonNullable, 4).unwrap();
188        let rhs = ConstantArray::new(1i64, lhs.len());
189        let result = lhs
190            .into_array()
191            .binary(rhs.into_array(), Operator::Eq)
192            .unwrap();
193        let expected = BoolArray::from_iter([false, false, false, false]);
194        assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx());
195    }
196}
vortex_sequence/compute/compare.rs

vortex_sequence/compute/
compare.rs
