vortex_sparse/compute/
between.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::ConstantArray;
9use vortex_array::builtins::ArrayBuiltins;
10use vortex_array::scalar_fn::fns::between::BetweenKernel;
11use vortex_array::scalar_fn::fns::between::BetweenOptions;
12use vortex_error::VortexResult;
13
14use crate::Sparse;
15use crate::SparseExt as _;
16
17/// Sparse-specific between kernel.
18///
19/// `lower <= x <= upper` (with per-bound strictness) over a Sparse column with constant
20/// bounds is itself sparse: every unpatched position resolves to `between(F, lo, hi)` and
21/// every patched position to `between(patch, lo, hi)`. We push the range check into the
22/// patches and rebuild a `Sparse<Bool>` with the new fill, preserving downstream sparsity.
23///
24/// Declines (falls back to canonical) unless both bounds are constants.
25impl BetweenKernel for Sparse {
26    fn between(
27        array: ArrayView<'_, Self>,
28        lower: &ArrayRef,
29        upper: &ArrayRef,
30        options: &BetweenOptions,
31        ctx: &mut ExecutionCtx,
32    ) -> VortexResult<Option<ArrayRef>> {
33        let (Some(lo), Some(hi)) = (lower.as_constant(), upper.as_constant()) else {
34            return Ok(None);
35        };
36
37        let patches = array.patches();
38
39        let fill_bool = ConstantArray::new(array.fill_scalar().clone(), 1)
40            .into_array()
41            .between(
42                ConstantArray::new(lo.clone(), 1).into_array(),
43                ConstantArray::new(hi.clone(), 1).into_array(),
44                options.clone(),
45            )?
46            .execute_scalar(0, ctx)?;
47
48        let new_patches = patches.map_values(|values| {
49            let len = values.len();
50            values.between(
51                ConstantArray::new(lo.clone(), len).into_array(),
52                ConstantArray::new(hi.clone(), len).into_array(),
53                options.clone(),
54            )
55        })?;
56
57        Ok(Some(
58            Sparse::try_new_from_patches(new_patches, fill_bool)?.into_array(),
59        ))
60    }
61}
62
63#[cfg(test)]
64mod tests {
65    use std::sync::LazyLock;
66
67    use rstest::rstest;
68    use vortex_array::Canonical;
69    use vortex_array::IntoArray;
70    use vortex_array::VortexSessionExecute;
71    use vortex_array::arrays::ConstantArray;
72    use vortex_array::assert_arrays_eq;
73    use vortex_array::builtins::ArrayBuiltins;
74    use vortex_array::scalar::Scalar;
75    use vortex_array::scalar_fn::fns::between::BetweenOptions;
76    use vortex_array::scalar_fn::fns::between::StrictComparison;
77    use vortex_buffer::buffer;
78    use vortex_session::VortexSession;
79
80    use crate::Sparse;
81    use crate::initialize;
82
83    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
84        let session = vortex_array::array_session();
85        initialize(&session);
86        session
87    });
88
89    #[rstest]
90    #[case(0i32, 100i32, StrictComparison::NonStrict, StrictComparison::NonStrict)]
91    #[case(5i32, 25i32, StrictComparison::Strict, StrictComparison::Strict)]
92    #[case(1i32, 20i32, StrictComparison::NonStrict, StrictComparison::Strict)]
93    fn between_matches_canonical(
94        #[case] lo: i32,
95        #[case] hi: i32,
96        #[case] lower_strict: StrictComparison,
97        #[case] upper_strict: StrictComparison,
98    ) {
99        let array = Sparse::try_new(
100            buffer![1u64, 3, 5].into_array(),
101            buffer![10i32, 20, 30].into_array(),
102            8,
103            Scalar::from(1i32),
104        )
105        .unwrap()
106        .into_array();
107        let len = array.len();
108        let options = BetweenOptions {
109            lower_strict,
110            upper_strict,
111        };
112
113        let lower = ConstantArray::new(Scalar::from(lo), len).into_array();
114        let upper = ConstantArray::new(Scalar::from(hi), len).into_array();
115
116        let mut ctx = SESSION.create_execution_ctx();
117
118        // Kernel path: between pushes through the Sparse encoding.
119        let kernel = array
120            .clone()
121            .between(lower.clone(), upper.clone(), options.clone())
122            .unwrap()
123            .execute::<Canonical>(&mut ctx)
124            .unwrap();
125
126        // Baseline: canonicalize the input first so between runs on a PrimitiveArray.
127        let canonical_input = array.execute::<Canonical>(&mut ctx).unwrap().into_array();
128        let baseline = canonical_input
129            .between(lower, upper, options)
130            .unwrap()
131            .execute::<Canonical>(&mut ctx)
132            .unwrap();
133
134        assert_arrays_eq!(kernel, baseline, &mut ctx);
135    }
136}
vortex_sparse/compute/between.rs

vortex_sparse/compute/
between.rs
