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::Nullability;
11use vortex_array::scalar::PValue;
12use vortex_array::scalar::Scalar;
13use vortex_array::scalar_fn::fns::binary::CompareKernel;
14use vortex_array::scalar_fn::fns::operators::CompareOperator;
15use vortex_buffer::BitBufferMut;
16use vortex_error::VortexExpect;
17use vortex_error::VortexResult;
18
19use crate::array::Sequence;
20use crate::eval;
21
22impl CompareKernel for Sequence {
23    fn compare(
24        lhs: ArrayView<'_, Self>,
25        rhs: &ArrayRef,
26        operator: CompareOperator,
27        _ctx: &mut ExecutionCtx,
28    ) -> VortexResult<Option<ArrayRef>> {
29        // TODO(joe): support other operators (NotEq, Lt, Lte, Gt, Gte) in encoded space.
30        if operator != CompareOperator::Eq {
31            return Ok(None);
32        }
33
34        let Some(constant) = rhs.as_constant() else {
35            return Ok(None);
36        };
37
38        let value = constant
39            .as_primitive()
40            .pvalue()
41            .vortex_expect("null constant handled in adaptor");
42
43        // Fall back for non-integer constants.
44        let Some(intersection) = find_intersection(lhs.base(), lhs.multiplier(), lhs.len(), value)
45        else {
46            return Ok(None);
47        };
48
49        let nullability = lhs.dtype().nullability() | rhs.dtype().nullability();
50        let validity = match nullability {
51            Nullability::NonNullable => vortex_array::validity::Validity::NonNullable,
52            Nullability::Nullable => vortex_array::validity::Validity::AllValid,
53        };
54
55        let array = match intersection {
56            Intersection::None => {
57                ConstantArray::new(Scalar::bool(false, nullability), lhs.len()).into_array()
58            }
59            Intersection::All => {
60                ConstantArray::new(Scalar::bool(true, nullability), lhs.len()).into_array()
61            }
62            Intersection::At(idx) => {
63                let mut buffer = BitBufferMut::new_unset(lhs.len());
64                buffer.set(idx);
65                BoolArray::new(buffer.freeze(), validity).into_array()
66            }
67        };
68
69        Ok(Some(array))
70    }
71}
72
73/// The intersection of a sequence with a value.
74pub(crate) enum Intersection {
75    None,
76    At(usize),
77    All,
78}
79
80/// Finds where a sequence equals `value`, or returns `None` for non-integer inputs.
81pub(crate) fn find_intersection(
82    base: PValue,
83    multiplier: PValue,
84    len: usize,
85    value: PValue,
86) -> Option<Intersection> {
87    if !value.ptype().is_int() || len == 0 {
88        return (len == 0).then_some(Intersection::None);
89    }
90    let (ascending, magnitude) = eval::step_parts(multiplier)?;
91
92    // Use the base's signedness to preserve values above `i64::MAX`.
93    let (towards, offset) = if base.ptype().is_signed_int() {
94        let base = base.cast::<i64>().vortex_expect("base fits its ptype");
95        let Ok(value) = value.cast::<i64>() else {
96            return Some(Intersection::None);
97        };
98        (value >= base, base.abs_diff(value))
99    } else {
100        let base = base.cast::<u64>().vortex_expect("base fits its ptype");
101        let Ok(value) = value.cast::<u64>() else {
102            return Some(Intersection::None);
103        };
104        (value >= base, base.abs_diff(value))
105    };
106
107    if offset == 0 {
108        return Some(if magnitude == 0 {
109            Intersection::All
110        } else {
111            Intersection::At(0)
112        });
113    }
114    if magnitude == 0 || towards != ascending || offset % magnitude != 0 {
115        return Some(Intersection::None);
116    }
117
118    Some(match usize::try_from(offset / magnitude) {
119        Ok(idx) if idx < len => Intersection::At(idx),
120        _ => Intersection::None,
121    })
122}
123
124#[cfg(test)]
125mod tests {
126    use std::sync::LazyLock;
127
128    use vortex_array::IntoArray;
129    use vortex_array::VortexSessionExecute;
130    use vortex_array::arrays::BoolArray;
131    use vortex_array::arrays::ConstantArray;
132    use vortex_array::assert_arrays_eq;
133    use vortex_array::builtins::ArrayBuiltins;
134    use vortex_array::dtype::Nullability::NonNullable;
135    use vortex_array::dtype::Nullability::Nullable;
136    use vortex_array::dtype::PType;
137    use vortex_array::scalar::PValue;
138    use vortex_array::scalar_fn::fns::operators::Operator;
139    use vortex_error::VortexResult;
140    use vortex_session::VortexSession;
141
142    use crate::Sequence;
143
144    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
145        let session = vortex_array::array_session();
146        crate::initialize(&session);
147        session
148    });
149
150    #[test]
151    fn test_compare_match() {
152        let lhs = Sequence::try_new_typed(2i64, 1, NonNullable, 4).unwrap();
153        let rhs = ConstantArray::new(4i64, lhs.len());
154        let result = lhs
155            .into_array()
156            .binary(rhs.into_array(), Operator::Eq)
157            .unwrap();
158        let expected = BoolArray::from_iter([false, false, true, false]);
159        assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx());
160    }
161
162    #[test]
163    fn test_compare_match_scale() {
164        let lhs = Sequence::try_new_typed(2i64, 3, Nullable, 4).unwrap();
165        let rhs = ConstantArray::new(8i64, lhs.len());
166        let result = lhs
167            .into_array()
168            .binary(rhs.into_array(), Operator::Eq)
169            .unwrap();
170        let expected = BoolArray::from_iter([Some(false), Some(false), Some(true), Some(false)]);
171        assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx());
172    }
173
174    #[test]
175    fn test_compare_no_match() {
176        let lhs = Sequence::try_new_typed(2i64, 1, NonNullable, 4).unwrap();
177        let rhs = ConstantArray::new(1i64, lhs.len());
178        let result = lhs
179            .into_array()
180            .binary(rhs.into_array(), Operator::Eq)
181            .unwrap();
182        let expected = BoolArray::from_iter([false, false, false, false]);
183        assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx());
184    }
185
186    #[test]
187    fn test_compare_descending_unsigned() -> VortexResult<()> {
188        let lhs = Sequence::try_new(
189            PValue::from(100i32),
190            PValue::from(-10i32),
191            PType::U8,
192            NonNullable,
193            5,
194        )?;
195        let rhs = ConstantArray::new(80u8, lhs.len());
196        let result = lhs.into_array().binary(rhs.into_array(), Operator::Eq)?;
197        let expected = BoolArray::from_iter([false, false, true, false, false]);
198        assert_arrays_eq!(result, expected, &mut SESSION.create_execution_ctx());
199
200        Ok(())
201    }
202
203    #[test]
204    fn test_compare_past_i64_max() -> VortexResult<()> {
205        let mut ctx = SESSION.create_execution_ctx();
206        let step = (1u64 << 63) + 1;
207        let lhs = Sequence::try_new(
208            PValue::from(1u64 << 62),
209            PValue::from(step),
210            PType::U64,
211            NonNullable,
212            2,
213        )?;
214
215        let hit = lhs.clone().into_array().binary(
216            ConstantArray::new((1u64 << 62) + step, lhs.len()).into_array(),
217            Operator::Eq,
218        )?;
219        assert_arrays_eq!(hit, BoolArray::from_iter([false, true]), &mut ctx);
220
221        let miss = lhs
222            .into_array()
223            .binary(ConstantArray::new(u64::MAX, 2).into_array(), Operator::Eq)?;
224        assert_arrays_eq!(miss, BoolArray::from_iter([false, false]), &mut ctx);
225
226        Ok(())
227    }
228
229    #[test]
230    fn test_compare_constant_sequence() -> VortexResult<()> {
231        let mut ctx = SESSION.create_execution_ctx();
232        let lhs = Sequence::try_new_typed(100i32, 0i32, NonNullable, 5)?;
233
234        let matches = lhs.clone().into_array().binary(
235            ConstantArray::new(100i32, lhs.len()).into_array(),
236            Operator::Eq,
237        )?;
238        assert_arrays_eq!(matches, BoolArray::from_iter([true; 5]), &mut ctx);
239
240        let misses = lhs
241            .into_array()
242            .binary(ConstantArray::new(7i32, 5).into_array(), Operator::Eq)?;
243        assert_arrays_eq!(misses, BoolArray::from_iter([false; 5]), &mut ctx);
244
245        Ok(())
246    }
247}
vortex_sequence/compute/compare.rs

vortex_sequence/compute/
compare.rs
