diff --git a/datafusion/physical-expr/benches/in_list_strategy.rs b/datafusion/physical-expr/benches/in_list_strategy.rs index c69af192b9cd..762c2c97a111 100644 --- a/datafusion/physical-expr/benches/in_list_strategy.rs +++ b/datafusion/physical-expr/benches/in_list_strategy.rs @@ -38,6 +38,7 @@ //! | 32-bit primitive cases | Int32, Float32 | small and large lists | 4, 32, 64, 256 | //! | 64-bit primitive cases | Int64, TimestampNs | small and large lists | 4, 16, 32, 128 | //! | 128-bit interval cases | IntervalMonthDayNano | small lists | 4 | +//! | Decimal128 cases | Decimal128 | larger lists | 5, 64 | //! | Utf8 short-string cases | Utf8 | 8-byte strings | 4, 64, 256 | //! | Utf8 long-string cases | Utf8 | 24-byte strings | 4, 64, 256 | //! | Utf8View short-string cases | Utf8View | 8-byte strings | 4, 16, 64, 256 | @@ -463,6 +464,23 @@ fn bench_primitive(c: &mut Criterion) { } } + // Decimal128: benchmark the first hash-set list size (5) and a larger list (64). + for list_size in [5, 64] { + for match_pct in MATCH_RATES { + bench_numeric::( + c, + "primitive", + &format!("decimal128/large_list/list={list_size}/match={match_pct}%"), + &NumericBenchConfig::new( + list_size, + match_pct as f64 / 100.0, + |rng| i128::from(rng.random::()), + |v| ScalarValue::Decimal128(Some(v), 38, 10), + ), + ); + } + } + // NOT IN benchmark: test negated path bench_numeric::( c, diff --git a/datafusion/physical-expr/src/expressions/in_list.rs b/datafusion/physical-expr/src/expressions/in_list.rs index 874e149b5832..0fb978cd0baf 100644 --- a/datafusion/physical-expr/src/expressions/in_list.rs +++ b/datafusion/physical-expr/src/expressions/in_list.rs @@ -43,7 +43,7 @@ mod result; mod static_filter; mod strategy; -use static_filter::StaticFilter; +use static_filter::StaticFilterRef; use strategy::instantiate_static_filter; /// InList @@ -51,7 +51,7 @@ pub struct InListExpr { expr: Arc, list: Vec>, negated: bool, - static_filter: Option>, + static_filter: Option, } impl Debug for InListExpr { @@ -148,7 +148,7 @@ impl InListExpr { expr: Arc, list: Vec>, negated: bool, - static_filter: Option>, + static_filter: Option, ) -> Self { Self { expr, @@ -222,8 +222,8 @@ impl InListExpr { /// Create a new InList expression, using a static filter when possible. /// /// This validates data types and attempts to create a static filter for constant - /// list expressions. Uses specialized StaticFilter implementations for better - /// performance (e.g., Int32StaticFilter for Int32). + /// list expressions. Uses specialized branchless, bitmap, or hash-set filters + /// when the list's physical representation supports them. /// /// Returns an error if data types don't match. If the list contains non-constant /// expressions, falls back to dynamic evaluation at runtime. @@ -2592,7 +2592,7 @@ mod tests { // Create IN list with Int32 literals: (100, 200, 300) let list = vec![lit(100i32), lit(200i32), lit(300i32)]; - // Create InListExpr via in_list() - this uses Int32StaticFilter for Int32 lists + // Create InListExpr via in_list(), which selects a primitive static filter. let expr = in_list(col_a, list, &false, &schema)?; // Create dictionary-encoded batch with values [100, 200, 500] diff --git a/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs b/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs index cd0cbd0de59a..8539dc7956dc 100644 --- a/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs @@ -42,7 +42,7 @@ //! different NaN values. [`BranchlessFilterType`] defines these safe, //! same-sized mappings and checks their sizes at compile time. //! -//! The fast path is intentionally limited to short lists: +//! The fast path is limited to short lists: //! //! - 16 values for 1-byte types //! - 8 values for 2-byte types diff --git a/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs b/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs index 8f8d9bad04af..bfafae31ed62 100644 --- a/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs @@ -20,21 +20,133 @@ //! This module provides membership tests for Arrow primitive types. use arrow::array::{Array, ArrayRef, AsArray, BooleanArray}; -use arrow::buffer::{BooleanBuffer, NullBuffer}; use arrow::datatypes::*; use arrow::util::bit_iterator::BitIndexIterator; use datafusion_common::{HashSet, Result, exec_datafusion_err}; use std::hash::{Hash, Hasher}; +use std::marker::PhantomData; +use std::sync::Arc; +use super::branchless_filter::{BranchlessFilter, BranchlessFilterType}; use super::result::build_in_list_result; -use super::static_filter::{StaticFilter, handle_dictionary}; +use super::static_filter::{StaticFilter, StaticFilterRef, handle_dictionary}; + +/// Selects an optimized filter for a primitive representation. +/// +/// Supported short lists use a branchless filter. Larger supported lists use a +/// bitmap or hash-set filter. Returns `None` for primitive types without a +/// specialized filter. Representation adapters call this after conversion and +/// use the same cutoffs and fallback policy as native primitive arrays. +pub(super) fn instantiate_primitive_filter( + in_array: &ArrayRef, +) -> Result> { + if let Some(filter) = instantiate_branchless_filter(in_array)? { + return Ok(Some(filter)); + } + + let filter: StaticFilterRef = match in_array.data_type() { + DataType::Int8 => Arc::new(BitmapFilter::::try_new(in_array)?), + DataType::UInt8 => Arc::new(BitmapFilter::::try_new(in_array)?), + DataType::Int16 => Arc::new(BitmapFilter::::try_new(in_array)?), + DataType::UInt16 => Arc::new(BitmapFilter::::try_new(in_array)?), + DataType::Float16 => Arc::new(BitmapFilter::::try_new(in_array)?), + DataType::Int32 => { + Arc::new(PrimitiveHashSetFilter::::try_new(in_array)?) + } + DataType::Int64 => { + Arc::new(PrimitiveHashSetFilter::::try_new(in_array)?) + } + DataType::UInt32 => { + Arc::new(PrimitiveHashSetFilter::::try_new(in_array)?) + } + DataType::UInt64 => { + Arc::new(PrimitiveHashSetFilter::::try_new(in_array)?) + } + DataType::Decimal128(_, _) => { + Arc::new(PrimitiveHashSetFilter::::try_new(in_array)?) + } + // Float primitive types use ordered wrapper keys for Hash/Eq. + DataType::Float32 => Arc::new(PrimitiveHashSetFilter::< + Float32Type, + OrderedFloat32, + >::try_new(in_array)?), + DataType::Float64 => Arc::new(PrimitiveHashSetFilter::< + Float64Type, + OrderedFloat64, + >::try_new(in_array)?), + _ => return Ok(None), + }; + + Ok(Some(filter)) +} + +fn instantiate_branchless_filter(in_array: &ArrayRef) -> Result> { + let non_null_count = in_array.len() - in_array.null_count(); + + macro_rules! branchless { + ($arrow_type:ty) => {{ + // Larger lists use the standard primitive filter. `try_new` + // checks the limit again when called directly. + if non_null_count > <$arrow_type as BranchlessFilterType>::MAX_LIST_LEN { + Ok(None) + } else { + let filter: StaticFilterRef = + Arc::new(BranchlessFilter::<$arrow_type>::try_new(in_array)?); + Ok(Some(filter)) + } + }}; + } + + match in_array.data_type() { + DataType::Int8 => branchless!(Int8Type), + DataType::UInt8 => branchless!(UInt8Type), + DataType::Int16 => branchless!(Int16Type), + DataType::UInt16 => branchless!(UInt16Type), + DataType::Float16 => branchless!(Float16Type), + DataType::Int32 => branchless!(Int32Type), + DataType::UInt32 => branchless!(UInt32Type), + DataType::Float32 => branchless!(Float32Type), + DataType::Date32 => branchless!(Date32Type), + DataType::Time32(unit) => match unit { + TimeUnit::Second => branchless!(Time32SecondType), + TimeUnit::Millisecond => branchless!(Time32MillisecondType), + _ => Ok(None), + }, + DataType::Int64 => branchless!(Int64Type), + DataType::UInt64 => branchless!(UInt64Type), + DataType::Float64 => branchless!(Float64Type), + DataType::Date64 => branchless!(Date64Type), + DataType::Time64(unit) => match unit { + TimeUnit::Microsecond => branchless!(Time64MicrosecondType), + TimeUnit::Nanosecond => branchless!(Time64NanosecondType), + _ => Ok(None), + }, + DataType::Timestamp(unit, _) => match unit { + TimeUnit::Second => branchless!(TimestampSecondType), + TimeUnit::Millisecond => branchless!(TimestampMillisecondType), + TimeUnit::Microsecond => branchless!(TimestampMicrosecondType), + TimeUnit::Nanosecond => branchless!(TimestampNanosecondType), + }, + DataType::Duration(unit) => match unit { + TimeUnit::Second => branchless!(DurationSecondType), + TimeUnit::Millisecond => branchless!(DurationMillisecondType), + TimeUnit::Microsecond => branchless!(DurationMicrosecondType), + TimeUnit::Nanosecond => branchless!(DurationNanosecondType), + }, + DataType::Decimal128(_, _) => branchless!(Decimal128Type), + DataType::Interval(IntervalUnit::MonthDayNano) => { + branchless!(IntervalMonthDayNanoType) + } + _ => Ok(None), + } +} /// Storage for the bits used by [`BitmapFilter`]. /// /// `BitmapFilter` represents an `IN` list with one bit for each possible /// value, so membership checks become direct bit tests. This trait lets the /// same filter code use different storage sizes for different integer widths. -pub(super) trait BitmapStorage: Send + Sync { +trait BitmapStorage: Send + Sync { fn new_zeroed() -> Self; fn set_bit(&mut self, index: usize); fn get_bit(&self, index: usize) -> bool; @@ -81,9 +193,7 @@ impl BitmapStorage for Box<[u64; 1024]> { /// supplies the bitmap storage size and maps values to their bit-pattern index /// for the primitive domains that are small enough to represent with one bit /// per possible value. -pub(super) trait BitmapFilterType: - ArrowPrimitiveType + Send + Sync + 'static -{ +trait BitmapFilterType: ArrowPrimitiveType + Send + Sync + 'static { type Storage: BitmapStorage; /// Returns the index in the bitmap to check for this value. @@ -151,7 +261,7 @@ impl BitmapFilterType for Float16Type { /// the bit selected by each value. Evaluating input values checks the same bit /// position. Null handling and `NOT IN` inversion are handled by /// `build_in_list_result`. -pub(super) struct BitmapFilter { +struct BitmapFilter { null_count: usize, bits: T::Storage, } @@ -160,7 +270,7 @@ impl BitmapFilter where T: BitmapFilterType, { - pub(super) fn try_new(in_array: &ArrayRef) -> Result { + fn try_new(in_array: &ArrayRef) -> Result { let prim_array = in_array.as_primitive_opt::().ok_or_else(|| { exec_datafusion_err!("BitmapFilter: expected {} array", T::DATA_TYPE) })?; @@ -272,157 +382,77 @@ impl From for OrderedFloat64 { } } -// Macro to generate specialized StaticFilter implementations for primitive types -macro_rules! primitive_static_filter { - ($Name:ident, $ArrowType:ty) => { - primitive_static_filter!( - $Name, - $ArrowType, - <$ArrowType as ArrowPrimitiveType>::Native, - |v| v - ); - }; - ($Name:ident, $ArrowType:ty, $SetValueType:ty, $to_set_value:expr) => { - pub(super) struct $Name { - null_count: usize, - values: HashSet<$SetValueType>, - } - - impl $Name { - pub(super) fn try_new(in_array: &ArrayRef) -> Result { - 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(($to_set_value)(v)); - } +/// Hash-set membership for primitive types. +/// +/// `K` defaults to the Arrow type's native value. Floats use an ordered wrapper +/// key because their native values do not implement [`Eq`] and [`Hash`]. +struct PrimitiveHashSetFilter< + T: ArrowPrimitiveType, + K = ::Native, +> { + null_count: usize, + values: HashSet, + _marker: PhantomData, +} - Ok(Self { null_count, values }) - } +impl PrimitiveHashSetFilter +where + T: ArrowPrimitiveType, + T::Native: Copy, + K: From + Eq + Hash, +{ + fn try_new(in_array: &ArrayRef) -> Result { + let in_array = in_array.as_primitive_opt::().ok_or_else(|| { + exec_datafusion_err!( + "PrimitiveHashSetFilter: expected {} array", + T::DATA_TYPE + ) + })?; + let mut values = HashSet::with_capacity(in_array.len() - in_array.null_count()); + for value in in_array.iter().flatten() { + values.insert(K::from(value)); } - impl StaticFilter for $Name { - fn null_count(&self) -> usize { - self.null_count - } - - fn contains(&self, v: &dyn Array, negated: bool) -> Result { - handle_dictionary!(self, v, negated); - - 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(&($to_set_value)(needle_values[i])) - }) - } else { - BooleanBuffer::collect_bool(needle_values.len(), |i| { - self.values.contains(&($to_set_value)(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)) - } - } - }; + Ok(Self { + null_count: in_array.null_count(), + values, + _marker: PhantomData, + }) + } } -primitive_static_filter!(Int32StaticFilter, Int32Type); -primitive_static_filter!(Int64StaticFilter, Int64Type); -primitive_static_filter!(UInt32StaticFilter, UInt32Type); -primitive_static_filter!(UInt64StaticFilter, UInt64Type); +impl StaticFilter for PrimitiveHashSetFilter +where + T: ArrowPrimitiveType + Send + Sync + 'static, + T::Native: Copy + Send + Sync, + K: From + Eq + Hash + Send + Sync + 'static, +{ + fn null_count(&self) -> usize { + self.null_count + } -// Macro to generate specialized StaticFilter implementations for float types -// Floats require a wrapper type (OrderedFloat*) to implement Hash/Eq due to NaN semantics -macro_rules! float_static_filter { - ($Name:ident, $ArrowType:ty, $OrderedType:ty) => { - primitive_static_filter!($Name, $ArrowType, $OrderedType, <$OrderedType>::from); - }; -} + fn contains(&self, v: &dyn Array, negated: bool) -> Result { + handle_dictionary!(self, v, negated); -// Generate specialized filters for float types using ordered wrappers -float_static_filter!(Float32StaticFilter, Float32Type, OrderedFloat32); -float_static_filter!(Float64StaticFilter, Float64Type, OrderedFloat64); + let v = v.as_primitive_opt::().ok_or_else(|| { + exec_datafusion_err!( + "PrimitiveHashSetFilter: expected {} array", + T::DATA_TYPE + ) + })?; + let input_values = v.values(); + Ok(build_in_list_result( + v.len(), + v.nulls(), + self.null_count > 0, + negated, + |index| { + let key = K::from(input_values[index]); + self.values.contains(&key) + }, + )) + } +} #[cfg(test)] mod tests { @@ -430,10 +460,15 @@ mod tests { use std::sync::Arc; use arrow::array::{ - DictionaryArray, Float16Array, Int8Array, Int16Array, UInt8Array, UInt16Array, + DictionaryArray, Float16Array, Float32Array, Float64Array, Int8Array, Int16Array, + UInt8Array, UInt16Array, UInt32Array, }; use half::f16; + fn uint32_array(values: Vec>) -> ArrayRef { + Arc::new(UInt32Array::from(values)) + } + fn assert_contains( filter: &dyn StaticFilter, needles: &dyn Array, @@ -446,6 +481,60 @@ mod tests { Ok(()) } + #[test] + fn branchless_routing_respects_max_list_len() -> Result<()> { + let max_len = ::MAX_LIST_LEN; + + let values = (0..max_len) + .map(|value| Some(value as u32)) + .collect::>(); + assert!(instantiate_branchless_filter(&uint32_array(values))?.is_some()); + + let values = (0..=max_len) + .map(|value| Some(value as u32)) + .collect::>(); + assert!(instantiate_branchless_filter(&uint32_array(values))?.is_none()); + + Ok(()) + } + + #[test] + fn branchless_routing_handles_zero_non_null_values() -> Result<()> { + let array = uint32_array(vec![None; 3]); + + assert!(instantiate_branchless_filter(&array)?.is_some()); + + Ok(()) + } + + #[test] + fn primitive_hash_filter_handles_float_keys() -> Result<()> { + let nan32 = f32::NAN; + let other_nan32 = f32::from_bits(nan32.to_bits() + 1); + let haystack: ArrayRef = Arc::new(Float32Array::from(vec![0.0, nan32])); + let filter = + PrimitiveHashSetFilter::::try_new(&haystack)?; + let needles = Float32Array::from(vec![ + Some(0.0), + Some(-0.0), + Some(nan32), + Some(other_nan32), + None, + ]); + assert_contains( + &filter, + &needles, + vec![Some(true), Some(false), Some(true), Some(false), None], + )?; + + let nan64 = f64::NAN; + let haystack: ArrayRef = Arc::new(Float64Array::from(vec![1.0, nan64])); + let filter = + PrimitiveHashSetFilter::::try_new(&haystack)?; + let needles = Float64Array::from(vec![Some(1.0), Some(nan64), Some(2.0)]); + assert_contains(&filter, &needles, vec![Some(true), Some(true), Some(false)]) + } + #[test] fn bitmap_filter_u8_handles_nulls() -> Result<()> { let haystack: ArrayRef = Arc::new(UInt8Array::from(vec![Some(1), None, Some(3)])); diff --git a/datafusion/physical-expr/src/expressions/in_list/result.rs b/datafusion/physical-expr/src/expressions/in_list/result.rs index 3ebdbfe19f74..1048963b6350 100644 --- a/datafusion/physical-expr/src/expressions/in_list/result.rs +++ b/datafusion/physical-expr/src/expressions/in_list/result.rs @@ -24,15 +24,21 @@ use arrow::array::BooleanArray; use arrow::buffer::{BooleanBuffer, NullBuffer}; -// Truth table for (needle_nulls, haystack_has_nulls, negated): -// (Some, true, false) => values: valid & contains, nulls: valid & contains -// (None, true, false) => values: contains, nulls: contains -// (Some, true, true) => values: valid & !contains, nulls: valid & contains -// (None, true, true) => values: !contains, nulls: contains -// (Some, false, false) => values: valid & contains, nulls: valid -// (Some, false, true) => values: valid & !contains, nulls: valid -// (None, false, false) => values: contains, nulls: none -// (None, false, true) => values: !contains, nulls: none +// Truth table for `value [NOT] IN (set)` with SQL three-valued logic: +// ("-" means the value does not affect the result) +// +// | needle null | set has null | negated | found in set | result | +// |-------------|--------------|---------|--------------|--------| +// | true | - | false | - | null | +// | true | - | true | - | null | +// | false | true | false | true | true | +// | false | true | false | false | null | +// | false | true | true | true | false | +// | false | true | true | false | null | +// | false | false | false | true | true | +// | false | false | false | false | false | +// | false | false | true | true | false | +// | false | false | true | false | true | /// Builds a BooleanArray result for IN list operations. /// @@ -44,7 +50,7 @@ use arrow::buffer::{BooleanBuffer, NullBuffer}; /// This version computes contains for all positions, including nulls, then applies /// null masking via bitmap operations. #[inline] -pub(crate) fn build_in_list_result( +pub(super) fn build_in_list_result( len: usize, needle_nulls: Option<&NullBuffer>, haystack_has_nulls: bool, @@ -63,7 +69,7 @@ where /// This version does not assume contains_buf is pre-masked at null positions. /// It handles nulls using bitmap operations. #[inline] -pub(crate) fn build_result_from_contains( +pub(super) fn build_result_from_contains( needle_nulls: Option<&NullBuffer>, haystack_has_nulls: bool, negated: bool, diff --git a/datafusion/physical-expr/src/expressions/in_list/static_filter.rs b/datafusion/physical-expr/src/expressions/in_list/static_filter.rs index 3c964d418347..e74ea9131e63 100644 --- a/datafusion/physical-expr/src/expressions/in_list/static_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/static_filter.rs @@ -15,9 +15,13 @@ // specific language governing permissions and limitations // under the License. +use std::sync::Arc; + use arrow::array::{Array, BooleanArray}; use datafusion_common::Result; +pub(super) type StaticFilterRef = Arc; + /// Trait for InList static filters. /// /// Static filters store a pre-computed set of values (the haystack) and check diff --git a/datafusion/physical-expr/src/expressions/in_list/strategy.rs b/datafusion/physical-expr/src/expressions/in_list/strategy.rs index d5ca8154a92f..d008217ee19f 100644 --- a/datafusion/physical-expr/src/expressions/in_list/strategy.rs +++ b/datafusion/physical-expr/src/expressions/in_list/strategy.rs @@ -19,179 +19,29 @@ use std::sync::Arc; use arrow::array::ArrayRef; use arrow::compute::cast; -use arrow::datatypes::{ - DataType, Date32Type, Date64Type, Decimal128Type, DurationMicrosecondType, - DurationMillisecondType, DurationNanosecondType, DurationSecondType, Float16Type, - Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, - IntervalMonthDayNanoType, IntervalUnit, Time32MillisecondType, Time32SecondType, - Time64MicrosecondType, Time64NanosecondType, TimeUnit, TimestampMicrosecondType, - TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, UInt8Type, - UInt16Type, UInt32Type, UInt64Type, -}; +use arrow::datatypes::DataType; use datafusion_common::Result; use super::array_static_filter::ArrayStaticFilter; -use super::branchless_filter::{ - BranchlessFilter, BranchlessFilterType, BranchlessNative, -}; -use super::primitive_filter::*; -use super::static_filter::StaticFilter; - -type StaticFilterRef = Arc; +use super::primitive_filter::instantiate_primitive_filter; +use super::static_filter::StaticFilterRef; pub(super) fn instantiate_static_filter(in_array: ArrayRef) -> Result { let in_array = flatten_dictionary_haystack(in_array)?; - if let Some(filter) = instantiate_branchless_filter(&in_array)? { + if let Some(filter) = instantiate_primitive_filter(&in_array)? { return Ok(filter); } - instantiate_standard_filter(in_array) + Ok(Arc::new(ArrayStaticFilter::try_new(in_array)?)) } fn flatten_dictionary_haystack(in_array: ArrayRef) -> Result { // Flatten dictionary-encoded haystacks to their value type so that - // specialized filters (e.g. Int32StaticFilter) are used instead of - // falling through to the generic ArrayStaticFilter. + // specialized primitive filters are used instead of falling through to the + // generic ArrayStaticFilter. match in_array.data_type() { DataType::Dictionary(_, value_type) => Ok(cast(&in_array, value_type.as_ref())?), _ => Ok(in_array), } } - -fn instantiate_branchless_filter(in_array: &ArrayRef) -> Result> { - let non_null_count = in_array.len() - in_array.null_count(); - - macro_rules! filter { - ($arrow_type:ty) => { - branchless_filter::<$arrow_type>(in_array, non_null_count) - }; - } - - match in_array.data_type() { - DataType::Int8 => filter!(Int8Type), - DataType::UInt8 => filter!(UInt8Type), - DataType::Int16 => filter!(Int16Type), - DataType::UInt16 => filter!(UInt16Type), - DataType::Float16 => filter!(Float16Type), - DataType::Int32 => filter!(Int32Type), - DataType::UInt32 => filter!(UInt32Type), - DataType::Float32 => filter!(Float32Type), - DataType::Date32 => filter!(Date32Type), - DataType::Time32(unit) => match unit { - TimeUnit::Second => filter!(Time32SecondType), - TimeUnit::Millisecond => filter!(Time32MillisecondType), - _ => Ok(None), - }, - DataType::Int64 => filter!(Int64Type), - DataType::UInt64 => filter!(UInt64Type), - DataType::Float64 => filter!(Float64Type), - DataType::Date64 => filter!(Date64Type), - DataType::Time64(unit) => match unit { - TimeUnit::Microsecond => filter!(Time64MicrosecondType), - TimeUnit::Nanosecond => filter!(Time64NanosecondType), - _ => Ok(None), - }, - DataType::Timestamp(unit, _) => match unit { - TimeUnit::Second => filter!(TimestampSecondType), - TimeUnit::Millisecond => filter!(TimestampMillisecondType), - TimeUnit::Microsecond => filter!(TimestampMicrosecondType), - TimeUnit::Nanosecond => filter!(TimestampNanosecondType), - }, - DataType::Duration(unit) => match unit { - TimeUnit::Second => filter!(DurationSecondType), - TimeUnit::Millisecond => filter!(DurationMillisecondType), - TimeUnit::Microsecond => filter!(DurationMicrosecondType), - TimeUnit::Nanosecond => filter!(DurationNanosecondType), - }, - DataType::Decimal128(_, _) => filter!(Decimal128Type), - DataType::Interval(IntervalUnit::MonthDayNano) => { - filter!(IntervalMonthDayNanoType) - } - _ => Ok(None), - } -} - -fn instantiate_standard_filter(in_array: ArrayRef) -> Result { - match in_array.data_type() { - DataType::Int8 => bitmap_filter::(&in_array), - DataType::UInt8 => bitmap_filter::(&in_array), - DataType::Int16 => bitmap_filter::(&in_array), - DataType::UInt16 => bitmap_filter::(&in_array), - DataType::Float16 => bitmap_filter::(&in_array), - DataType::Int32 => Ok(Arc::new(Int32StaticFilter::try_new(&in_array)?)), - DataType::Int64 => Ok(Arc::new(Int64StaticFilter::try_new(&in_array)?)), - DataType::UInt32 => Ok(Arc::new(UInt32StaticFilter::try_new(&in_array)?)), - DataType::UInt64 => Ok(Arc::new(UInt64StaticFilter::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)?)), - _ => { - // Fall through to generic implementation for unsupported types - // (Struct, etc.). - Ok(Arc::new(ArrayStaticFilter::try_new(in_array)?)) - } - } -} - -fn bitmap_filter(in_array: &ArrayRef) -> Result -where - T: BitmapFilterType, -{ - Ok(Arc::new(BitmapFilter::::try_new(in_array)?)) -} - -fn branchless_filter( - in_array: &ArrayRef, - non_null_count: usize, -) -> Result> -where - T: BranchlessFilterType, - BranchlessNative: Copy + PartialEq + Send + Sync, -{ - // Larger lists use the standard filter. `try_new` checks the limit again. - if non_null_count > T::MAX_LIST_LEN { - return Ok(None); - } - - Ok(Some(Arc::new(BranchlessFilter::::try_new(in_array)?))) -} - -#[cfg(test)] -mod tests { - use arrow::array::UInt32Array; - use arrow::datatypes::UInt32Type; - - use super::super::branchless_filter::BranchlessFilterType; - use super::*; - - fn uint32_array(values: Vec>) -> ArrayRef { - Arc::new(UInt32Array::from(values)) - } - - #[test] - fn branchless_routing_respects_max_list_len() -> Result<()> { - let max_len = ::MAX_LIST_LEN; - - let values = (0..max_len) - .map(|value| Some(value as u32)) - .collect::>(); - assert!(instantiate_branchless_filter(&uint32_array(values))?.is_some()); - - let values = (0..=max_len) - .map(|value| Some(value as u32)) - .collect::>(); - assert!(instantiate_branchless_filter(&uint32_array(values))?.is_none()); - - Ok(()) - } - - #[test] - fn branchless_routing_handles_zero_non_null_values() -> Result<()> { - let array = uint32_array(vec![None; 3]); - - assert!(instantiate_branchless_filter(&array)?.is_some()); - - Ok(()) - } -}