specialize fsst compare to empty const value (#2234)

AdamGS · web-flow · commit 4781ca5d5d93 · 2025-02-04T19:26:16.000Z
followup to #2227, I also tried to extract some of the shared logic up to the `compute` module.
diff --git a/encodings/fsst/src/compute/compare.rs b/encodings/fsst/src/compute/compare.rs
@@ -1,10 +1,11 @@
 use fsst::Symbol;
-use vortex_array::array::ConstantArray;
-use vortex_array::compute::{compare, CompareFn, Operator};
-use vortex_array::{Array, IntoArrayVariant};
+use vortex_array::array::{BoolArray, BooleanBuffer, ConstantArray};
+use vortex_array::compute::{compare, compare_lengths_to_empty, CompareFn, Operator};
+use vortex_array::variants::PrimitiveArrayTrait;
+use vortex_array::{Array, IntoArray, IntoArrayVariant, IntoCanonical};
 use vortex_buffer::ByteBuffer;
-use vortex_dtype::DType;
-use vortex_error::{VortexExpect, VortexResult};
+use vortex_dtype::{match_each_native_ptype, DType};
+use vortex_error::{vortex_bail, VortexExpect, VortexResult};
 
 use crate::{FSSTArray, FSSTEncoding};
 
@@ -15,26 +16,61 @@ impl CompareFn<FSSTArray> for FSSTEncoding {
         rhs: &Array,
         operator: Operator,
     ) -> VortexResult<Option<Array>> {
-        match (rhs.as_constant(), operator) {
-            (Some(constant), Operator::Eq | Operator::NotEq) => compare_fsst_constant(
-                lhs,
-                &ConstantArray::new(constant, lhs.len()),
-                operator == Operator::Eq,
-            )
-            .map(Some),
+        match rhs.as_constant() {
+            Some(constant) => {
+                compare_fsst_constant(lhs, &ConstantArray::new(constant, lhs.len()), operator)
+            }
             // Otherwise, fall back to the default comparison behavior.
             _ => Ok(None),
         }
     }
 }
 
-/// Specialized compare function implementation used when performing equals or not equals against
-/// a constant.
+/// Specialized compare function implementation used when performing against a constant
 fn compare_fsst_constant(
     left: &FSSTArray,
     right: &ConstantArray,
-    equal: bool,
-) -> VortexResult<Array> {
+    operator: Operator,
+) -> VortexResult<Option<Array>> {
+    let rhs_scalar = right.scalar();
+    let is_rhs_empty = match rhs_scalar.dtype() {
+        DType::Binary(_) => rhs_scalar
+            .as_binary()
+            .is_empty()
+            .vortex_expect("RHS should not be null"),
+        DType::Utf8(_) => rhs_scalar
+            .as_utf8()
+            .is_empty()
+            .vortex_expect("RHS should not be null"),
+        _ => vortex_bail!("VarBinArray can only have type of Binary or Utf8"),
+    };
+    if is_rhs_empty {
+        let buffer = match operator {
+            // Every possible value is gte ""
+            Operator::Gte => BooleanBuffer::new_set(left.len()),
+            // No value is lt ""
+            Operator::Lt => BooleanBuffer::new_unset(left.len()),
+            _ => {
+                let uncompressed_lengths = left
+                    .uncompressed_lengths()
+                    .into_canonical()?
+                    .into_primitive()?;
+                match_each_native_ptype!(uncompressed_lengths.ptype(), |$P| {
+                    compare_lengths_to_empty(uncompressed_lengths.as_slice::<$P>().iter().copied(), operator)
+                })
+            }
+        };
+
+        return Ok(Some(
+            BoolArray::try_new(buffer, left.validity())?.into_array(),
+        ));
+    }
+
+    // The following section only supports Eq/NotEq
+    if !matches!(operator, Operator::Eq | Operator::NotEq) {
+        return Ok(None);
+    }
+
     let symbols = left.symbols().into_primitive()?;
     let symbols_u64 = symbols.as_slice::<u64>();
 
@@ -68,11 +104,7 @@ fn compare_fsst_constant(
     };
 
     let rhs = ConstantArray::new(encoded_scalar, left.len());
-    compare(
-        left.codes(),
-        rhs,
-        if equal { Operator::Eq } else { Operator::NotEq },
-    )
+    compare(left.codes(), rhs, operator).map(Some)
 }
 
 #[cfg(test)]
diff --git a/vortex-array/src/array/varbin/compute/compare.rs b/vortex-array/src/array/varbin/compute/compare.rs
@@ -2,13 +2,12 @@ use arrow_array::{BinaryArray, StringArray};
 use arrow_buffer::BooleanBuffer;
 use arrow_ord::cmp;
 use itertools::Itertools;
-use num_traits::Zero;
 use vortex_dtype::{match_each_native_ptype, DType, NativePType};
 use vortex_error::{vortex_bail, vortex_err, VortexExpect as _, VortexResult};
 
 use crate::array::{BoolArray, PrimitiveArray, VarBinArray, VarBinEncoding};
 use crate::arrow::{from_arrow_array_with_len, Datum};
-use crate::compute::{CompareFn, Operator};
+use crate::compute::{compare_lengths_to_empty, CompareFn, Operator};
 use crate::variants::PrimitiveArrayTrait as _;
 use crate::{Array, IntoArray, IntoCanonical};
 
@@ -25,24 +24,15 @@ impl CompareFn<VarBinArray> for VarBinEncoding {
             let len = lhs.len();
 
             let rhs_is_empty = match rhs_const.dtype() {
-                DType::Utf8(_) => {
-                    let v = rhs_const
-                        .as_utf8()
-                        .value()
-                        .vortex_expect("null scalar handled in top-level");
-                    v.is_empty()
-                }
-                DType::Binary(_) => {
-                    let v = rhs_const
-                        .as_binary()
-                        .value()
-                        .vortex_expect("null scalar handled in top-level");
-                    v.is_empty()
-                }
-                _ => vortex_bail!(
-                    "VarBin array RHS can only be Utf8 or Binary, given {}",
-                    rhs_const.dtype()
-                ),
+                DType::Binary(_) => rhs_const
+                    .as_binary()
+                    .is_empty()
+                    .vortex_expect("RHS should not be null"),
+                DType::Utf8(_) => rhs_const
+                    .as_utf8()
+                    .is_empty()
+                    .vortex_expect("RHS should not be null"),
+                _ => vortex_bail!("VarBinArray can only have type of Binary or Utf8"),
             };
 
             if rhs_is_empty {
@@ -54,7 +44,7 @@ impl CompareFn<VarBinArray> for VarBinEncoding {
                     _ => {
                         let lhs_offsets = lhs.offsets().into_canonical()?.into_primitive()?;
                         match_each_native_ptype!(lhs_offsets.ptype(), |$P| {
-                            compare_to_empty::<$P>(lhs_offsets, operator)
+                            compare_offsets_to_empty::<$P>(lhs_offsets, operator)
                         })
                     }
                 };
@@ -101,25 +91,14 @@ impl CompareFn<VarBinArray> for VarBinEncoding {
     }
 }
 
-/// All comparison can be expressed in terms of equality. "" is the absolute min of possible value.
-const fn cmp_to_empty<T: PartialEq + PartialOrd + Zero>(op: Operator) -> fn(T) -> bool {
-    match op {
-        Operator::Eq | Operator::Lte => |v| v == T::zero(),
-        Operator::NotEq | Operator::Gt => |v| v != T::zero(),
-        Operator::Gte => |_| true,
-        Operator::Lt => |_| false,
-    }
-}
-
-fn compare_to_empty<T>(array: PrimitiveArray, op: Operator) -> BooleanBuffer
-where
-    T: NativePType,
-{
-    let cmp_fn = cmp_to_empty::<T>(op);
-    let slice = array.as_slice::<T>();
-    slice
+fn compare_offsets_to_empty<P: NativePType>(
+    offsets: PrimitiveArray,
+    operator: Operator,
+) -> BooleanBuffer {
+    let lengths_iter = offsets
+        .as_slice::<P>()
         .iter()
         .tuple_windows()
-        .map(|(&s, &e)| cmp_fn(e - s))
-        .collect::<BooleanBuffer>()
+        .map(|(&s, &e)| e - s);
+    compare_lengths_to_empty(lengths_iter, operator)
 }
diff --git a/vortex-array/src/compute/compare.rs b/vortex-array/src/compute/compare.rs
@@ -1,8 +1,9 @@
 use core::fmt;
 use std::fmt::{Display, Formatter};
 
+use arrow_buffer::BooleanBuffer;
 use arrow_ord::cmp;
-use vortex_dtype::{DType, Nullability};
+use vortex_dtype::{DType, NativePType, Nullability};
 use vortex_error::{vortex_bail, VortexError, VortexResult};
 use vortex_scalar::Scalar;
 
@@ -166,6 +167,24 @@ pub fn compare(
     Ok(result)
 }
 
+/// Helper function to compare empty values with arrays that have external value length information
+/// like `VarBin`.
+pub fn compare_lengths_to_empty<P, I>(lengths: I, op: Operator) -> BooleanBuffer
+where
+    P: NativePType,
+    I: Iterator<Item = P>,
+{
+    // All comparison can be expressed in terms of equality. "" is the absolute min of possible value.
+    let cmp_fn = match op {
+        Operator::Eq | Operator::Lte => |v| v == P::zero(),
+        Operator::NotEq | Operator::Gt => |v| v != P::zero(),
+        Operator::Gte => |_| true,
+        Operator::Lt => |_| false,
+    };
+
+    lengths.map(cmp_fn).collect::<BooleanBuffer>()
+}
+
 /// Implementation of `CompareFn` using the Arrow crate.
 fn arrow_compare(left: &Array, right: &Array, operator: Operator) -> VortexResult<Array> {
     let nullable = left.dtype().is_nullable() || right.dtype().is_nullable();
@@ -324,4 +343,18 @@ mod tests {
         assert_eq!(res.as_bool().value(), Some(false));
         assert_eq!(compare.len(), 10);
     }
+
+    #[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])]
+    #[case(Operator::Gte, vec![true, true, true, true])]
+    #[case(Operator::Lt, vec![false, false, false, false])]
+    #[case(Operator::Lte, vec![false, false, false, true])]
+    fn test_cmp_to_empty(#[case] op: Operator, #[case] expected: Vec<bool>) {
+        let lengths: Vec<i32> = vec![1, 5, 7, 0];
+
+        let output = compare_lengths_to_empty(lengths.iter().copied(), op);
+        assert_eq!(Vec::from_iter(output.iter()), expected);
+    }
 }
diff --git a/vortex-array/src/compute/mod.rs b/vortex-array/src/compute/mod.rs
@@ -14,7 +14,7 @@ pub use boolean::{
     and, and_kleene, binary_boolean, or, or_kleene, BinaryBooleanFn, BinaryOperator,
 };
 pub use cast::{try_cast, CastFn};
-pub use compare::{compare, scalar_cmp, CompareFn, Operator};
+pub use compare::{compare, compare_lengths_to_empty, scalar_cmp, CompareFn, Operator};
 pub use fill_forward::{fill_forward, FillForwardFn};
 pub use fill_null::{fill_null, FillNullFn};
 pub use filter::{filter, FilterFn};
diff --git a/vortex-scalar/src/binary.rs b/vortex-scalar/src/binary.rs
@@ -46,6 +46,11 @@ impl<'a> BinaryScalar<'a> {
             ))),
         ))
     }
