From 6795d52be54734263a0dbe4c5344f648490d7a1b Mon Sep 17 00:00:00 2001 From: Geoffrey Claude Date: Mon, 6 Jul 2026 13:34:14 +0200 Subject: [PATCH 1/3] Add branchless IN-list filter for small primitive lists --- .../physical-expr/benches/in_list_strategy.rs | 28 +- .../physical-expr/src/expressions/in_list.rs | 18 +- .../expressions/in_list/primitive_filter.rs | 434 +++++++++++++++++- .../src/expressions/in_list/strategy.rs | 152 +++++- 4 files changed, 602 insertions(+), 30 deletions(-) diff --git a/datafusion/physical-expr/benches/in_list_strategy.rs b/datafusion/physical-expr/benches/in_list_strategy.rs index c70f6da2a40d9..c69af192b9cdd 100644 --- a/datafusion/physical-expr/benches/in_list_strategy.rs +++ b/datafusion/physical-expr/benches/in_list_strategy.rs @@ -37,6 +37,7 @@ //! | Narrow integer cases | Int16, Float16 | larger value domain | 4, 64, 256 | //! | 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 | //! | 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 | @@ -45,8 +46,9 @@ //! | Shared-prefix string cases | Utf8, Utf8View | same prefix, different suffix | 16, 32, 64 | //! | Fixed-size binary cases | FixedSizeBinary(16) | fixed-width binary values | 4, 64, 256, 10000 | +use arrow::array::types::IntervalMonthDayNano; use arrow::array::*; -use arrow::datatypes::{Field, Int32Type, Schema}; +use arrow::datatypes::{Field, Int32Type, IntervalMonthDayNanoType, Schema}; use arrow::record_batch::RecordBatch; use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; use datafusion_common::ScalarValue; @@ -528,6 +530,28 @@ fn bench_timestamp_ns(c: &mut Criterion) { } } +fn bench_interval_month_day_nano(c: &mut Criterion) { + for match_pct in MATCH_RATES { + bench_numeric::( + c, + "interval_month_day_nano", + &format!("small_list/list=4/match={match_pct}%"), + &NumericBenchConfig::new( + 4, + match_pct as f64 / 100.0, + |rng| { + IntervalMonthDayNanoType::make_value( + rng.random_range(-120..=120), + rng.random_range(-31..=31), + rng.random_range(-1_000_000_000..=1_000_000_000), + ) + }, + |v| ScalarValue::IntervalMonthDayNano(Some(v)), + ), + ); + } +} + // ============================================================================= // UTF8 STRING CASE BENCHMARKS // ============================================================================= @@ -1049,7 +1073,7 @@ fn bench_fixed_size_binary(c: &mut Criterion) { criterion_group! { name = benches; config = Criterion::default(); - targets = bench_narrow_integer, bench_primitive, bench_f32, bench_timestamp_ns, bench_utf8, bench_utf8view, bench_dictionary, bench_nulls, bench_fixed_size_binary + targets = bench_narrow_integer, bench_primitive, bench_f32, bench_timestamp_ns, bench_interval_month_day_nano, bench_utf8, bench_utf8view, bench_dictionary, bench_nulls, bench_fixed_size_binary } criterion_main!(benches); diff --git a/datafusion/physical-expr/src/expressions/in_list.rs b/datafusion/physical-expr/src/expressions/in_list.rs index 2764083f31b09..e4ec72285cd3f 100644 --- a/datafusion/physical-expr/src/expressions/in_list.rs +++ b/datafusion/physical-expr/src/expressions/in_list.rs @@ -2900,10 +2900,9 @@ mod tests { #[test] fn test_in_list_esoteric_types() -> Result<()> { - // Test esoteric/less common types to validate the transform and mapping flow. - // These types are reinterpreted to base primitive types (e.g., Timestamp -> UInt64, - // Interval -> Decimal128, Float16 -> UInt16). We just need to verify basic - // functionality works - no need for comprehensive null handling tests. + // Test less common types covered by IN-list evaluation. Some of these + // use specialized filters, and others fall back to the generic path; + // this keeps the end-to-end behavior covered either way. // Helper: simple IN test that expects [Some(true), Some(false)] let test_type = |data_type: DataType, @@ -2926,7 +2925,7 @@ mod tests { Ok(()) }; - // Timestamp types (all units map to Int64 -> UInt64) + // Timestamp types test_type( DataType::Timestamp(TimeUnit::Second, None), Arc::new(TimestampSecondArray::from(vec![Some(1000), Some(2000)])), @@ -2960,7 +2959,7 @@ mod tests { ], )?; - // Time32 and Time64 (map to Int32 -> UInt32 and Int64 -> UInt64 respectively) + // Time32 and Time64 test_type( DataType::Time32(TimeUnit::Second), Arc::new(Time32SecondArray::from(vec![Some(3600), Some(7200)])), @@ -3006,7 +3005,7 @@ mod tests { ], )?; - // Duration types (map to Int64 -> UInt64) + // Duration types test_type( DataType::Duration(TimeUnit::Second), Arc::new(DurationSecondArray::from(vec![Some(86400), Some(172800)])), @@ -3052,7 +3051,7 @@ mod tests { ], )?; - // Interval types (map to 16-byte Decimal128Type) + // Interval types test_type( DataType::Interval(IntervalUnit::YearMonth), Arc::new(IntervalYearMonthArray::from(vec![Some(12), Some(24)])), @@ -3114,8 +3113,7 @@ mod tests { ], )?; - // Decimal256 (maps to Decimal128Type for 16-byte width) - // Need to use with_precision_and_scale() to set the metadata + // Decimal256. Need to use with_precision_and_scale() to set the metadata. let precision = 38; let scale = 10; test_type( 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 8f8d9bad04afa..ce1a9ff8b86d1 100644 --- a/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs @@ -19,12 +19,13 @@ //! //! This module provides membership tests for Arrow primitive types. -use arrow::array::{Array, ArrayRef, AsArray, BooleanArray}; -use arrow::buffer::{BooleanBuffer, NullBuffer}; +use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, PrimitiveArray}; +use arrow::buffer::{BooleanBuffer, NullBuffer, ScalarBuffer}; use arrow::datatypes::*; use arrow::util::bit_iterator::BitIndexIterator; use datafusion_common::{HashSet, Result, exec_datafusion_err}; use std::hash::{Hash, Hasher}; +use std::mem::size_of; use super::result::build_in_list_result; use super::static_filter::{StaticFilter, handle_dictionary}; @@ -222,6 +223,239 @@ where } } +pub(super) type BranchlessNative = + <::CompareType as ArrowPrimitiveType>::Native; + +/// Maximum list size for branchless lookup on 1-byte primitives. +/// +/// Sixteen 1-byte values fit in one 128-bit SIMD vector, so this keeps the +/// branchless list small enough for a single vectorized membership check. +const BRANCHLESS_MAX_1B: usize = 16; + +/// Maximum list size for branchless lookup on 2-byte primitives. +/// +/// Eight 2-byte values fit in one 128-bit SIMD vector, so this keeps the +/// branchless list small enough for a single vectorized membership check. +const BRANCHLESS_MAX_2B: usize = 8; + +/// Maximum list size for branchless lookup on 4-byte primitives. +/// +/// Thirty-two 4-byte values keep the inline list at 128 bytes. Beyond that, +/// the comparison chain and filter footprint grow enough that the hash/generic +/// fallback is a better fit. +const BRANCHLESS_MAX_4B: usize = 32; + +/// Maximum list size for branchless lookup on 8-byte primitives. +/// +/// Sixteen 8-byte values use the same 128-byte inline-list budget as 4-byte +/// primitives. Larger lists are left to the hash/generic fallback. +const BRANCHLESS_MAX_8B: usize = 16; + +/// Maximum list size for branchless lookup on 16-byte primitives. +/// +/// These comparisons are wider, so this path is limited to four values. +/// Larger lists are left to the generic fallback. +const BRANCHLESS_MAX_16B: usize = 4; + +/// Arrow primitive types supported by [`BranchlessFilter`]. +/// +/// `T` is the logical Arrow type accepted by the filter. `CompareType` is the +/// same-width type used for the fixed comparison chain. Signed integers, +/// floats, and temporal values use an unsigned comparison type so they compare +/// by their raw bit pattern. +pub(super) trait BranchlessFilterType: + ArrowPrimitiveType + Send + Sync + 'static +{ + type CompareType: ArrowPrimitiveType + Send + Sync + 'static; + + /// Maximum number of non-null IN-list values to handle with + /// [`BranchlessFilter`] for this primitive type. + const MAX_LIST_LEN: usize; +} + +macro_rules! branchless_filter_type { + ($logical:ty, $compare:ty, $max_len:expr) => { + // The branchless filter reads the same Arrow value buffer as the + // comparison type. That is only valid when both native types have the + // same width, so catch any bad mapping here at compile time. + const _: () = assert!( + size_of::<<$logical as ArrowPrimitiveType>::Native>() + == size_of::<<$compare as ArrowPrimitiveType>::Native>(), + "BranchlessFilterType::CompareType must use the same native width" + ); + + impl BranchlessFilterType for $logical { + type CompareType = $compare; + const MAX_LIST_LEN: usize = $max_len; + } + }; +} + +branchless_filter_type!(Int8Type, UInt8Type, BRANCHLESS_MAX_1B); +branchless_filter_type!(UInt8Type, UInt8Type, BRANCHLESS_MAX_1B); +branchless_filter_type!(Int16Type, UInt16Type, BRANCHLESS_MAX_2B); +branchless_filter_type!(UInt16Type, UInt16Type, BRANCHLESS_MAX_2B); +branchless_filter_type!(Float16Type, UInt16Type, BRANCHLESS_MAX_2B); + +branchless_filter_type!(Int32Type, UInt32Type, BRANCHLESS_MAX_4B); +branchless_filter_type!(UInt32Type, UInt32Type, BRANCHLESS_MAX_4B); +branchless_filter_type!(Float32Type, UInt32Type, BRANCHLESS_MAX_4B); +branchless_filter_type!(Date32Type, UInt32Type, BRANCHLESS_MAX_4B); +branchless_filter_type!(Time32SecondType, UInt32Type, BRANCHLESS_MAX_4B); +branchless_filter_type!(Time32MillisecondType, UInt32Type, BRANCHLESS_MAX_4B); + +branchless_filter_type!(Int64Type, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(UInt64Type, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(Float64Type, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(Date64Type, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(Time64MicrosecondType, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(Time64NanosecondType, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(TimestampSecondType, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(TimestampMillisecondType, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(TimestampMicrosecondType, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(TimestampNanosecondType, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(DurationSecondType, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(DurationMillisecondType, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(DurationMicrosecondType, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(DurationNanosecondType, UInt64Type, BRANCHLESS_MAX_8B); + +branchless_filter_type!(Decimal128Type, Decimal128Type, BRANCHLESS_MAX_16B); +branchless_filter_type!( + IntervalMonthDayNanoType, + IntervalMonthDayNanoType, + BRANCHLESS_MAX_16B +); + +/// A branchless filter for fixed-width primitive `IN` lists up to +/// `T::MAX_LIST_LEN` values. +/// +/// The filter stores the non-null values in a fixed-size array and checks each +/// input value with a fixed comparison chain. This avoids hash table work for +/// short lists where a handful of direct comparisons is cheaper. +pub(super) struct BranchlessFilter { + expected_data_type: DataType, + null_count: usize, + values: [BranchlessNative; N], +} + +impl BranchlessFilter +where + T: BranchlessFilterType, + BranchlessNative: Copy + PartialEq, +{ + pub(super) fn try_new(in_array: &ArrayRef) -> Result { + let in_array = in_array.as_primitive_opt::().ok_or_else(|| { + exec_datafusion_err!("BranchlessFilter: expected {} array", T::DATA_TYPE) + })?; + let non_null_count = in_array.len() - in_array.null_count(); + if non_null_count != N { + return Err(exec_datafusion_err!( + "BranchlessFilter: expected {N} non-null values, got {non_null_count}" + )); + } + + let input_values = branchless_values::(in_array); + let mut values = [::default_value(); N]; + let mut i = 0; + + match in_array.nulls() { + None => { + for &value in input_values.iter() { + values[i] = value; + i += 1; + } + } + Some(nulls) => { + for row in + BitIndexIterator::new(nulls.validity(), nulls.offset(), nulls.len()) + { + values[i] = input_values[row]; + i += 1; + } + } + } + + debug_assert_eq!(i, N); + Ok(Self { + expected_data_type: in_array.data_type().clone(), + null_count: in_array.null_count(), + values, + }) + } + + #[inline(always)] + fn check(&self, needle: BranchlessNative) -> bool { + // `N` is known at compile time, and this small method is always + // inlined. In optimized builds, the compiler can turn this into the + // same kind of code we would write by hand: + // `(value0 == needle) | (value1 == needle) | ...`. + // + // This simple shape also lets the optimizer use vector instructions + // (SIMD) when the caller applies the filter to many input rows. + // + // Use `|` instead of `||` so every comparison is evaluated. With `||`, + // the code would branch as soon as one value matched. + self.values + .iter() + .fold(false, |acc, &v| acc | (v == needle)) + } + + #[inline] + fn contains_slice( + &self, + values: &[BranchlessNative], + nulls: Option<&NullBuffer>, + negated: bool, + ) -> BooleanArray { + build_in_list_result(values.len(), nulls, self.null_count > 0, negated, |i| { + // SAFETY: `build_in_list_result` invokes this closure for + // indices in `0..values.len()`. + let needle = unsafe { *values.get_unchecked(i) }; + self.check(needle) + }) + } +} + +impl StaticFilter for BranchlessFilter +where + T: BranchlessFilterType, + BranchlessNative: Copy + PartialEq + Send + Sync, +{ + fn null_count(&self) -> usize { + self.null_count + } + + fn contains(&self, v: &dyn Array, negated: bool) -> Result { + handle_dictionary!(self, v, negated); + + if v.data_type() != &self.expected_data_type { + return Err(exec_datafusion_err!( + "BranchlessFilter: expected {} array, got {}", + self.expected_data_type, + v.data_type() + )); + } + + let v = v.as_primitive_opt::().ok_or_else(|| { + exec_datafusion_err!("BranchlessFilter: expected {} array", T::DATA_TYPE) + })?; + let values = branchless_values::(v); + Ok(self.contains_slice(values.as_ref(), v.nulls(), negated)) + } +} + +fn branchless_values(array: &PrimitiveArray) -> ScalarBuffer> +where + T: BranchlessFilterType, +{ + let data = array.to_data(); + ScalarBuffer::>::new( + data.buffers()[0].clone(), + data.offset(), + data.len(), + ) +} + /// Wrapper for f32 that implements Hash and Eq using bit comparison. /// This treats NaN values as equal to each other when they have the same bit pattern. #[derive(Clone, Copy)] @@ -430,7 +664,9 @@ mod tests { use std::sync::Arc; use arrow::array::{ - DictionaryArray, Float16Array, Int8Array, Int16Array, UInt8Array, UInt16Array, + Decimal128Array, DictionaryArray, Float16Array, Float32Array, Float64Array, + Int8Array, Int16Array, IntervalMonthDayNanoArray, TimestampNanosecondArray, + UInt8Array, UInt16Array, }; use half::f16; @@ -584,4 +820,196 @@ mod tests { Ok(()) } + + #[test] + fn branchless_filter_u8_handles_nulls() -> Result<()> { + let haystack: ArrayRef = Arc::new(UInt8Array::from(vec![Some(1), None, Some(3)])); + let filter = BranchlessFilter::::try_new(&haystack)?; + let needles = UInt8Array::from(vec![Some(1), Some(2), None, Some(3)]); + + assert_contains(&filter, &needles, vec![Some(true), None, None, Some(true)])?; + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![Some(false), None, None, Some(false)]) + ); + + Ok(()) + } + + #[test] + fn branchless_filter_i8_handles_signed_boundaries_and_slices() -> Result<()> { + let haystack: ArrayRef = Arc::new( + Int8Array::from(vec![Some(99), Some(i8::MIN), None, Some(-1), Some(42)]) + .slice(1, 3), + ); + let filter = BranchlessFilter::::try_new(&haystack)?; + let needles = + Int8Array::from(vec![Some(7), Some(i8::MIN), Some(-1), None]).slice(1, 3); + + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![Some(true), Some(true), None]) + ); + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![Some(false), Some(false), None]) + ); + + let wrong_type = UInt8Array::from(vec![Some(128), Some(u8::MAX)]); + let err = filter.contains(&wrong_type, false).unwrap_err().to_string(); + assert!(err.contains("expected Int8 array, got UInt8"), "{err}"); + + Ok(()) + } + + #[test] + fn branchless_filter_f16_handles_bit_patterns_and_slices() -> Result<()> { + let nan_a = f16::from_bits(0x7e01); + let nan_b = f16::from_bits(0x7e02); + let haystack: ArrayRef = Arc::new( + Float16Array::from(vec![ + Some(f16::from_f32(9.0)), + Some(f16::from_f32(-0.0)), + Some(nan_a), + None, + ]) + .slice(1, 3), + ); + let filter = BranchlessFilter::::try_new(&haystack)?; + let needles = Float16Array::from(vec![ + Some(f16::from_f32(0.0)), + Some(f16::from_f32(-0.0)), + Some(nan_a), + Some(nan_b), + None, + ]); + + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![None, Some(true), Some(true), None, None]) + ); + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![None, Some(false), Some(false), None, None]) + ); + + let wrong_type = UInt16Array::from(vec![Some(0x8000), Some(0x7e01)]); + let err = filter.contains(&wrong_type, false).unwrap_err().to_string(); + assert!(err.contains("expected Float16 array, got UInt16"), "{err}"); + + Ok(()) + } + + #[test] + fn branchless_filter_floats_use_bit_equality() -> Result<()> { + let nan_a = f32::from_bits(0x7fc0_0001); + let nan_b = f32::from_bits(0x7fc0_0002); + let haystack: ArrayRef = + Arc::new(Float32Array::from(vec![Some(-0.0), Some(nan_a)])); + let filter = BranchlessFilter::::try_new(&haystack)?; + let needles = + Float32Array::from(vec![Some(0.0), Some(-0.0), Some(nan_a), Some(nan_b)]); + + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![Some(false), Some(true), Some(true), Some(false)]) + ); + + let nan_a = f64::from_bits(0x7ff8_0000_0000_0001); + let nan_b = f64::from_bits(0x7ff8_0000_0000_0002); + let haystack: ArrayRef = + Arc::new(Float64Array::from(vec![Some(-0.0), Some(nan_a)])); + let filter = BranchlessFilter::::try_new(&haystack)?; + let needles = + Float64Array::from(vec![Some(0.0), Some(-0.0), Some(nan_a), Some(nan_b)]); + + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![Some(false), Some(true), Some(true), Some(false)]) + ); + + Ok(()) + } + + #[test] + fn branchless_filter_timestamp_preserves_exact_timezone() -> Result<()> { + let haystack: ArrayRef = Arc::new( + TimestampNanosecondArray::from(vec![Some(1), Some(3)]).with_timezone("UTC"), + ); + let filter = BranchlessFilter::::try_new(&haystack)?; + let needles = TimestampNanosecondArray::from(vec![Some(1), Some(2), None]) + .with_timezone("UTC"); + + assert_contains(&filter, &needles, vec![Some(true), Some(false), None])?; + + let different_timezone = + TimestampNanosecondArray::from(vec![Some(1)]).with_timezone("Europe/Paris"); + let err = filter + .contains(&different_timezone, false) + .unwrap_err() + .to_string(); + assert!(err.contains("expected Timestamp"), "{err}"); + assert!(err.contains("UTC"), "{err}"); + assert!(err.contains("Europe/Paris"), "{err}"); + + Ok(()) + } + + #[test] + fn branchless_filter_decimal128_handles_precision_scale_and_nulls() -> Result<()> { + let haystack: ArrayRef = Arc::new( + Decimal128Array::from(vec![Some(12345), None, Some(-700), Some(42)]) + .with_precision_and_scale(10, 2)?, + ); + let filter = BranchlessFilter::::try_new(&haystack)?; + let needles = + Decimal128Array::from(vec![Some(12345), Some(999), None, Some(-700)]) + .with_precision_and_scale(10, 2)?; + + assert_contains(&filter, &needles, vec![Some(true), None, None, Some(true)])?; + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![Some(false), None, None, Some(false)]) + ); + + let different_scale = + Decimal128Array::from(vec![Some(12345)]).with_precision_and_scale(10, 3)?; + let err = filter + .contains(&different_scale, false) + .unwrap_err() + .to_string(); + assert!(err.contains("expected Decimal128(10, 2) array"), "{err}"); + assert!(err.contains("got Decimal128(10, 3)"), "{err}"); + + Ok(()) + } + + #[test] + fn branchless_filter_interval_month_day_nano_handles_nulls() -> Result<()> { + let one_month = IntervalMonthDayNanoType::make_value(1, 0, 0); + let two_days = IntervalMonthDayNanoType::make_value(0, 2, 0); + let three_nanos = IntervalMonthDayNanoType::make_value(0, 0, 3); + let absent = IntervalMonthDayNanoType::make_value(4, 5, 6); + let haystack: ArrayRef = Arc::new(IntervalMonthDayNanoArray::from(vec![ + Some(one_month), + None, + Some(two_days), + Some(three_nanos), + ])); + let filter = BranchlessFilter::::try_new(&haystack)?; + let needles = IntervalMonthDayNanoArray::from(vec![ + Some(one_month), + Some(absent), + None, + Some(three_nanos), + ]); + + assert_contains(&filter, &needles, vec![Some(true), None, None, Some(true)])?; + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![Some(false), None, None, Some(false)]) + ); + + Ok(()) + } } diff --git a/datafusion/physical-expr/src/expressions/in_list/strategy.rs b/datafusion/physical-expr/src/expressions/in_list/strategy.rs index 9db90ea4faf13..eb234604c34d1 100644 --- a/datafusion/physical-expr/src/expressions/in_list/strategy.rs +++ b/datafusion/physical-expr/src/expressions/in_list/strategy.rs @@ -20,7 +20,13 @@ use std::sync::Arc; use arrow::array::ArrayRef; use arrow::compute::cast; use arrow::datatypes::{ - DataType, Float16Type, Int8Type, Int16Type, UInt8Type, UInt16Type, + 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 datafusion_common::Result; @@ -28,24 +34,88 @@ use super::array_static_filter::ArrayStaticFilter; use super::primitive_filter::*; use super::static_filter::StaticFilter; -pub(super) fn instantiate_static_filter( - in_array: ArrayRef, -) -> Result> { +type StaticFilterRef = Arc; + +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)? { + return Ok(filter); + } + + instantiate_standard_filter(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. - let in_array = match in_array.data_type() { - DataType::Dictionary(_, value_type) => cast(&in_array, value_type.as_ref())?, - _ => in_array, - }; match in_array.data_type() { - DataType::Int8 => Ok(Arc::new(BitmapFilter::::try_new(&in_array)?)), - DataType::UInt8 => Ok(Arc::new(BitmapFilter::::try_new(&in_array)?)), - DataType::Int16 => Ok(Arc::new(BitmapFilter::::try_new(&in_array)?)), - DataType::UInt16 => Ok(Arc::new(BitmapFilter::::try_new(&in_array)?)), - DataType::Float16 => { - Ok(Arc::new(BitmapFilter::::try_new(&in_array)?)) + 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)?)), @@ -54,8 +124,60 @@ pub(super) fn instantiate_static_filter( 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.) */ + // 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, +{ + if non_null_count > T::MAX_LIST_LEN { + return Ok(None); + } + + branchless_filter_for_len::(in_array, non_null_count).map(Some) +} + +fn branchless_filter_for_len( + in_array: &ArrayRef, + non_null_count: usize, +) -> Result +where + T: BranchlessFilterType, + BranchlessNative: Copy + PartialEq + Send + Sync, +{ + macro_rules! dispatch { + ($($n:literal),* $(,)?) => { + match non_null_count { + $($n => Ok(Arc::new(BranchlessFilter::::try_new(in_array)?)),)* + _ => unreachable!("validated branchless list length"), + } + }; + } + + match T::MAX_LIST_LEN { + 4 => dispatch!(0, 1, 2, 3, 4), + 8 => dispatch!(0, 1, 2, 3, 4, 5, 6, 7, 8), + 16 => dispatch!(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16), + 32 => dispatch!( + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, + 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, + ), + _ => unreachable!("known branchless max list length"), + } +} From 3ab971a6c40d435769ee3a77e9da354d2573961c Mon Sep 17 00:00:00 2001 From: Geoffrey Claude Date: Wed, 8 Jul 2026 13:51:16 +0200 Subject: [PATCH 2/3] Prepare primitive IN-list filters for shared lookup code --- .../expressions/in_list/primitive_filter.rs | 110 +++++++++--------- .../src/expressions/in_list/strategy.rs | 12 +- 2 files changed, 61 insertions(+), 61 deletions(-) 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 ce1a9ff8b86d1..57ef2895c2987 100644 --- a/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs @@ -223,8 +223,8 @@ where } } -pub(super) type BranchlessNative = - <::CompareType as ArrowPrimitiveType>::Native; +pub(super) type PrimitiveFilterNative = + <::CompareType as ArrowPrimitiveType>::Native; /// Maximum list size for branchless lookup on 1-byte primitives. /// @@ -263,85 +263,85 @@ const BRANCHLESS_MAX_16B: usize = 4; /// same-width type used for the fixed comparison chain. Signed integers, /// floats, and temporal values use an unsigned comparison type so they compare /// by their raw bit pattern. -pub(super) trait BranchlessFilterType: +pub(super) trait PrimitiveFilterType: ArrowPrimitiveType + Send + Sync + 'static { type CompareType: ArrowPrimitiveType + Send + Sync + 'static; /// Maximum number of non-null IN-list values to handle with /// [`BranchlessFilter`] for this primitive type. - const MAX_LIST_LEN: usize; + const BRANCHLESS_MAX_LIST_LEN: usize; } -macro_rules! branchless_filter_type { +macro_rules! primitive_filter_type { ($logical:ty, $compare:ty, $max_len:expr) => { - // The branchless filter reads the same Arrow value buffer as the - // comparison type. That is only valid when both native types have the - // same width, so catch any bad mapping here at compile time. + // These filters read the same Arrow value buffer as the comparison + // type. That is only valid when both native types have the same width, + // so catch any bad mapping here at compile time. const _: () = assert!( size_of::<<$logical as ArrowPrimitiveType>::Native>() == size_of::<<$compare as ArrowPrimitiveType>::Native>(), - "BranchlessFilterType::CompareType must use the same native width" + "PrimitiveFilterType::CompareType must use the same native width" ); - impl BranchlessFilterType for $logical { + impl PrimitiveFilterType for $logical { type CompareType = $compare; - const MAX_LIST_LEN: usize = $max_len; + const BRANCHLESS_MAX_LIST_LEN: usize = $max_len; } }; } -branchless_filter_type!(Int8Type, UInt8Type, BRANCHLESS_MAX_1B); -branchless_filter_type!(UInt8Type, UInt8Type, BRANCHLESS_MAX_1B); -branchless_filter_type!(Int16Type, UInt16Type, BRANCHLESS_MAX_2B); -branchless_filter_type!(UInt16Type, UInt16Type, BRANCHLESS_MAX_2B); -branchless_filter_type!(Float16Type, UInt16Type, BRANCHLESS_MAX_2B); - -branchless_filter_type!(Int32Type, UInt32Type, BRANCHLESS_MAX_4B); -branchless_filter_type!(UInt32Type, UInt32Type, BRANCHLESS_MAX_4B); -branchless_filter_type!(Float32Type, UInt32Type, BRANCHLESS_MAX_4B); -branchless_filter_type!(Date32Type, UInt32Type, BRANCHLESS_MAX_4B); -branchless_filter_type!(Time32SecondType, UInt32Type, BRANCHLESS_MAX_4B); -branchless_filter_type!(Time32MillisecondType, UInt32Type, BRANCHLESS_MAX_4B); - -branchless_filter_type!(Int64Type, UInt64Type, BRANCHLESS_MAX_8B); -branchless_filter_type!(UInt64Type, UInt64Type, BRANCHLESS_MAX_8B); -branchless_filter_type!(Float64Type, UInt64Type, BRANCHLESS_MAX_8B); -branchless_filter_type!(Date64Type, UInt64Type, BRANCHLESS_MAX_8B); -branchless_filter_type!(Time64MicrosecondType, UInt64Type, BRANCHLESS_MAX_8B); -branchless_filter_type!(Time64NanosecondType, UInt64Type, BRANCHLESS_MAX_8B); -branchless_filter_type!(TimestampSecondType, UInt64Type, BRANCHLESS_MAX_8B); -branchless_filter_type!(TimestampMillisecondType, UInt64Type, BRANCHLESS_MAX_8B); -branchless_filter_type!(TimestampMicrosecondType, UInt64Type, BRANCHLESS_MAX_8B); -branchless_filter_type!(TimestampNanosecondType, UInt64Type, BRANCHLESS_MAX_8B); -branchless_filter_type!(DurationSecondType, UInt64Type, BRANCHLESS_MAX_8B); -branchless_filter_type!(DurationMillisecondType, UInt64Type, BRANCHLESS_MAX_8B); -branchless_filter_type!(DurationMicrosecondType, UInt64Type, BRANCHLESS_MAX_8B); -branchless_filter_type!(DurationNanosecondType, UInt64Type, BRANCHLESS_MAX_8B); - -branchless_filter_type!(Decimal128Type, Decimal128Type, BRANCHLESS_MAX_16B); -branchless_filter_type!( +primitive_filter_type!(Int8Type, UInt8Type, BRANCHLESS_MAX_1B); +primitive_filter_type!(UInt8Type, UInt8Type, BRANCHLESS_MAX_1B); +primitive_filter_type!(Int16Type, UInt16Type, BRANCHLESS_MAX_2B); +primitive_filter_type!(UInt16Type, UInt16Type, BRANCHLESS_MAX_2B); +primitive_filter_type!(Float16Type, UInt16Type, BRANCHLESS_MAX_2B); + +primitive_filter_type!(Int32Type, UInt32Type, BRANCHLESS_MAX_4B); +primitive_filter_type!(UInt32Type, UInt32Type, BRANCHLESS_MAX_4B); +primitive_filter_type!(Float32Type, UInt32Type, BRANCHLESS_MAX_4B); +primitive_filter_type!(Date32Type, UInt32Type, BRANCHLESS_MAX_4B); +primitive_filter_type!(Time32SecondType, UInt32Type, BRANCHLESS_MAX_4B); +primitive_filter_type!(Time32MillisecondType, UInt32Type, BRANCHLESS_MAX_4B); + +primitive_filter_type!(Int64Type, UInt64Type, BRANCHLESS_MAX_8B); +primitive_filter_type!(UInt64Type, UInt64Type, BRANCHLESS_MAX_8B); +primitive_filter_type!(Float64Type, UInt64Type, BRANCHLESS_MAX_8B); +primitive_filter_type!(Date64Type, UInt64Type, BRANCHLESS_MAX_8B); +primitive_filter_type!(Time64MicrosecondType, UInt64Type, BRANCHLESS_MAX_8B); +primitive_filter_type!(Time64NanosecondType, UInt64Type, BRANCHLESS_MAX_8B); +primitive_filter_type!(TimestampSecondType, UInt64Type, BRANCHLESS_MAX_8B); +primitive_filter_type!(TimestampMillisecondType, UInt64Type, BRANCHLESS_MAX_8B); +primitive_filter_type!(TimestampMicrosecondType, UInt64Type, BRANCHLESS_MAX_8B); +primitive_filter_type!(TimestampNanosecondType, UInt64Type, BRANCHLESS_MAX_8B); +primitive_filter_type!(DurationSecondType, UInt64Type, BRANCHLESS_MAX_8B); +primitive_filter_type!(DurationMillisecondType, UInt64Type, BRANCHLESS_MAX_8B); +primitive_filter_type!(DurationMicrosecondType, UInt64Type, BRANCHLESS_MAX_8B); +primitive_filter_type!(DurationNanosecondType, UInt64Type, BRANCHLESS_MAX_8B); + +primitive_filter_type!(Decimal128Type, Decimal128Type, BRANCHLESS_MAX_16B); +primitive_filter_type!( IntervalMonthDayNanoType, IntervalMonthDayNanoType, BRANCHLESS_MAX_16B ); /// A branchless filter for fixed-width primitive `IN` lists up to -/// `T::MAX_LIST_LEN` values. +/// `T::BRANCHLESS_MAX_LIST_LEN` values. /// /// The filter stores the non-null values in a fixed-size array and checks each /// input value with a fixed comparison chain. This avoids hash table work for /// short lists where a handful of direct comparisons is cheaper. -pub(super) struct BranchlessFilter { +pub(super) struct BranchlessFilter { expected_data_type: DataType, null_count: usize, - values: [BranchlessNative; N], + values: [PrimitiveFilterNative; N], } impl BranchlessFilter where - T: BranchlessFilterType, - BranchlessNative: Copy + PartialEq, + T: PrimitiveFilterType, + PrimitiveFilterNative: Copy + PartialEq, { pub(super) fn try_new(in_array: &ArrayRef) -> Result { let in_array = in_array.as_primitive_opt::().ok_or_else(|| { @@ -354,7 +354,7 @@ where )); } - let input_values = branchless_values::(in_array); + let input_values = primitive_values::(in_array); let mut values = [::default_value(); N]; let mut i = 0; @@ -384,7 +384,7 @@ where } #[inline(always)] - fn check(&self, needle: BranchlessNative) -> bool { + fn check(&self, needle: PrimitiveFilterNative) -> bool { // `N` is known at compile time, and this small method is always // inlined. In optimized builds, the compiler can turn this into the // same kind of code we would write by hand: @@ -403,7 +403,7 @@ where #[inline] fn contains_slice( &self, - values: &[BranchlessNative], + values: &[PrimitiveFilterNative], nulls: Option<&NullBuffer>, negated: bool, ) -> BooleanArray { @@ -418,8 +418,8 @@ where impl StaticFilter for BranchlessFilter where - T: BranchlessFilterType, - BranchlessNative: Copy + PartialEq + Send + Sync, + T: PrimitiveFilterType, + PrimitiveFilterNative: Copy + PartialEq + Send + Sync, { fn null_count(&self) -> usize { self.null_count @@ -439,17 +439,17 @@ where let v = v.as_primitive_opt::().ok_or_else(|| { exec_datafusion_err!("BranchlessFilter: expected {} array", T::DATA_TYPE) })?; - let values = branchless_values::(v); + let values = primitive_values::(v); Ok(self.contains_slice(values.as_ref(), v.nulls(), negated)) } } -fn branchless_values(array: &PrimitiveArray) -> ScalarBuffer> +fn primitive_values(array: &PrimitiveArray) -> ScalarBuffer> where - T: BranchlessFilterType, + T: PrimitiveFilterType, { let data = array.to_data(); - ScalarBuffer::>::new( + ScalarBuffer::>::new( data.buffers()[0].clone(), data.offset(), data.len(), diff --git a/datafusion/physical-expr/src/expressions/in_list/strategy.rs b/datafusion/physical-expr/src/expressions/in_list/strategy.rs index eb234604c34d1..f07ae38f43390 100644 --- a/datafusion/physical-expr/src/expressions/in_list/strategy.rs +++ b/datafusion/physical-expr/src/expressions/in_list/strategy.rs @@ -143,10 +143,10 @@ fn branchless_filter( non_null_count: usize, ) -> Result> where - T: BranchlessFilterType, - BranchlessNative: Copy + PartialEq + Send + Sync, + T: PrimitiveFilterType, + PrimitiveFilterNative: Copy + PartialEq + Send + Sync, { - if non_null_count > T::MAX_LIST_LEN { + if non_null_count > T::BRANCHLESS_MAX_LIST_LEN { return Ok(None); } @@ -158,8 +158,8 @@ fn branchless_filter_for_len( non_null_count: usize, ) -> Result where - T: BranchlessFilterType, - BranchlessNative: Copy + PartialEq + Send + Sync, + T: PrimitiveFilterType, + PrimitiveFilterNative: Copy + PartialEq + Send + Sync, { macro_rules! dispatch { ($($n:literal),* $(,)?) => { @@ -170,7 +170,7 @@ where }; } - match T::MAX_LIST_LEN { + match T::BRANCHLESS_MAX_LIST_LEN { 4 => dispatch!(0, 1, 2, 3, 4), 8 => dispatch!(0, 1, 2, 3, 4, 5, 6, 7, 8), 16 => dispatch!(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16), From a6f0778464f2f987d5114abb194cb5d7f4a700eb Mon Sep 17 00:00:00 2001 From: Geoffrey Claude Date: Wed, 8 Jul 2026 13:52:48 +0200 Subject: [PATCH 3/3] Implement Direct Probe Filter for large primitive IN lists --- .../expressions/in_list/primitive_filter.rs | 290 +++++++++++++++++- .../src/expressions/in_list/strategy.rs | 69 +++++ 2 files changed, 352 insertions(+), 7 deletions(-) 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 57ef2895c2987..1ece5276873a3 100644 --- a/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs @@ -19,7 +19,9 @@ //! //! This module provides membership tests for Arrow primitive types. -use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, PrimitiveArray}; +use arrow::array::{ + Array, ArrayRef, AsArray, BooleanArray, PrimitiveArray, types::IntervalMonthDayNano, +}; use arrow::buffer::{BooleanBuffer, NullBuffer, ScalarBuffer}; use arrow::datatypes::*; use arrow::util::bit_iterator::BitIndexIterator; @@ -257,12 +259,12 @@ const BRANCHLESS_MAX_8B: usize = 16; /// Larger lists are left to the generic fallback. const BRANCHLESS_MAX_16B: usize = 4; -/// Arrow primitive types supported by [`BranchlessFilter`]. +/// Arrow primitive types supported by fixed-width primitive filters. /// /// `T` is the logical Arrow type accepted by the filter. `CompareType` is the -/// same-width type used for the fixed comparison chain. Signed integers, -/// floats, and temporal values use an unsigned comparison type so they compare -/// by their raw bit pattern. +/// same-width type used for the optimized lookup. Signed integers, floats, and +/// temporal values use an unsigned comparison type so they compare by their raw +/// bit pattern. pub(super) trait PrimitiveFilterType: ArrowPrimitiveType + Send + Sync + 'static { @@ -444,7 +446,9 @@ where } } -fn primitive_values(array: &PrimitiveArray) -> ScalarBuffer> +fn primitive_values( + array: &PrimitiveArray, +) -> ScalarBuffer> where T: PrimitiveFilterType, { @@ -456,6 +460,210 @@ where ) } +/// Load factor inverse for [`DirectProbeFilter`]. +/// +/// A value of 4 means the table is sized to around 25% full. The extra empty +/// slots are intentional: misses are common for `IN` filters, and a miss stops +/// as soon as probing reaches an empty slot. +const DIRECT_PROBE_LOAD_FACTOR_INVERSE: usize = 4; + +/// Minimum table size for [`DirectProbeFilter`]. +const DIRECT_PROBE_MIN_TABLE_SIZE: usize = 16; + +/// Golden ratio constant for 32-bit hash mixing. +const GOLDEN_RATIO_32: u32 = 0x9e3779b9; + +/// Golden ratio constant for 64-bit hash mixing. +const GOLDEN_RATIO_64: u64 = 0x9e3779b97f4a7c15; + +/// Secondary mixing constant for 128-bit hashing. +const SPLITMIX_CONSTANT: u64 = 0xbf58476d1ce4e5b9; + +/// Values that can be stored in [`DirectProbeFilter`]. +/// +/// The filter uses normal equality for matches. `probe_hash()` is only for the +/// lookup table used by the hot path, where the work is limited to fixed-width +/// primitive values and does not need a general-purpose `Hasher`. +pub(super) trait DirectProbeHashable: Copy + Eq { + fn probe_hash(self) -> usize; +} + +impl DirectProbeHashable for u32 { + #[inline(always)] + fn probe_hash(self) -> usize { + let x = self.wrapping_mul(GOLDEN_RATIO_32); + (x ^ (x >> 16)) as usize + } +} + +impl DirectProbeHashable for u64 { + #[inline(always)] + fn probe_hash(self) -> usize { + let x = self.wrapping_mul(GOLDEN_RATIO_64); + (x ^ (x >> 32)) as usize + } +} + +impl DirectProbeHashable for i128 { + #[inline(always)] + fn probe_hash(self) -> usize { + let lo = self as u64; + let hi = (self >> 64) as u64; + let x = lo.wrapping_mul(GOLDEN_RATIO_64) ^ hi.wrapping_mul(SPLITMIX_CONSTANT); + (x ^ (x >> 32)) as usize + } +} + +impl DirectProbeHashable for IntervalMonthDayNano { + #[inline(always)] + fn probe_hash(self) -> usize { + let months_days = (self.months as u32 as u64) | ((self.days as u32 as u64) << 32); + let x = months_days.wrapping_mul(GOLDEN_RATIO_64) + ^ (self.nanoseconds as u64).wrapping_mul(SPLITMIX_CONSTANT); + (x ^ (x >> 32)) as usize + } +} + +/// A compact open-addressing filter for larger fixed-width primitive `IN` lists. +/// +/// For list sizes above the branchless thresholds, a short comparison chain is +/// no longer the best fit. This filter stores the non-null IN-list values in a +/// power-of-two table and probes it directly with linear probing. +/// +/// The existing `HashSet` based primitive filters are still a good general +/// fallback. This type is narrower on purpose: values are fixed-width, already +/// normalized to their comparison representation, and can be checked with a +/// small table lookup plus equality test. +pub(super) struct DirectProbeFilter +where + PrimitiveFilterNative: DirectProbeHashable, +{ + expected_data_type: DataType, + null_count: usize, + table: Box<[Option>]>, + mask: usize, +} + +impl DirectProbeFilter +where + T: PrimitiveFilterType, + PrimitiveFilterNative: DirectProbeHashable, +{ + pub(super) fn try_new(in_array: &ArrayRef) -> Result { + let in_array = in_array.as_primitive_opt::().ok_or_else(|| { + exec_datafusion_err!("DirectProbeFilter: expected {} array", T::DATA_TYPE) + })?; + + let null_count = in_array.null_count(); + let value_count = in_array.len() - null_count; + // Keep the table sparse so misses usually stop after one or two probes. + // Duplicates make it even sparser, which is fine for this lookup path. + let table_size = (value_count.max(1) * DIRECT_PROBE_LOAD_FACTOR_INVERSE) + .next_power_of_two() + .max(DIRECT_PROBE_MIN_TABLE_SIZE); + let mask = table_size - 1; + let values = primitive_values::(in_array); + let mut filter = Self { + expected_data_type: in_array.data_type().clone(), + null_count, + table: vec![None; table_size].into_boxed_slice(), + mask, + }; + + match in_array.nulls() { + None => { + for &value in values.iter() { + filter.insert(value); + } + } + Some(nulls) => { + for i in + BitIndexIterator::new(nulls.validity(), nulls.offset(), nulls.len()) + { + filter.insert(values[i]); + } + } + } + + Ok(filter) + } + + fn insert(&mut self, value: PrimitiveFilterNative) { + let mut slot = value.probe_hash() & self.mask; + loop { + match self.table[slot] { + None => { + self.table[slot] = Some(value); + return; + } + Some(existing) if existing == value => return, + _ => slot = (slot + 1) & self.mask, + } + } + } + + #[inline(always)] + fn contains_single(&self, needle: PrimitiveFilterNative) -> bool { + let mut slot = needle.probe_hash() & self.mask; + loop { + // A miss stops at the first empty slot. There is always at least + // one empty slot because the table was built sparse. + // + // SAFETY: `slot` is always in bounds. The table length is a power + // of two, `mask` is `table.len() - 1`, and every probe wraps with + // `& self.mask`. + match unsafe { self.table.get_unchecked(slot) } { + None => return false, + Some(value) if *value == needle => return true, + _ => slot = (slot + 1) & self.mask, + } + } + } + + #[inline] + fn contains_slice( + &self, + values: &[PrimitiveFilterNative], + nulls: Option<&NullBuffer>, + negated: bool, + ) -> BooleanArray { + build_in_list_result(values.len(), nulls, self.null_count > 0, negated, |i| { + // SAFETY: `build_in_list_result` invokes this closure for + // indices in `0..values.len()`. + let needle = unsafe { *values.get_unchecked(i) }; + self.contains_single(needle) + }) + } +} + +impl StaticFilter for DirectProbeFilter +where + T: PrimitiveFilterType, + PrimitiveFilterNative: DirectProbeHashable + Send + Sync, +{ + fn null_count(&self) -> usize { + self.null_count + } + + fn contains(&self, v: &dyn Array, negated: bool) -> Result { + handle_dictionary!(self, v, negated); + + if v.data_type() != &self.expected_data_type { + return Err(exec_datafusion_err!( + "DirectProbeFilter: expected {} array, got {}", + self.expected_data_type, + v.data_type() + )); + } + + let v = v.as_primitive_opt::().ok_or_else(|| { + exec_datafusion_err!("DirectProbeFilter: expected {} array", T::DATA_TYPE) + })?; + let values = primitive_values::(v); + Ok(self.contains_slice(values.as_ref(), v.nulls(), negated)) + } +} + /// Wrapper for f32 that implements Hash and Eq using bit comparison. /// This treats NaN values as equal to each other when they have the same bit pattern. #[derive(Clone, Copy)] @@ -666,7 +874,7 @@ mod tests { use arrow::array::{ Decimal128Array, DictionaryArray, Float16Array, Float32Array, Float64Array, Int8Array, Int16Array, IntervalMonthDayNanoArray, TimestampNanosecondArray, - UInt8Array, UInt16Array, + UInt8Array, UInt16Array, UInt32Array, }; use half::f16; @@ -1012,4 +1220,72 @@ mod tests { Ok(()) } + + #[test] + fn direct_probe_filter_handles_slices_and_nulls() -> Result<()> { + let haystack: ArrayRef = Arc::new( + UInt32Array::from(vec![ + Some(999), + Some(10), + None, + Some(20), + Some(10), + Some(30), + ]) + .slice(1, 5), + ); + let filter = DirectProbeFilter::::try_new(&haystack)?; + let needles = + UInt32Array::from(vec![Some(0), Some(10), Some(11), Some(30), None]) + .slice(1, 4); + + assert_contains(&filter, &needles, vec![Some(true), None, Some(true), None])?; + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![Some(false), None, Some(false), None]) + ); + + Ok(()) + } + + #[test] + fn direct_probe_filter_floats_use_bit_equality() -> Result<()> { + let nan_a = f32::from_bits(0x7fc0_0001); + let nan_b = f32::from_bits(0x7fc0_0002); + let haystack: ArrayRef = + Arc::new(Float32Array::from(vec![Some(-0.0), Some(nan_a)])); + let filter = DirectProbeFilter::::try_new(&haystack)?; + let needles = + Float32Array::from(vec![Some(0.0), Some(-0.0), Some(nan_a), Some(nan_b)]); + + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![Some(false), Some(true), Some(true), Some(false)]) + ); + + Ok(()) + } + + #[test] + fn direct_probe_filter_interval_month_day_nano_handles_nulls() -> Result<()> { + let one_month = IntervalMonthDayNanoType::make_value(1, 0, 0); + let two_days = IntervalMonthDayNanoType::make_value(0, 2, 0); + let three_nanos = IntervalMonthDayNanoType::make_value(0, 0, 3); + let absent = IntervalMonthDayNanoType::make_value(4, 5, 6); + let haystack: ArrayRef = Arc::new(IntervalMonthDayNanoArray::from(vec![ + Some(one_month), + None, + Some(two_days), + Some(three_nanos), + ])); + let filter = DirectProbeFilter::::try_new(&haystack)?; + let needles = IntervalMonthDayNanoArray::from(vec![ + Some(one_month), + Some(absent), + None, + Some(three_nanos), + ]); + + assert_contains(&filter, &needles, vec![Some(true), None, None, Some(true)]) + } } diff --git a/datafusion/physical-expr/src/expressions/in_list/strategy.rs b/datafusion/physical-expr/src/expressions/in_list/strategy.rs index f07ae38f43390..8ef3376cfcd65 100644 --- a/datafusion/physical-expr/src/expressions/in_list/strategy.rs +++ b/datafusion/physical-expr/src/expressions/in_list/strategy.rs @@ -43,6 +43,10 @@ pub(super) fn instantiate_static_filter(in_array: ArrayRef) -> Result Result { } } +fn instantiate_direct_probe_filter( + in_array: &ArrayRef, +) -> Result> { + let non_null_count = in_array.len() - in_array.null_count(); + + macro_rules! filter { + ($arrow_type:ty) => { + direct_probe_filter::<$arrow_type>(in_array, non_null_count) + }; + } + + match in_array.data_type() { + 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 bitmap_filter(in_array: &ArrayRef) -> Result where T: BitmapFilterType, @@ -138,6 +192,21 @@ where Ok(Arc::new(BitmapFilter::::try_new(in_array)?)) } +fn direct_probe_filter( + in_array: &ArrayRef, + non_null_count: usize, +) -> Result> +where + T: PrimitiveFilterType, + PrimitiveFilterNative: DirectProbeHashable + Send + Sync, +{ + if non_null_count <= T::BRANCHLESS_MAX_LIST_LEN { + return Ok(None); + } + + Ok(Some(Arc::new(DirectProbeFilter::::try_new(in_array)?))) +} + fn branchless_filter( in_array: &ArrayRef, non_null_count: usize,