perf(in_list): add PrimitiveFilter for integer types

geoffreyclaude · geoffreyclaude · commit 88ad0f83ee58 · 2025-12-19T10:56:45.000+01:00
Adds PrimitiveFilter&lt;T&gt; - a generic HashSet-based filter for primitive types.
Also adds contains_slice() method for zero-copy buffer access used by
type reinterpretation.

The strategy module now uses PrimitiveFilter directly for Int32/Int64/UInt32/UInt64
instead of the macro-generated filters, while keeping Float32/Float64 with
OrderedFloat wrappers for now.
diff --git a/datafusion/physical-expr/src/expressions/in_list/primitive.rs b/datafusion/physical-expr/src/expressions/in_list/primitive.rs
@@ -19,9 +19,11 @@
 //!
 //! This module provides high-performance membership testing for Arrow primitive types.
 
+use std::hash::Hash;
+
 use arrow::array::{Array, ArrayRef, AsArray, BooleanArray};
 use arrow::datatypes::ArrowPrimitiveType;
-use datafusion_common::{Result, exec_datafusion_err};
+use datafusion_common::{HashSet, Result, exec_datafusion_err};
 
 use super::array_filter::StaticFilter;
 use super::result::{build_in_list_result, handle_dictionary};
@@ -171,3 +173,78 @@ impl<C: BitmapFilterConfig> StaticFilter for BitmapFilter<C> {
         ))
     }
 }
+
+// =============================================================================
+// PRIMITIVE FILTER (Hash-based)
+// =============================================================================
+
+/// Hash-based filter for primitive types with larger IN lists.
+pub(crate) struct PrimitiveFilter<T: ArrowPrimitiveType> {
+    null_count: usize,
+    set: HashSet<T::Native>,
+}
+
+impl<T: ArrowPrimitiveType> PrimitiveFilter<T>
+where
+    T::Native: Hash + Eq,
+{
+    pub(crate) fn try_new(in_array: &ArrayRef) -> Result<Self> {
+        let arr = in_array.as_primitive_opt::<T>().ok_or_else(|| {
+            exec_datafusion_err!(
+                "PrimitiveFilter: expected {} array",
+                std::any::type_name::<T>()
+            )
+        })?;
+        Ok(Self {
+            null_count: arr.null_count(),
+            set: arr.iter().flatten().collect(),
+        })
+    }
+
+    /// Check membership using a raw values slice (zero-copy path for type reinterpretation).
+    #[inline]
+    pub(crate) fn contains_slice(
+        &self,
+        values: &[T::Native],
+        nulls: Option<&arrow::buffer::NullBuffer>,
+        negated: bool,
+    ) -> BooleanArray {
+        build_in_list_result(
+            values.len(),
+            nulls,
+            self.null_count > 0,
+            negated,
+            // SAFETY: i is in bounds since we iterate 0..values.len()
+            |i| self.set.contains(unsafe { values.get_unchecked(i) }),
+        )
+    }
+}
+
+impl<T> StaticFilter for PrimitiveFilter<T>
+where
+    T: ArrowPrimitiveType + 'static,
+    T::Native: Hash + Eq + Send + Sync + 'static,
+{
+    fn null_count(&self) -> usize {
+        self.null_count
+    }
+
+    fn contains(&self, v: &dyn Array, negated: bool) -> Result<BooleanArray> {
+        handle_dictionary!(self, v, negated);
+        let v = v.as_primitive_opt::<T>().ok_or_else(|| {
+            exec_datafusion_err!(
+                "PrimitiveFilter: expected {} array",
+                std::any::type_name::<T>()
+            )
+        })?;
+        let values = v.values();
+        Ok(build_in_list_result(
+            v.len(),
+            v.nulls(),
+            self.null_count > 0,
+            negated,
+            // SAFETY: i is in bounds since we iterate 0..v.len()
+            |i| self.set.contains(unsafe { values.get_unchecked(i) }),
+        ))
+    }
+}
diff --git a/datafusion/physical-expr/src/expressions/in_list/strategy.rs b/datafusion/physical-expr/src/expressions/in_list/strategy.rs
@@ -27,7 +27,7 @@ use arrow::datatypes::*;
 use datafusion_common::{HashSet, Result, exec_datafusion_err};
 
 use super::array_filter::{ArrayStaticFilter, StaticFilter};
