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

vortex_fsst/compute/
compare.rs
