vortex_sparse/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::ConstantArray;
9use vortex_array::builtins::ArrayBuiltins;
10use vortex_array::scalar_fn::fns::binary::CompareKernel;
11use vortex_array::scalar_fn::fns::binary::scalar_cmp;
12use vortex_array::scalar_fn::fns::operators::CompareOperator;
13use vortex_array::scalar_fn::fns::operators::Operator;
14use vortex_error::VortexResult;
15
16use crate::Sparse;
17use crate::SparseExt as _;
18
19/// Sparse-specific compare kernel.
20///
21/// When the RHS is a constant scalar, the result of any comparison is itself sparse:
22/// every unpatched position resolves to `compare(fill, rhs)`, and every patched position
23/// to `compare(patch, rhs)`. We push the comparison into the patches and rebuild a
24/// `Sparse<Bool>` with the new fill, preserving downstream sparsity (filter masks, etc.).
25///
26/// For non-constant RHS we decline and let the canonical fallback handle it.
27impl CompareKernel for Sparse {
28    fn compare(
29        lhs: ArrayView<'_, Self>,
30        rhs: &ArrayRef,
31        operator: CompareOperator,
32        _ctx: &mut ExecutionCtx,
33    ) -> VortexResult<Option<ArrayRef>> {
34        let Some(rhs_scalar) = rhs.as_constant() else {
35            return Ok(None);
36        };
37
38        let fill_bool = scalar_cmp(lhs.fill_scalar(), &rhs_scalar, operator)?;
39        let patches = lhs.patches();
40
41        let new_patches = patches.map_values(|values| {
42            let len = values.len();
43            values.binary(
44                ConstantArray::new(rhs_scalar.clone(), len).into_array(),
45                Operator::from(operator),
46            )
47        })?;
48
49        Ok(Some(
50            Sparse::try_new_from_patches(new_patches, fill_bool)?.into_array(),
51        ))
52    }
53}
54
55#[cfg(test)]
56mod tests {
57    use std::sync::LazyLock;
58
59    use rstest::rstest;
60    use vortex_array::Canonical;
61    use vortex_array::IntoArray;
62    use vortex_array::VortexSessionExecute;
63    use vortex_array::arrays::ConstantArray;
64    use vortex_array::assert_arrays_eq;
65    use vortex_array::builtins::ArrayBuiltins;
66    use vortex_array::scalar::Scalar;
67    use vortex_array::scalar_fn::fns::operators::Operator;
68    use vortex_buffer::buffer;
69    use vortex_session::VortexSession;
70
71    use crate::Sparse;
72    use crate::SparseArray;
73    use crate::initialize;
74
75    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
76        let session = vortex_array::array_session();
77        initialize(&session);
78        session
79    });
80
81    #[rstest]
82    #[case::eq_fill(Scalar::from(1i32), Operator::Eq)]
83    #[case::eq_patch(Scalar::from(10i32), Operator::Eq)]
84    #[case::gt(Scalar::from(5i32), Operator::Gt)]
85    #[case::lte(Scalar::from(10i32), Operator::Lte)]
86    #[case::neq(Scalar::from(1i32), Operator::NotEq)]
87    fn compare_matches_canonical(#[case] rhs: Scalar, #[case] op: Operator) {
88        let array: SparseArray = Sparse::try_new(
89            buffer![1u64, 3, 5].into_array(),
90            buffer![10i32, 20, 30].into_array(),
91            8,
92            Scalar::from(1i32),
93        )
94        .unwrap();
95        let arr = array.into_array();
96        let len = arr.len();
97        let mut ctx = SESSION.create_execution_ctx();
98
99        // Kernel path: compare pushes through the Sparse encoding.
100        let kernel_bool = arr
101            .binary(ConstantArray::new(rhs.clone(), len).into_array(), op)
102            .unwrap()
103            .execute::<Canonical>(&mut ctx)
104            .unwrap();
105
106        // Baseline: canonicalize first, then compare on the PrimitiveArray.
107        let canonical_input = arr.execute::<Canonical>(&mut ctx).unwrap().into_array();
108        let canonical_bool = canonical_input
109            .binary(ConstantArray::new(rhs, len).into_array(), op)
110            .unwrap()
111            .execute::<Canonical>(&mut ctx)
112            .unwrap();
113
114        assert_arrays_eq!(kernel_bool, canonical_bool, &mut ctx);
115    }
116}
vortex_sparse/compute/compare.rs

vortex_sparse/compute/
compare.rs
