fix: Handle fallback to arrow_compare between canonical and non canonical encodings of Utf8 and Binary (#2696)

robert3005 · web-flow · commit 4f7e98d69187 · 2025-03-25T12:00:03.000Z
diff --git a/fuzz/fuzz_targets/file_io.rs b/fuzz/fuzz_targets/file_io.rs
@@ -14,7 +14,7 @@ use vortex_array::stream::ArrayStreamArrayExt;
 use vortex_array::{Array, ArrayRef, ToCanonical};
 use vortex_buffer::ByteBufferMut;
 use vortex_dtype::{DType, StructDType};
-use vortex_error::VortexUnwrap;
+use vortex_error::{VortexUnwrap, vortex_panic};
 use vortex_file::{VortexOpenOptions, VortexWriteOptions};
 
 fuzz_target!(|array_data: ArbitraryArray| -> Corpus {
@@ -66,13 +66,21 @@ fuzz_target!(|array_data: ArbitraryArray| -> Corpus {
         if matches!(array_data.dtype(), DType::Struct(_, _) | DType::List(_, _)) {
             compare_struct(array_data, output);
         } else {
-            let r = compare(&array_data, &output, Operator::Eq).vortex_unwrap();
-            let true_count = r
-                .to_bool()
+            let bool_result = compare(&array_data, &output, Operator::Eq)
                 .vortex_unwrap()
-                .boolean_buffer()
-                .count_set_bits();
-            assert_eq!(true_count, array_data.len());
+                .to_bool()
+                .vortex_unwrap();
+            let true_count = bool_result.boolean_buffer().count_set_bits();
+            if true_count != array_data.len()
+                && (bool_result.all_valid().vortex_unwrap()
+                    || array_data.all_valid().vortex_unwrap())
+            {
+                vortex_panic!(
+                    "Failed to match original array {}with{}",
+                    array_data.tree_display(),
+                    output.tree_display()
+                );
+            }
         }
     });
 
diff --git a/vortex-array/src/arrays/varbin/compute/compare.rs b/vortex-array/src/arrays/varbin/compute/compare.rs
@@ -54,7 +54,7 @@ impl CompareFn<&VarBinArray> for VarBinEncoding {
                 ));
             }
 
-            let lhs = Datum::try_new(lhs.clone().into_array())?;
+            let lhs = Datum::try_new(lhs)?;
 
             // TODO(robert): Handle LargeString/Binary arrays
             let arrow_rhs: &dyn arrow_array::Datum = match rhs_const.dtype() {
diff --git a/vortex-array/src/arrays/varbin/compute/to_arrow.rs b/vortex-array/src/arrays/varbin/compute/to_arrow.rs
@@ -31,10 +31,14 @@ impl ToArrowFn<&VarBinArray> for VarBinEncoding {
         data_type: &DataType,
     ) -> VortexResult<Option<ArrayRef>> {
         let array_ref = match data_type {
-            DataType::BinaryView | DataType::FixedSizeBinary(_) | DataType::Utf8View => {
+            DataType::FixedSizeBinary(_) => {
                 // TODO(ngates): we should support converting VarBin into these Arrow arrays.
                 return Ok(None);
             }
+            DataType::BinaryView | DataType::Utf8View => Ok(arrow_cast::cast(
+                &*varbin_to_arrow::<i32>(array)?,
+                data_type,
+            )?),
             DataType::Binary | DataType::Utf8 => {
                 // These are both supported with a zero-copy cast, see below
                 varbin_to_arrow::<i32>(array)
diff --git a/vortex-array/src/arrays/varbinview/compute/mod.rs b/vortex-array/src/arrays/varbinview/compute/mod.rs
@@ -42,6 +42,10 @@ impl ComputeVTable for VarBinViewEncoding {
         Some(self)
     }
 
+    fn min_max_fn(&self) -> Option<&dyn MinMaxFn<&dyn Array>> {
+        Some(self)
+    }
+
     fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn<&dyn Array>> {
         Some(self)
     }
@@ -58,10 +62,6 @@ impl ComputeVTable for VarBinViewEncoding {
         Some(self)
     }
 
-    fn min_max_fn(&self) -> Option<&dyn MinMaxFn<&dyn Array>> {
-        Some(self)
-    }
-
     fn uncompressed_size_fn(&self) -> Option<&dyn UncompressedSizeFn<&dyn Array>> {
         Some(self)
     }
diff --git a/vortex-array/src/arrow/datum.rs b/vortex-array/src/arrow/datum.rs
@@ -1,4 +1,5 @@
 use arrow_array::{Array as ArrowArray, ArrayRef as ArrowArrayRef, Datum as ArrowDatum};
+use arrow_schema::DataType;
 use vortex_error::{VortexResult, vortex_panic};
 
 use crate::arrays::ConstantArray;
@@ -15,15 +16,32 @@ pub struct Datum {
 
 impl Datum {
     /// Create a new [`Datum`] from an [`ArrayRef`], which can then be passed to Arrow compute.
-    pub fn try_new(array: ArrayRef) -> VortexResult<Self> {
+    pub fn try_new(array: &dyn Array) -> VortexResult<Self> {
         if array.is_constant() {
             Ok(Self {
-                array: slice(&array, 0, 1)?.into_arrow_preferred()?,
+                array: slice(array, 0, 1)?.into_arrow_preferred()?,
                 is_scalar: true,
             })
         } else {
             Ok(Self {
-                array: array.into_arrow_preferred()?,
+                array: array.to_array().into_arrow_preferred()?,
+                is_scalar: false,
+            })
+        }
+    }
+
+    pub fn with_target_datatype(
+        array: &dyn Array,
+        target_datatype: &DataType,
+    ) -> VortexResult<Self> {
+        if array.is_constant() {
+            Ok(Self {
+                array: slice(array, 0, 1)?.into_arrow(target_datatype)?,
+                is_scalar: true,
+            })
+        } else {
+            Ok(Self {
+                array: array.to_array().into_arrow(target_datatype)?,
                 is_scalar: false,
             })
         }
diff --git a/vortex-array/src/compute/binary_numeric.rs b/vortex-array/src/compute/binary_numeric.rs
@@ -150,8 +150,8 @@ fn arrow_numeric(
     let nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable();
     let len = lhs.len();
 
-    let left = Datum::try_new(lhs.to_array())?;
-    let right = Datum::try_new(rhs.to_array())?;
+    let left = Datum::try_new(lhs)?;
+    let right = Datum::try_new(rhs)?;
 
     let array = match operator {
         BinaryNumericOperator::Add => arrow_arith::numeric::add(&left, &right)?,
diff --git a/vortex-array/src/compute/compare.rs b/vortex-array/src/compute/compare.rs
@@ -3,6 +3,7 @@ use std::fmt::{Display, Formatter};
 
 use arrow_buffer::BooleanBuffer;
 use arrow_ord::cmp;
+use arrow_schema::DataType;
 use vortex_dtype::{DType, NativePType, Nullability};
 use vortex_error::{VortexExpect, VortexResult, vortex_bail};
 use vortex_scalar::Scalar;
@@ -194,8 +195,8 @@ fn arrow_compare(
     operator: Operator,
 ) -> VortexResult<ArrayRef> {
     let nullable = left.dtype().is_nullable() || right.dtype().is_nullable();
-    let lhs = Datum::try_new(left.to_array())?;
-    let rhs = Datum::try_new(right.to_array())?;
+    let lhs = datum_for_cmp(left)?;
+    let rhs = datum_for_cmp(right)?;
 
     let array = match operator {
         Operator::Eq => cmp::eq(&lhs, &rhs)?,
@@ -250,14 +251,26 @@ pub fn scalar_cmp(lhs: &Scalar, rhs: &Scalar, operator: Operator) -> Scalar {
     }
 }
 
+// Make sure both of the arrays end up with the same arrow data type
+fn datum_for_cmp(array: &dyn Array) -> VortexResult<Datum> {
+    if matches!(array.dtype(), DType::Utf8(_)) {
+        Datum::with_target_datatype(array, &DataType::Utf8View)
+    } else if matches!(array.dtype(), DType::Binary(_)) {
+        Datum::with_target_datatype(array, &DataType::BinaryView)
+    } else {
+        Datum::try_new(array)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use arrow_buffer::BooleanBuffer;
     use itertools::Itertools;
+    use rstest::rstest;
 
     use super::*;
     use crate::ToCanonical;
-    use crate::arrays::{BoolArray, ConstantArray};
+    use crate::arrays::{BoolArray, ConstantArray, VarBinArray, VarBinViewArray};
     use crate::validity::Validity;
 
     fn to_int_indices(indices_bits: BoolArray) -> Vec<u64> {
@@ -337,7 +350,7 @@ mod tests {
         assert_eq!(compare.len(), 10);
     }
 
-    #[rstest::rstest]
+    #[rstest]
     #[case(Operator::Eq, vec![false, false, false, true])]
     #[case(Operator::NotEq, vec![true, true, true, false])]
     #[case(Operator::Gt, vec![true, true, true, false])]
@@ -350,4 +363,17 @@ mod tests {
         let output = compare_lengths_to_empty(lengths.iter().copied(), op);
         assert_eq!(Vec::from_iter(output.iter()), expected);
     }
+
+    #[rstest]
+    #[case(VarBinArray::from(vec!["a", "b"]).into_array(), VarBinViewArray::from_iter_str(["a", "b"]).into_array())]
+    #[case(VarBinViewArray::from_iter_str(["a", "b"]).into_array(), VarBinArray::from(vec!["a", "b"]).into_array())]
+    #[case(VarBinArray::from(vec!["a".as_bytes(), "b".as_bytes()]).into_array(), VarBinViewArray::from_iter_bin(["a".as_bytes(), "b".as_bytes()]).into_array())]
+    #[case(VarBinViewArray::from_iter_bin(["a".as_bytes(), "b".as_bytes()]).into_array(), VarBinArray::from(vec!["a".as_bytes(), "b".as_bytes()]).into_array())]
+    fn arrow_compare_different_encodings(#[case] left: ArrayRef, #[case] right: ArrayRef) {
+        let res = arrow_compare(&left, &right, Operator::Eq).unwrap();
+        assert_eq!(
+            res.to_bool().unwrap().boolean_buffer().count_set_bits(),
+            left.len()
+        );
+    }
 }
diff --git a/vortex-array/src/compute/like.rs b/vortex-array/src/compute/like.rs
@@ -116,8 +116,8 @@ pub(crate) fn arrow_like(
         "Arrow Like: length mismatch for {}",
         array.encoding()
     );
-    let lhs = Datum::try_new(array.to_array())?;
-    let rhs = Datum::try_new(pattern.to_array())?;
+    let lhs = Datum::try_new(array)?;
+    let rhs = Datum::try_new(pattern)?;
 
     let result = match (options.negated, options.case_insensitive) {
         (false, false) => arrow_string::like::like(&lhs, &rhs)?,

Original file line number	Diff line number	Diff line change
`@@ -54,7 +54,7 @@ impl CompareFn<&VarBinArray> for VarBinEncoding {`
`54`	`54`	`));`
`55`	`55`	`}`
`56`	`56`
`57`		`- let lhs = Datum::try_new(lhs.clone().into_array())?;`
	`57`	`+ let lhs = Datum::try_new(lhs)?;`
`58`	`58`
`59`	`59`	`// TODO(robert): Handle LargeString/Binary arrays`
`60`	`60`	`let arrow_rhs: &dyn arrow_array::Datum = match rhs_const.dtype() {`
Original file line number	Diff line number	Diff line change
`@@ -42,6 +42,10 @@ impl ComputeVTable for VarBinViewEncoding {`
`42`	`42`	`Some(self)`
`43`	`43`	`}`
`44`	`44`
	`45`	`+ fn min_max_fn(&self) -> Option<&dyn MinMaxFn<&dyn Array>> {`
	`46`	`+ Some(self)`
	`47`	`+ }`
	`48`	`+`
`45`	`49`	`fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn<&dyn Array>> {`
`46`	`50`	`Some(self)`
`47`	`51`	`}`
`@@ -58,10 +62,6 @@ impl ComputeVTable for VarBinViewEncoding {`
`58`	`62`	`Some(self)`
`59`	`63`	`}`
`60`	`64`
`61`		`- fn min_max_fn(&self) -> Option<&dyn MinMaxFn<&dyn Array>> {`
`62`		`- Some(self)`
`63`		`- }`
`64`		`-`
`65`	`65`	`fn uncompressed_size_fn(&self) -> Option<&dyn UncompressedSizeFn<&dyn Array>> {`
`66`	`66`	`Some(self)`
`67`	`67`	`}`