-use super::primitive::{U8Config, U16Config};
+use super::primitive::{PrimitiveFilter, U8Config, U16Config};
 use super::transform::make_bitmap_filter;
 
 pub(crate) fn instantiate_static_filter(
@@ -39,11 +39,11 @@ pub(crate) fn instantiate_static_filter(
         // 2-byte types: use bitmap (65536 bits = 8 KB)
         DataType::Int16 | DataType::UInt16 => make_bitmap_filter::<U16Config>(&in_array),
         // 4-byte integer types
-        DataType::Int32 => Ok(Arc::new(Int32StaticFilter::try_new(&in_array)?)),
-        DataType::UInt32 => Ok(Arc::new(UInt32StaticFilter::try_new(&in_array)?)),
+        DataType::Int32 => Ok(Arc::new(PrimitiveFilter::<Int32Type>::try_new(&in_array)?)),
+        DataType::UInt32 => Ok(Arc::new(PrimitiveFilter::<UInt32Type>::try_new(&in_array)?)),
         // 8-byte integer types
-        DataType::Int64 => Ok(Arc::new(Int64StaticFilter::try_new(&in_array)?)),
-        DataType::UInt64 => Ok(Arc::new(UInt64StaticFilter::try_new(&in_array)?)),
+        DataType::Int64 => Ok(Arc::new(PrimitiveFilter::<Int64Type>::try_new(&in_array)?)),
+        DataType::UInt64 => Ok(Arc::new(PrimitiveFilter::<UInt64Type>::try_new(&in_array)?)),
         // Float primitive types (use ordered wrappers for Hash/Eq)
         DataType::Float32 => Ok(Arc::new(Float32StaticFilter::try_new(&in_array)?)),
         DataType::Float64 => Ok(Arc::new(Float64StaticFilter::try_new(&in_array)?)),
@@ -104,139 +104,6 @@ impl From<f64> for OrderedFloat64 {
     }
 }
 
-// Macro to generate specialized StaticFilter implementations for primitive types
-macro_rules! primitive_static_filter {
-    ($Name:ident, $ArrowType:ty) => {
-        struct $Name {
-            null_count: usize,
-            values: HashSet<<$ArrowType as ArrowPrimitiveType>::Native>,
-        }
-
-        impl $Name {
-            fn try_new(in_array: &ArrayRef) -> Result<Self> {
-                let in_array = in_array
-                    .as_primitive_opt::<$ArrowType>()
-                    .ok_or_else(|| exec_datafusion_err!("Failed to downcast an array to a '{}' array", stringify!($ArrowType)))?;
-
-                let mut values = HashSet::with_capacity(in_array.len());
-                let null_count = in_array.null_count();
-
-                for v in in_array.iter().flatten() {
-                    values.insert(v);
-                }
-
-                Ok(Self { null_count, values })
-            }
-        }
-
-        impl StaticFilter for $Name {
-            fn null_count(&self) -> usize {
-                self.null_count
-            }
-
-            fn contains(&self, v: &dyn Array, negated: bool) -> Result<BooleanArray> {
-                // Handle dictionary arrays by recursing on the values
-                downcast_dictionary_array! {
-                    v => {
-                        let values_contains = self.contains(v.values().as_ref(), negated)?;
-                        let result = take(&values_contains, v.keys(), None)?;
-                        return Ok(downcast_array(result.as_ref()))
-                    }
-                    _ => {}
-                }
-
-                let v = v
-                    .as_primitive_opt::<$ArrowType>()
-                    .ok_or_else(|| exec_datafusion_err!("Failed to downcast an array to a '{}' array", stringify!($ArrowType)))?;
-
-                let haystack_has_nulls = self.null_count > 0;
-
-                let needle_values = v.values();
-                let needle_nulls = v.nulls();
-                let needle_has_nulls = v.null_count() > 0;
-
-                // Truth table for `value [NOT] IN (set)` with SQL three-valued logic:
-                // ("-" means the value doesn't affect the result)
-                //
-                // | needle_null | haystack_null | negated | in set? | result |
-                // |-------------|---------------|---------|---------|--------|
-                // | true        | -             | false   | -       | null   |
-                // | true        | -             | true    | -       | null   |
-                // | false       | true          | false   | yes     | true   |
-                // | false       | true          | false   | no      | null   |
-                // | false       | true          | true    | yes     | false  |
-                // | false       | true          | true    | no      | null   |
-                // | false       | false         | false   | yes     | true   |
-                // | false       | false         | false   | no      | false  |
-                // | false       | false         | true    | yes     | false  |
-                // | false       | false         | true    | no      | true   |
-
-                // Compute the "contains" result using collect_bool (fast batched approach)
-                // This ignores nulls - we handle them separately
-                let contains_buffer = if negated {
-                    BooleanBuffer::collect_bool(needle_values.len(), |i| {
-                        !self.values.contains(&needle_values[i])
-                    })
-                } else {
-                    BooleanBuffer::collect_bool(needle_values.len(), |i| {
-                        self.values.contains(&needle_values[i])
-                    })
-                };
-
-                // Compute the null mask
-                // Output is null when:
-                // 1. needle value is null, OR
-                // 2. needle value is not in set AND haystack has nulls
-                let result_nulls = match (needle_has_nulls, haystack_has_nulls) {
-                    (false, false) => {
-                        // No nulls anywhere
-                        None
-                    }
-                    (true, false) => {
-                        // Only needle has nulls - just use needle's null mask
-                        needle_nulls.cloned()
-                    }
-                    (false, true) => {
-                        // Only haystack has nulls - result is null when value not in set
-                        // Valid (not null) when original "in set" is true
-                        // For NOT IN: contains_buffer = !original, so validity = !contains_buffer
-                        let validity = if negated {
-                            !&contains_buffer
-                        } else {
-                            contains_buffer.clone()
-                        };
-                        Some(NullBuffer::new(validity))
-                    }
-                    (true, true) => {
-                        // Both have nulls - combine needle nulls with haystack-induced nulls
-                        let needle_validity = needle_nulls.map(|n| n.inner().clone())
-                            .unwrap_or_else(|| BooleanBuffer::new_set(needle_values.len()));
-
-                        // Valid when original "in set" is true (see above)
-                        let haystack_validity = if negated {
-                            !&contains_buffer
-                        } else {
-                            contains_buffer.clone()
-                        };
-
-                        // Combined validity: valid only where both are valid
-                        let combined_validity = &needle_validity & &haystack_validity;
-                        Some(NullBuffer::new(combined_validity))
-                    }
-                };
-
-                Ok(BooleanArray::new(contains_buffer, result_nulls))
-            }
-        }
-    };
-}
-
-// Generate specialized filters for 4-byte and 8-byte integer primitive types
-// (1-byte and 2-byte types use BitmapFilter instead)
-primitive_static_filter!(Int32StaticFilter, Int32Type);
-primitive_static_filter!(Int64StaticFilter, Int64Type);
-primitive_static_filter!(UInt32StaticFilter, UInt32Type);
-primitive_static_filter!(UInt64StaticFilter, UInt64Type);
 
 // Macro to generate specialized StaticFilter implementations for float types
 // Floats require a wrapper type (OrderedFloat*) to implement Hash/Eq due to NaN semantics
diff --git a/datafusion/physical-expr/src/expressions/in_list/transform.rs b/datafusion/physical-expr/src/expressions/in_list/transform.rs
@@ -20,10 +20,13 @@
 //! This module provides type reinterpretation for optimizing filter dispatch.
 //! For equality comparison, only the bit pattern matters, so we can:
 //! - Reinterpret signed integers as unsigned (Int8 → UInt8)
+//! - Reinterpret floats as unsigned integers (Float64 → UInt64)
 //!
-//! This allows using a single filter implementation (e.g., BitmapFilter for UInt8)
-//! to handle multiple types (Int8, UInt8) that share the same byte width.
+//! This allows using a single filter implementation (e.g., for UInt64) to handle
+//! multiple types (Int64, Float64, Timestamp, Duration) that share the same
+//! byte width, reducing code duplication.
 
+use std::hash::Hash;
 use std::sync::Arc;
 
 use arrow::array::{Array, ArrayRef, BooleanArray, PrimitiveArray};
@@ -32,7 +35,7 @@ use arrow::datatypes::ArrowPrimitiveType;
 use datafusion_common::Result;
 
 use super::array_filter::StaticFilter;
-use super::primitive::{BitmapFilter, BitmapFilterConfig};
+use super::primitive::{BitmapFilter, BitmapFilterConfig, PrimitiveFilter};
 use super::result::handle_dictionary;
 
 // =============================================================================
@@ -59,6 +62,32 @@ impl<C: BitmapFilterConfig> StaticFilter for ReinterpretedBitmap<C> {
     }
 }
 
