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_array::session::ArraySession;
78    use vortex_buffer::buffer;
79    use vortex_session::VortexSession;
80
81    use crate::Sparse;
82    use crate::initialize;
83
84    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
85        let session = VortexSession::empty().with::<ArraySession>();
86        initialize(&session);
87        session
88    });
89
90    #[rstest]
91    #[case(0i32, 100i32, StrictComparison::NonStrict, StrictComparison::NonStrict)]
92    #[case(5i32, 25i32, StrictComparison::Strict, StrictComparison::Strict)]
93    #[case(1i32, 20i32, StrictComparison::NonStrict, StrictComparison::Strict)]
94    fn between_matches_canonical(
95        #[case] lo: i32,
96        #[case] hi: i32,
97        #[case] lower_strict: StrictComparison,
98        #[case] upper_strict: StrictComparison,
99    ) {
100        let array = Sparse::try_new(
101            buffer![1u64, 3, 5].into_array(),
102            buffer![10i32, 20, 30].into_array(),
103            8,
104            Scalar::from(1i32),
105        )
106        .unwrap()
107        .into_array();
108        let len = array.len();
109        let options = BetweenOptions {
110            lower_strict,
111            upper_strict,
112        };
113
114        let lower = ConstantArray::new(Scalar::from(lo), len).into_array();
115        let upper = ConstantArray::new(Scalar::from(hi), len).into_array();
116
117        let mut ctx = SESSION.create_execution_ctx();
118
119        // Kernel path: between pushes through the Sparse encoding.
120        let kernel = array
121            .clone()
122            .between(lower.clone(), upper.clone(), options.clone())
123            .unwrap()
124            .execute::<Canonical>(&mut ctx)
125            .unwrap();
126
127        // Baseline: canonicalize the input first so between runs on a PrimitiveArray.
128        let canonical_input = array.execute::<Canonical>(&mut ctx).unwrap().into_array();
129        let baseline = canonical_input
130            .between(lower, upper, options)
131            .unwrap()
132            .execute::<Canonical>(&mut ctx)
133            .unwrap();
134
135        assert_arrays_eq!(kernel, baseline);
136    }
137}
vortex_sparse/compute/between.rs

vortex_sparse/compute/
between.rs
