vortex_array/arrays/varbin/compute/
compare.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use arrow_array::BinaryArray;
5use arrow_array::LargeBinaryArray;
6use arrow_array::LargeStringArray;
7use arrow_array::StringArray;
8use arrow_ord::cmp;
9use arrow_schema::DataType;
10use vortex_buffer::BitBuffer;
11use vortex_error::VortexExpect as _;
12use vortex_error::VortexResult;
13use vortex_error::vortex_bail;
14use vortex_error::vortex_err;
15
16use crate::ArrayRef;
17use crate::ExecutionCtx;
18use crate::IntoArray;
19use crate::array::ArrayView;
20use crate::arrays::BoolArray;
21use crate::arrays::PrimitiveArray;
22use crate::arrays::VarBin;
23use crate::arrays::VarBinViewArray;
24use crate::arrays::varbin::VarBinArrayExt;
25use crate::arrow::Datum;
26use crate::arrow::from_arrow_columnar;
27use crate::builtins::ArrayBuiltins;
28use crate::dtype::DType;
29use crate::dtype::IntegerPType;
30use crate::match_each_integer_ptype;
31use crate::scalar_fn::fns::binary::CompareKernel;
32use crate::scalar_fn::fns::operators::CompareOperator;
33use crate::scalar_fn::fns::operators::Operator;
34
35// This implementation exists so we can have custom translation of RHS to arrow that's not the same as IntoCanonical
36impl CompareKernel for VarBin {
37    fn compare(
38        lhs: ArrayView<'_, VarBin>,
39        rhs: &ArrayRef,
40        operator: CompareOperator,
41        ctx: &mut ExecutionCtx,
42    ) -> VortexResult<Option<ArrayRef>> {
43        if let Some(rhs_const) = rhs.as_constant() {
44            let nullable = lhs.dtype().is_nullable() || rhs_const.dtype().is_nullable();
45            let len = lhs.len();
46
47            let rhs_is_empty = match rhs_const.dtype() {
48                DType::Binary(_) => rhs_const
49                    .as_binary()
50                    .is_empty()
51                    .vortex_expect("RHS should not be null"),
52                DType::Utf8(_) => rhs_const
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
59            if rhs_is_empty {
60                let buffer = match operator {
61                    CompareOperator::Gte => BitBuffer::new_set(len), /* Every possible value is >= "" */
62                    CompareOperator::Lt => BitBuffer::new_unset(len), // No value is < ""
63                    CompareOperator::Eq | CompareOperator::Lte => {
64                        let lhs_offsets = lhs.offsets().clone().execute::<PrimitiveArray>(ctx)?;
65                        match_each_integer_ptype!(lhs_offsets.ptype(), |P| {
66                            compare_offsets_to_empty::<P>(lhs_offsets, true)
67                        })
68                    }
69                    CompareOperator::NotEq | CompareOperator::Gt => {
70                        let lhs_offsets = lhs.offsets().clone().execute::<PrimitiveArray>(ctx)?;
71                        match_each_integer_ptype!(lhs_offsets.ptype(), |P| {
72                            compare_offsets_to_empty::<P>(lhs_offsets, false)
73                        })
74                    }
75                };
76
77                return Ok(Some(
78                    BoolArray::new(
79                        buffer,
80                        lhs.validity()?.union_nullability(rhs.dtype().nullability()),
81                    )
82                    .into_array(),
83                ));
84            }
85
86            let lhs = Datum::try_new(lhs.array(), ctx)?;
87
88            // The RHS scalar must match the LHS Arrow data type. VarBin with i64 offsets is
89            // converted to LargeBinary/LargeUtf8 (see `preferred_arrow_type`), and Arrow refuses to
90            // compare LargeBinary with Binary (or LargeUtf8 with Utf8).
91            let arrow_rhs: &dyn arrow_array::Datum = match (rhs_const.dtype(), lhs.data_type()) {
92                (DType::Utf8(_), DataType::LargeUtf8) => &rhs_const
93                    .as_utf8()
94                    .value()
95                    .map(LargeStringArray::new_scalar)
96                    .unwrap_or_else(|| arrow_array::Scalar::new(LargeStringArray::new_null(1))),
97                (DType::Utf8(_), _) => &rhs_const
98                    .as_utf8()
99                    .value()
100                    .map(StringArray::new_scalar)
101                    .unwrap_or_else(|| arrow_array::Scalar::new(StringArray::new_null(1))),
102                (DType::Binary(_), DataType::LargeBinary) => &rhs_const
103                    .as_binary()
104                    .value()
105                    .map(LargeBinaryArray::new_scalar)
106                    .unwrap_or_else(|| arrow_array::Scalar::new(LargeBinaryArray::new_null(1))),
107                (DType::Binary(_), _) => &rhs_const
108                    .as_binary()
109                    .value()
110                    .map(BinaryArray::new_scalar)
111                    .unwrap_or_else(|| arrow_array::Scalar::new(BinaryArray::new_null(1))),
112                _ => vortex_bail!(
113                    "VarBin array RHS can only be Utf8 or Binary, given {}",
114                    rhs_const.dtype()
115                ),
116            };
117
118            let array = match operator {
119                CompareOperator::Eq => cmp::eq(&lhs, arrow_rhs),
120                CompareOperator::NotEq => cmp::neq(&lhs, arrow_rhs),
121                CompareOperator::Gt => cmp::gt(&lhs, arrow_rhs),
122                CompareOperator::Gte => cmp::gt_eq(&lhs, arrow_rhs),
123                CompareOperator::Lt => cmp::lt(&lhs, arrow_rhs),
124                CompareOperator::Lte => cmp::lt_eq(&lhs, arrow_rhs),
125            }
126            .map_err(|err| vortex_err!("Failed to compare VarBin array: {}", err))?;
127
128            Ok(Some(from_arrow_columnar(&array, len, nullable, ctx)?))
129        } else if !rhs.is::<VarBin>() {
130            // NOTE: If the rhs is not a VarBin array it will be canonicalized to a VarBinView
131            // Arrow doesn't support comparing VarBin to VarBinView arrays, so we convert ourselves
132            // to VarBinView and re-invoke.
133            Ok(Some(
134                lhs.array()
135                    .clone()
136                    .execute::<VarBinViewArray>(ctx)?
137                    .into_array()
138                    .binary(rhs.clone(), Operator::from(operator))?,
139            ))
140        } else {
141            Ok(None)
142        }
143    }
144}
145
146fn compare_offsets_to_empty<P: IntegerPType>(offsets: PrimitiveArray, eq: bool) -> BitBuffer {
147    let fn_ = if eq { P::eq } else { P::ne };
148    let offsets = offsets.as_slice::<P>();
149    BitBuffer::collect_bool(offsets.len() - 1, |idx| {
150        let left = unsafe { offsets.get_unchecked(idx) };
151        let right = unsafe { offsets.get_unchecked(idx + 1) };
152        fn_(left, right)
153    })
154}
155
156#[cfg(test)]
157mod test {
158    use vortex_buffer::BitBuffer;
159    use vortex_buffer::ByteBuffer;
160
161    use crate::IntoArray;
162    use crate::VortexSessionExecute;
163    use crate::array_session;
164    use crate::arrays::BoolArray;
165    use crate::arrays::ConstantArray;
166    use crate::arrays::VarBinArray;
167    use crate::arrays::VarBinViewArray;
168    use crate::arrays::bool::BoolArrayExt;
169    use crate::builtins::ArrayBuiltins;
170    use crate::dtype::DType;
171    use crate::dtype::Nullability;
172    use crate::scalar::Scalar;
173    use crate::scalar_fn::fns::operators::Operator;
174
175    #[test]
176    fn test_binary_compare() {
177        let mut ctx = array_session().create_execution_ctx();
178        let array = VarBinArray::from_iter(
179            [Some(b"abc".to_vec()), None, Some(b"def".to_vec())],
180            DType::Binary(Nullability::Nullable),
181        );
182        let result = array
183            .into_array()
184            .binary(
185                ConstantArray::new(
186                    Scalar::binary(ByteBuffer::copy_from(b"abc"), Nullability::Nullable),
187                    3,
188                )
189                .into_array(),
190                Operator::Eq,
191            )
192            .unwrap()
193            .execute::<BoolArray>(&mut ctx)
194            .unwrap();
195
196        assert_eq!(
197            &result
198                .as_ref()
199                .validity()
200                .unwrap()
201                .execute_mask(result.as_ref().len(), &mut ctx)
202                .unwrap()
203                .to_bit_buffer(),
204            &BitBuffer::from_iter([true, false, true])
205        );
206        assert_eq!(
207            result.to_bit_buffer(),
208            BitBuffer::from_iter([true, false, false])
209        );
210    }
211
212    #[test]
213    fn varbinview_compare() {
214        let mut ctx = array_session().create_execution_ctx();
215        let array = VarBinArray::from_iter(
216            [Some(b"abc".to_vec()), None, Some(b"def".to_vec())],
217            DType::Binary(Nullability::Nullable),
218        );
219        let vbv = VarBinViewArray::from_iter(
220            [None, None, Some(b"def".to_vec())],
221            DType::Binary(Nullability::Nullable),
222        );
223        let result = array
224            .into_array()
225            .binary(vbv.into_array(), Operator::Eq)
226            .unwrap()
227            .execute::<BoolArray>(&mut ctx)
228            .unwrap();
229
230        assert_eq!(
231            result
232                .as_ref()
233                .validity()
234                .unwrap()
235                .execute_mask(result.as_ref().len(), &mut ctx)
236                .unwrap()
237                .to_bit_buffer(),
238            BitBuffer::from_iter([false, false, true])
239        );
240        assert_eq!(
241            result.to_bit_buffer(),
242            BitBuffer::from_iter([false, true, true])
243        );
244    }
245}
246
247#[cfg(test)]
248mod tests {
249    use vortex_buffer::ByteBuffer;
250
251    use crate::IntoArray;
252    use crate::VortexSessionExecute;
253    use crate::array_session;
254    use crate::arrays::BoolArray;
255    use crate::arrays::ConstantArray;
256    use crate::arrays::VarBinArray;
257    use crate::arrays::varbin::builder::VarBinBuilder;
258    use crate::assert_arrays_eq;
259    use crate::builtins::ArrayBuiltins;
260    use crate::dtype::DType;
261    use crate::dtype::Nullability;
262    use crate::scalar::Scalar;
263    use crate::scalar_fn::fns::operators::Operator;
264
265    #[test]
266    fn test_null_compare() {
267        let arr = VarBinArray::from_iter([Some("h")], DType::Utf8(Nullability::NonNullable));
268
269        let const_ = ConstantArray::new(Scalar::utf8("", Nullability::Nullable), 1);
270
271        assert_eq!(
272            arr.into_array()
273                .binary(const_.into_array(), Operator::Eq)
274                .unwrap()
275                .dtype(),
276            &DType::Bool(Nullability::Nullable)
277        );
278    }
279
280    /// Regression: a [`VarBinArray`] built with `i64` offsets is canonicalised to
281    /// Arrow `LargeUtf8` / `LargeBinary` by `preferred_arrow_type`. Without an explicit
282    /// branch in [`CompareKernel`], the constant RHS is wrapped in a `StringArray` /
283    /// `BinaryArray` and Arrow rejects the `LargeUtf8 == Utf8` mismatch. Triggering
284    /// this only requires `i64` offsets, not large data.
285    ///
286    /// [`CompareKernel`]: super::CompareKernel
287    #[test]
288    fn varbin_i64_offsets_compare_constant() {
289        let mut ctx = array_session().create_execution_ctx();
290        let mut builder = VarBinBuilder::<i64>::with_capacity(3);
291        builder.append_value(b"abc");
292        builder.append_value(b"xyz");
293        builder.append_value(b"abc");
294        let array = builder.finish(DType::Utf8(Nullability::NonNullable));
295
296        let result = array
297            .into_array()
298            .binary(
299                ConstantArray::new(Scalar::utf8("abc", Nullability::NonNullable), 3).into_array(),
300                Operator::Eq,
301            )
302            .unwrap();
303
304        let expected = BoolArray::from_iter([true, false, true]);
305        assert_arrays_eq!(result, expected, &mut ctx);
306    }
307
308    #[test]
309    fn varbin_i64_offsets_compare_constant_binary() {
310        let mut ctx = array_session().create_execution_ctx();
311        let mut builder = VarBinBuilder::<i64>::with_capacity(3);
312        builder.append_value(b"abc");
313        builder.append_value(b"xyz");
314        builder.append_value(b"abc");
315        let array = builder.finish(DType::Binary(Nullability::NonNullable));
316
317        let result = array
318            .into_array()
319            .binary(
320                ConstantArray::new(
321                    Scalar::binary(ByteBuffer::copy_from(b"abc"), Nullability::NonNullable),
322                    3,
323                )
324                .into_array(),
325                Operator::Eq,
326            )
327            .unwrap();
328
329        let expected = BoolArray::from_iter([true, false, true]);
330        assert_arrays_eq!(result, expected, &mut ctx);
331    }
332}
vortex_array/arrays/varbin/compute/compare.rs

vortex_array/arrays/varbin/compute/
compare.rs
