vortex_fsst/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::builtins::ArrayBuiltins;
11use vortex_array::dtype::DType;
12use vortex_array::scalar::Scalar;
13use vortex_array::scalar_fn::fns::binary::CompareKernel;
14use vortex_array::scalar_fn::fns::operators::CompareOperator;
15use vortex_array::scalar_fn::fns::operators::Operator;
16use vortex_buffer::BitBuffer;
17use vortex_buffer::ByteBuffer;
18use vortex_error::VortexExpect;
19use vortex_error::VortexResult;
20use vortex_error::vortex_bail;
21
22use crate::FSST;
23use crate::FSSTArrayExt;
24impl CompareKernel for FSST {
25    fn compare(
26        lhs: ArrayView<'_, Self>,
27        rhs: &ArrayRef,
28        operator: CompareOperator,
29        ctx: &mut ExecutionCtx,
30    ) -> VortexResult<Option<ArrayRef>> {
31        match rhs.as_constant() {
32            Some(constant) => compare_fsst_constant(lhs, &constant, operator, ctx),
33            // Otherwise, fall back to the default comparison behavior.
34            _ => Ok(None),
35        }
36    }
37}
38
39/// Specialized compare function implementation used when performing against a constant
40fn compare_fsst_constant(
41    left: ArrayView<'_, FSST>,
42    right: &Scalar,
43    operator: CompareOperator,
44    ctx: &mut ExecutionCtx,
45) -> VortexResult<Option<ArrayRef>> {
46    let is_rhs_empty = match right.dtype() {
47        DType::Binary(_) => right
48            .as_binary()
49            .is_empty()
50            .vortex_expect("RHS should not be null"),
51        DType::Utf8(_) => right
52            .as_utf8()
53            .is_empty()
54            .vortex_expect("RHS should not be null"),
55        _ => vortex_bail!("VarBinArray can only have type of Binary or Utf8"),
56    };
57    if is_rhs_empty {
58        let buffer = match operator {
59            // Every possible value is gte ""
60            CompareOperator::Gte => BitBuffer::new_set(left.len()),
61            // No value is lt ""
62            CompareOperator::Lt => BitBuffer::new_unset(left.len()),
63            _ => left
64                .uncompressed_lengths()
65                .binary(
66                    ConstantArray::new(
67                        Scalar::zero_value(left.uncompressed_lengths().dtype()),
68                        left.uncompressed_lengths().len(),
69                    )
70                    .into_array(),
71                    operator.into(),
72                )?
73                .execute(ctx)?,
74        };
75
76        return Ok(Some(
77            BoolArray::new(
78                buffer,
79                left.array()
80                    .validity()?
81                    .union_nullability(right.dtype().nullability()),
82            )
83            .into_array(),
84        ));
85    }
86
87    // The following section only supports Eq/NotEq
88    if !matches!(operator, CompareOperator::Eq | CompareOperator::NotEq) {
89        return Ok(None);
90    }
91
92    let compressor = left.compressor();
93    let encoded_buffer = match left.dtype() {
94        DType::Utf8(_) => {
95            let value = right
96                .as_utf8()
97                .value()
98                .vortex_expect("Expected non-null scalar");
99            ByteBuffer::from(compressor.compress(value.as_bytes()))
100        }
101        DType::Binary(_) => {
102            let value = right
103                .as_binary()
104                .value()
105                .vortex_expect("Expected non-null scalar");
106            ByteBuffer::from(compressor.compress(value.as_slice()))
107        }
108        _ => unreachable!("FSSTArray can only have string or binary data type"),
109    };
110
111    let encoded_scalar = Scalar::binary(
112        encoded_buffer,
113        left.dtype().nullability() | right.dtype().nullability(),
114    );
115
116    let rhs = ConstantArray::new(encoded_scalar, left.len());
117    left.codes()
118        .into_array()
119        .binary(rhs.into_array(), Operator::from(operator))
120        .map(Some)
121}
122
123#[cfg(test)]
124mod tests {
125    use std::sync::LazyLock;
126
127    use vortex_array::IntoArray;
128    use vortex_array::VortexSessionExecute;
129    use vortex_array::arrays::BoolArray;
130    use vortex_array::arrays::ConstantArray;
131    use vortex_array::arrays::VarBinArray;
132    use vortex_array::assert_arrays_eq;
133    use vortex_array::builtins::ArrayBuiltins;
134    use vortex_array::dtype::DType;
135    use vortex_array::dtype::Nullability;
136    use vortex_array::scalar::Scalar;
137    use vortex_array::scalar_fn::fns::operators::Operator;
138    use vortex_error::VortexResult;
139    use vortex_session::VortexSession;
140
141    use crate::fsst_compress;
142    use crate::fsst_train_compressor;
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    #[cfg_attr(miri, ignore)]
152    fn test_compare_fsst() -> VortexResult<()> {
153        let mut ctx = SESSION.create_execution_ctx();
154        let lhs = VarBinArray::from_iter(
155            [
156                Some("hello"),
157                None,
158                Some("world"),
159                None,
160                Some("this is a very long string"),
161            ],
162            DType::Utf8(Nullability::Nullable),
163        )
164        .into_array();
165        let compressor = fsst_train_compressor(&lhs, &mut ctx)?;
166        let lhs = fsst_compress(&lhs, &compressor, &mut ctx)?;
167
168        let rhs = ConstantArray::new("world", lhs.len());
169
170        // Ensure fastpath for Eq exists, and returns correct answer
171        let equals = lhs
172            .clone()
173            .into_array()
174            .binary(rhs.clone().into_array(), Operator::Eq)?
175            .execute::<BoolArray>(&mut ctx)?;
176
177        assert_eq!(equals.dtype(), &DType::Bool(Nullability::Nullable));
178
179        assert_arrays_eq!(
180            &equals,
181            &BoolArray::from_iter([Some(false), None, Some(true), None, Some(false)]),
182            &mut ctx
183        );
184
185        // Ensure fastpath for Eq exists, and returns correct answer
186        let not_equals = lhs
187            .clone()
188            .into_array()
189            .binary(rhs.into_array(), Operator::NotEq)?
190            .execute::<BoolArray>(&mut ctx)?;
191
192        assert_eq!(not_equals.dtype(), &DType::Bool(Nullability::Nullable));
193        assert_arrays_eq!(
194            &not_equals,
195            &BoolArray::from_iter([Some(true), None, Some(false), None, Some(true)]),
196            &mut ctx
197        );
198
199        // Ensure null constants are handled correctly.
200        let null_rhs =
201            ConstantArray::new(Scalar::null(DType::Utf8(Nullability::Nullable)), lhs.len());
202        let equals_null = lhs
203            .clone()
204            .into_array()
205            .binary(null_rhs.clone().into_array(), Operator::Eq)?;
206        assert_arrays_eq!(
207            &equals_null,
208            &BoolArray::from_iter([None::<bool>, None, None, None, None]),
209            &mut ctx
210        );
211
212        let noteq_null = lhs
213            .into_array()
214            .binary(null_rhs.into_array(), Operator::NotEq)?;
215        assert_arrays_eq!(
216            &noteq_null,
217            &BoolArray::from_iter([None::<bool>, None, None, None, None]),
218            &mut ctx
219        );
220        Ok(())
221    }
222}
vortex_fsst/compute/compare.rs

vortex_fsst/compute/
compare.rs