+/// Reinterpreting filter for hash-based lookups.
+///
+/// Zero-copy: reinterprets input buffer directly as target type slice.
+struct ReinterpretedPrimitive<D: ArrowPrimitiveType> {
+    inner: PrimitiveFilter<D>,
+}
+
+impl<D> StaticFilter for ReinterpretedPrimitive<D>
+where
+    D: ArrowPrimitiveType + 'static,
+    D::Native: Hash + Eq + Send + Sync + 'static,
+{
+    fn null_count(&self) -> usize {
+        self.inner.null_count()
+    }
+
+    fn contains(&self, v: &dyn Array, negated: bool) -> Result<BooleanArray> {
+        handle_dictionary!(self, v, negated);
+
+        let data = v.to_data();
+        let values: &[D::Native] = data.buffers()[0].typed_data();
+
+        Ok(self.inner.contains_slice(values, data.nulls(), negated))
+    }
+}
+
 /// Reinterprets any primitive-like array as the target primitive type T by extracting
 /// the underlying buffer.
 ///
@@ -87,3 +116,20 @@ where
     let inner = BitmapFilter::<C>::try_new(&reinterpreted)?;
     Ok(Arc::new(ReinterpretedBitmap { inner }))
 }
+
+/// Creates a hash-based filter, reinterpreting types if needed.
+pub(crate) fn make_primitive_filter<D>(
+    in_array: &ArrayRef,
+) -> Result<Arc<dyn StaticFilter + Send + Sync>>
+where
+    D: ArrowPrimitiveType + 'static,
+    D::Native: Hash + Eq + Send + Sync + 'static,
+{
+    if in_array.data_type() == &D::DATA_TYPE {
+        return Ok(Arc::new(PrimitiveFilter::<D>::try_new(in_array)?));
+    }
+
+    let reinterpreted = reinterpret_any_primitive_to::<D>(in_array.as_ref());
+    let inner = PrimitiveFilter::<D>::try_new(&reinterpreted)?;
+    Ok(Arc::new(ReinterpretedPrimitive { inner }))
+}
