Improvement: Faster primitive is_constant (#2873)

Author: robert3005

Cargo.lock

@@ -26,6 +26,7 @@ arrow-ord = { workspace = true }
 arrow-schema = { workspace = true }
 arrow-select = { workspace = true }
 arrow-string = { workspace = true }
+cfg-if = { workspace = true }
 enum-iterator = { workspace = true }
 flatbuffers = { workspace = true }
 flexbuffers = { workspace = true }

vortex-array/Cargo.toml

@@ -1,32 +1,67 @@
+use vortex_dtype::half::f16;
 use vortex_dtype::{NativePType, match_each_native_ptype};
+use vortex_error::VortexResult;
 
 use crate::arrays::{PrimitiveArray, PrimitiveEncoding};
 use crate::compute::{IsConstantFn, IsConstantOpts};
 use crate::variants::PrimitiveArrayTrait;
 
+cfg_if::cfg_if! {
+    if #[cfg(target_feature = "avx2")] {
+        const IS_CONST_LANE_WIDTH: usize = 32;
+    } else {
+        const IS_CONST_LANE_WIDTH: usize = 16;
+    }
+}
+
 impl IsConstantFn<&PrimitiveArray> for PrimitiveEncoding {
     fn is_constant(
         &self,
         array: &PrimitiveArray,
         _opts: &IsConstantOpts,
-    ) -> vortex_error::VortexResult<Option<bool>> {
-        let is_constant = match_each_native_ptype!(array.ptype(), |$P| {
-            compute_is_constant(array.as_slice::<$P>())
+    ) -> VortexResult<Option<bool>> {
+        let is_constant = match_each_native_ptype!(array.ptype(), integral: |$P| {
+            compute_is_constant::<_, {IS_CONST_LANE_WIDTH / size_of::<$P>()}>(array.as_slice::<$P>())
+        } floating_point: |$P| {
+            compute_is_constant::<_, {IS_CONST_LANE_WIDTH / size_of::<$P>()}>(unsafe { std::mem::transmute::<&[$P], &[<$P as EqFloat>::IntType]>(array.as_slice::<$P>()) })
         });
 
         Ok(Some(is_constant))
     }
 }
 
+// Assumes any floating point has been cast into its bit representation for which != and !is_eq are the same
 // Assumes there's at least 1 value in the slice, which is an invariant of the entry level function.
-fn compute_is_constant<T: NativePType>(values: &[T]) -> bool {
+fn compute_is_constant<T: NativePType, const WIDTH: usize>(values: &[T]) -> bool {
     let first_value = values[0];
+    let first_vec = &[first_value; WIDTH];
+
+    let mut chunks = values[1..].array_chunks::<WIDTH>();
+    for chunk in &mut chunks {
+        if first_vec != chunk {
+            return false;
+        }
+    }
 
-    for value in &values[1..] {
+    for value in chunks.remainder() {
         if !value.is_eq(first_value) {
             return false;
         }
     }
 
     true
 }
+
+trait EqFloat {
+    type IntType;
+}
+
+impl EqFloat for f16 {
+    type IntType = u16;
+}
+impl EqFloat for f32 {
+    type IntType = u32;
+}
+impl EqFloat for f64 {
+    type IntType = u64;
+}

vortex-array/src/arrays/primitive/compute/is_constant.rs

@@ -2,6 +2,7 @@
 #![feature(portable_simd)]
 #![feature(substr_range)]
 #![feature(trusted_len)]
+#![feature(array_chunks)]
 //! Vortex crate containing core logic for encoding and memory representation of [arrays](ArrayRef).
 //!
 //! At the heart of Vortex are [arrays](ArrayRef) and [encodings](vtable::EncodingVTable).

vortex-array/src/lib.rs

@@ -116,14 +116,17 @@ macro_rules! native_ptype {
         impl NativePType for $T {
             const PTYPE: PType = PType::$ptype;
 
+            #[inline]
             fn is_nan(self) -> bool {
                 false
             }
 
+            #[inline]
             fn total_compare(self, other: Self) -> Ordering {
                 self.cmp(&other)
             }
 
+            #[inline]
             fn is_eq(self, other: Self) -> bool {
                 self == other
             }
@@ -136,14 +139,17 @@ macro_rules! native_float_ptype {
         impl NativePType for $T {
             const PTYPE: PType = PType::$ptype;
 
+            #[inline]
             fn is_nan(self) -> bool {
                 <$T>::is_nan(self)
             }
 
+            #[inline]
             fn total_compare(self, other: Self) -> Ordering {
                 self.total_cmp(&other)
             }
 
+            #[inline]
             fn is_eq(self, other: Self) -> bool {
                 self.to_bits() == other.to_bits()
             }