+
+    /// Returns whether its value is non-null and empty, otherwise `None`.
+    pub fn is_empty(&self) -> Option<bool> {
+        self.value.as_ref().map(|v| v.is_empty())
+    }
 }
 
 impl Scalar {
diff --git a/vortex-scalar/src/utf8.rs b/vortex-scalar/src/utf8.rs
@@ -47,6 +47,11 @@ impl<'a> Utf8Scalar<'a> {
             ))),
         ))
     }
+
+    /// Returns whether its value is non-null and empty, otherwise `None`.
+    pub fn is_empty(&self) -> Option<bool> {
+        self.value.as_ref().map(|v| v.is_empty())
+    }
 }
 
 impl Scalar {

Original file line number	Diff line number	Diff line change
`@@ -46,6 +46,11 @@ impl<'a> BinaryScalar<'a> {`
`46`	`46`	`))),`
`47`	`47`	`))`
`48`	`48`	`}`
	`49`	`+`
	`50`	+ /// Returns whether its value is non-null and empty, otherwise `None`.
	`51`	`+ pub fn is_empty(&self) -> Option<bool> {`
	`52`	`+ self.value.as_ref().map(\|v\| v.is_empty())`
	`53`	`+ }`
`49`	`54`	`}`
`50`	`55`
`51`	`56`	`impl Scalar {`
Original file line number	Diff line number	Diff line change
`@@ -47,6 +47,11 @@ impl<'a> Utf8Scalar<'a> {`
`47`	`47`	`))),`
`48`	`48`	`))`
`49`	`49`	`}`
	`50`	`+`
	`51`	+ /// Returns whether its value is non-null and empty, otherwise `None`.
	`52`	`+ pub fn is_empty(&self) -> Option<bool> {`
	`53`	`+ self.value.as_ref().map(\|v\| v.is_empty())`
	`54`	`+ }`
`50`	`55`	`}`
`51`	`56`
`52`	`57`	`impl Scalar {`