From 48c32f6a60b4df38164fe5c7abd56eca7d9c9157 Mon Sep 17 00:00:00 2001 From: Geoffrey Claude Date: Tue, 21 Jul 2026 18:01:40 +0200 Subject: [PATCH] Add keyed frozen set for large primitive IN lists --- .../physical-expr/benches/in_list_strategy.rs | 8 +- .../physical-expr/src/expressions/in_list.rs | 7 +- .../expressions/in_list/branchless_filter.rs | 5 +- .../src/expressions/in_list/frozen_set.rs | 213 +++++++++++ .../expressions/in_list/primitive_filter.rs | 345 ++++++++---------- .../src/expressions/in_list/strategy.rs | 203 +++++++---- 6 files changed, 520 insertions(+), 261 deletions(-) create mode 100644 datafusion/physical-expr/src/expressions/in_list/frozen_set.rs diff --git a/datafusion/physical-expr/benches/in_list_strategy.rs b/datafusion/physical-expr/benches/in_list_strategy.rs index c69af192b9cdd..1a85d53915000 100644 --- a/datafusion/physical-expr/benches/in_list_strategy.rs +++ b/datafusion/physical-expr/benches/in_list_strategy.rs @@ -35,8 +35,8 @@ //! |------|-------|-----------------|-------------------| //! | Narrow integer cases | UInt8 | small value domain | 4, 16 | //! | 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 | +//! | 32-bit primitive cases | Int32, Float32 | small and large lists | 4, 32, 33, 64, 256, 1024, 10000 | +//! | 64-bit primitive cases | Int64, TimestampNs | small and large lists | 4, 16, 17, 32, 128, 1024, 10000 | //! | 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 | @@ -420,7 +420,7 @@ fn bench_narrow_integer(c: &mut Criterion) { fn bench_primitive(c: &mut Criterion) { // Int32: small and larger list sizes - for list_size in [4, 32, 64, 256] { + for list_size in [4, 32, 33, 64, 256, 1024, 10_000] { let list_case = if list_size <= 32 { "small_list" } else { @@ -442,7 +442,7 @@ fn bench_primitive(c: &mut Criterion) { } // Int64: small and larger list sizes - for list_size in [4, 16, 32, 128] { + for list_size in [4, 16, 17, 32, 128, 1024, 10_000] { let list_case = if list_size <= 16 { "small_list" } else { diff --git a/datafusion/physical-expr/src/expressions/in_list.rs b/datafusion/physical-expr/src/expressions/in_list.rs index 874e149b58328..0825b71627721 100644 --- a/datafusion/physical-expr/src/expressions/in_list.rs +++ b/datafusion/physical-expr/src/expressions/in_list.rs @@ -38,6 +38,7 @@ use datafusion_expr::{ColumnarValue, expr_vec_fmt}; mod array_static_filter; mod branchless_filter; +mod frozen_set; mod primitive_filter; mod result; mod static_filter; @@ -222,8 +223,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 `StaticFilter` implementations for better + /// performance (for example, primitive filters for fixed-width values). /// /// 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 +2593,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() - this uses a specialized primitive 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 cd0cbd0de59a8..b5886d22af506 100644 --- a/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs @@ -58,8 +58,9 @@ //! would be useful. Wider types have too many possible values for such a //! bitmap, so their limits are tuned separately. //! -//! Larger lists use the standard filter strategy, including bitmap filters for -//! one- and two-byte types. +//! Larger lists use another filter strategy: bitmap filters for one- and +//! two-byte types, frozen sets for supported four- and eight-byte types, and +//! the standard fallback for the remaining types. //! //! # What about nulls? //! diff --git a/datafusion/physical-expr/src/expressions/in_list/frozen_set.rs b/datafusion/physical-expr/src/expressions/in_list/frozen_set.rs new file mode 100644 index 0000000000000..c9945e8dc98db --- /dev/null +++ b/datafusion/physical-expr/src/expressions/in_list/frozen_set.rs @@ -0,0 +1,213 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Immutable set for fixed-width values. + +use std::hash::{BuildHasher, Hash}; + +use datafusion_common::{Result, exec_datafusion_err}; +use hashbrown::{DefaultHashBuilder, HashTable}; + +/// Immutable set optimized for repeated membership tests. +/// +/// Each hash bucket holds two values directly. Values that collide beyond those +/// two slots use a `HashTable`, which keeps the common lookup path to one hash +/// and two bucket comparisons without sacrificing collision handling. +/// +/// The first member is used as the empty-slot sentinel and handled before +/// lookup, so primary slots store `V` directly without an `Option` wrapper. +pub(super) struct FrozenSet { + hash_builder: DefaultHashBuilder, + sentinel: Option, + buckets: Box<[[V; 2]]>, + overflowed: Box<[bool]>, + overflow: HashTable, +} + +impl FrozenSet +where + V: Copy + Eq + Hash, +{ + pub(super) fn try_new(values: &[V]) -> Result { + let Some((&sentinel, values)) = values.split_first() else { + return Ok(Self { + hash_builder: DefaultHashBuilder::default(), + sentinel: None, + buckets: Box::default(), + overflowed: Box::default(), + overflow: HashTable::new(), + }); + }; + + if values.is_empty() { + return Ok(Self { + hash_builder: DefaultHashBuilder::default(), + sentinel: Some(sentinel), + buckets: Box::default(), + overflowed: Box::default(), + overflow: HashTable::new(), + }); + } + + let bucket_count = values + .len() + .checked_add(1) + .ok_or_else(|| exec_datafusion_err!("FrozenSet capacity overflow"))?; + bucket_count + .checked_mul(2) + .ok_or_else(|| exec_datafusion_err!("FrozenSet capacity overflow"))?; + let mut set = Self { + hash_builder: DefaultHashBuilder::default(), + sentinel: Some(sentinel), + buckets: vec![[sentinel; 2]; bucket_count].into_boxed_slice(), + overflowed: vec![false; bucket_count].into_boxed_slice(), + overflow: HashTable::with_capacity(bucket_count / 8), + }; + values.iter().copied().for_each(|value| set.insert(value)); + Ok(set) + } + + fn insert(&mut self, value: V) { + let sentinel = self.sentinel.expect("non-empty frozen set"); + if value == sentinel { + return; + } + + let hash = self.hash_builder.hash_one(value); + let bucket = reduce_hash(hash, self.buckets.len()); + let slots = &mut self.buckets[bucket]; + if slots[0] == sentinel { + slots[0] = value; + return; + } + if slots[0] == value { + return; + } + if slots[1] == sentinel { + slots[1] = value; + return; + } + if slots[1] == value { + return; + } + + let hash_builder = &self.hash_builder; + self.overflow + .entry( + hash, + |stored| *stored == value, + |stored| hash_builder.hash_one(stored), + ) + .or_insert(value); + self.overflowed[bucket] = true; + } + + #[inline(always)] + pub(super) fn contains(&self, value: V) -> bool { + let Some(sentinel) = self.sentinel else { + return false; + }; + if value == sentinel { + return true; + } + if self.buckets.is_empty() { + return false; + } + + let hash = self.hash_builder.hash_one(value); + let bucket = reduce_hash(hash, self.buckets.len()); + // SAFETY: `reduce_hash` returns an index in `0..buckets.len()`. + let slots = unsafe { self.buckets.get_unchecked(bucket) }; + if (slots[0] == value) | (slots[1] == value) { + return true; + } + + // SAFETY: `overflowed` has exactly one entry per bucket. + (unsafe { *self.overflowed.get_unchecked(bucket) }) + && self + .overflow + .find(hash, |stored| *stored == value) + .is_some() + } +} + +#[inline(always)] +fn reduce_hash(hash: u64, len: usize) -> usize { + #[cfg(target_pointer_width = "64")] + return ((hash as u128 * len as u128) >> 64) as usize; + + #[cfg(target_pointer_width = "32")] + return (((hash as u32) as u64 * len as u64) >> 32) as usize; +} + +#[cfg(test)] +mod tests { + use super::*; + use std::hash::{Hash, Hasher}; + + #[derive(Clone, Copy, Debug, Eq, PartialEq)] + struct Colliding(u64); + + impl Hash for Colliding { + fn hash(&self, state: &mut H) { + 0_u8.hash(state); + } + } + + #[test] + fn handles_empty_duplicates_and_sentinel_member() { + let empty = FrozenSet::::try_new(&[]).unwrap(); + assert!(!empty.contains(0)); + + let singleton = FrozenSet::try_new(&[42]).unwrap(); + assert!(singleton.contains(42)); + assert!(!singleton.contains(7)); + + let values = [42, 7, 42, 9, 11, 7]; + let set = FrozenSet::try_new(&values).unwrap(); + for value in [42, 7, 9, 11] { + assert!(set.contains(value)); + } + for value in [0, 8, 10, 12] { + assert!(!set.contains(value)); + } + } + + #[test] + fn handles_many_values() { + let values = (0_u64..10_000).collect::>(); + let set = FrozenSet::try_new(&values).unwrap(); + assert!(values.iter().all(|&value| set.contains(value))); + assert!((10_000..20_000).all(|value| !set.contains(value))); + } + + #[test] + fn handles_u128_values() { + let values = [1_u128, (1_u128 << 64) | 1, (2_u128 << 96) | 7, u128::MAX]; + let set = FrozenSet::try_new(&values).unwrap(); + assert!(values.iter().all(|&value| set.contains(value))); + assert!(!set.contains(1_u128 << 96)); + } + + #[test] + fn handles_collisions() { + let values = (0..128).map(Colliding).collect::>(); + let set = FrozenSet::try_new(&values).unwrap(); + assert!(values.iter().all(|&value| set.contains(value))); + assert!((128..256).all(|value| !set.contains(Colliding(value)))); + } +} 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..881cf628aedb2 100644 --- a/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs @@ -19,13 +19,15 @@ //! //! 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::{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 datafusion_common::{Result, exec_datafusion_err}; +use std::hash::Hash; +use super::branchless_filter::{BranchlessFilterType, BranchlessNative}; +use super::frozen_set::FrozenSet; use super::result::build_in_list_result; use super::static_filter::{StaticFilter, handle_dictionary}; @@ -222,215 +224,110 @@ where } } -/// 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)] -struct OrderedFloat32(f32); - -impl Hash for OrderedFloat32 { - fn hash(&self, state: &mut H) { - self.0.to_ne_bytes().hash(state); - } +fn primitive_values(array: &PrimitiveArray) -> ScalarBuffer> +where + T: BranchlessFilterType, +{ + let data = array.to_data(); + ScalarBuffer::>::new( + data.buffers()[0].clone(), + data.offset(), + data.len(), + ) } -impl PartialEq for OrderedFloat32 { - fn eq(&self, other: &Self) -> bool { - self.0.to_bits() == other.0.to_bits() - } +/// Frozen-set filter for larger fixed-width primitive `IN` lists. +pub(super) struct PrimitiveFrozenFilter +where + BranchlessNative: Copy + Eq + Hash, +{ + expected_data_type: DataType, + null_count: usize, + values: FrozenSet>, } -impl Eq for OrderedFloat32 {} +impl PrimitiveFrozenFilter +where + T: BranchlessFilterType, + BranchlessNative: Copy + Eq + Hash, +{ + pub(super) fn try_new(in_array: &ArrayRef) -> Result { + let in_array = in_array.as_primitive_opt::().ok_or_else(|| { + exec_datafusion_err!("PrimitiveFrozenFilter: expected {} array", T::DATA_TYPE) + })?; -impl From for OrderedFloat32 { - fn from(v: f32) -> Self { - Self(v) + let null_count = in_array.null_count(); + let values = primitive_values::(in_array); + let values = match in_array.nulls() { + None => FrozenSet::try_new(values.as_ref())?, + Some(nulls) => { + let values = + BitIndexIterator::new(nulls.validity(), nulls.offset(), nulls.len()) + .map(|i| values[i]) + .collect::>(); + FrozenSet::try_new(&values)? + } + }; + Ok(Self { + expected_data_type: in_array.data_type().clone(), + null_count, + values, + }) } -} - -/// Wrapper for f64 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)] -struct OrderedFloat64(f64); -impl Hash for OrderedFloat64 { - fn hash(&self, state: &mut H) { - self.0.to_ne_bytes().hash(state); - } -} - -impl PartialEq for OrderedFloat64 { - fn eq(&self, other: &Self) -> bool { - self.0.to_bits() == other.0.to_bits() + #[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.values.contains(needle) + }) } } -impl Eq for OrderedFloat64 {} - -impl From for OrderedFloat64 { - fn from(v: f64) -> Self { - Self(v) +impl StaticFilter for PrimitiveFrozenFilter +where + T: BranchlessFilterType, + BranchlessNative: Copy + Eq + Hash + Send + Sync, +{ + fn null_count(&self) -> usize { + self.null_count } -} - -// 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)); - } - 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 { + handle_dictionary!(self, v, negated); - 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)) - } + if !PrimitiveArray::::is_compatible(v.data_type()) { + return Err(exec_datafusion_err!( + "PrimitiveFrozenFilter: expected {} array, got {}", + self.expected_data_type, + v.data_type() + )); } - }; -} - -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 -macro_rules! float_static_filter { - ($Name:ident, $ArrowType:ty, $OrderedType:ty) => { - primitive_static_filter!($Name, $ArrowType, $OrderedType, <$OrderedType>::from); - }; + let v = v.as_primitive_opt::().ok_or_else(|| { + exec_datafusion_err!("PrimitiveFrozenFilter: expected {} array", T::DATA_TYPE) + })?; + let values = primitive_values::(v); + Ok(self.contains_slice(values.as_ref(), v.nulls(), negated)) + } } -// Generate specialized filters for float types using ordered wrappers -float_static_filter!(Float32StaticFilter, Float32Type, OrderedFloat32); -float_static_filter!(Float64StaticFilter, Float64Type, OrderedFloat64); - #[cfg(test)] mod tests { use super::*; use std::sync::Arc; use arrow::array::{ - DictionaryArray, Float16Array, Int8Array, Int16Array, UInt8Array, UInt16Array, + DictionaryArray, Float16Array, Float32Array, Int8Array, Int16Array, + TimestampMillisecondArray, TimestampNanosecondArray, UInt8Array, UInt16Array, + UInt32Array, }; use half::f16; @@ -584,4 +481,72 @@ mod tests { Ok(()) } + + #[test] + fn primitive_frozen_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 = PrimitiveFrozenFilter::::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 primitive_frozen_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 = PrimitiveFrozenFilter::::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 primitive_frozen_filter_timestamp_uses_physical_compatibility() -> Result<()> { + let haystack: ArrayRef = Arc::new( + TimestampNanosecondArray::from(vec![Some(1), Some(3)]).with_timezone("UTC"), + ); + let filter = + PrimitiveFrozenFilter::::try_new(&haystack)?; + + let different_timezone = TimestampNanosecondArray::from(vec![Some(1), Some(2)]) + .with_timezone("Europe/Paris"); + assert_contains(&filter, &different_timezone, vec![Some(true), Some(false)])?; + + let different_unit = TimestampMillisecondArray::from(vec![Some(1)]); + let err = filter + .contains(&different_unit, false) + .unwrap_err() + .to_string(); + assert!(err.contains("Timestamp(ns"), "{err}"); + assert!(err.contains("Timestamp(ms"), "{err}"); + + Ok(()) + } } diff --git a/datafusion/physical-expr/src/expressions/in_list/strategy.rs b/datafusion/physical-expr/src/expressions/in_list/strategy.rs index d5ca8154a92f6..87eeb7c5b4933 100644 --- a/datafusion/physical-expr/src/expressions/in_list/strategy.rs +++ b/datafusion/physical-expr/src/expressions/in_list/strategy.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use std::sync::Arc; +use std::{hash::Hash, sync::Arc}; use arrow::array::ArrayRef; use arrow::compute::cast; @@ -42,103 +42,147 @@ 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)? { + 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> { +fn instantiate_primitive_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) + ($arrow_type:ty, $strategy:ident) => { + $strategy::<$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::Int8 => filter!(Int8Type, branchless_or_bitmap_filter), + DataType::UInt8 => filter!(UInt8Type, branchless_or_bitmap_filter), + DataType::Int16 => filter!(Int16Type, branchless_or_bitmap_filter), + DataType::UInt16 => filter!(UInt16Type, branchless_or_bitmap_filter), + DataType::Float16 => filter!(Float16Type, branchless_or_bitmap_filter), + DataType::Int32 => filter!(Int32Type, branchless_or_frozen_filter), + DataType::UInt32 => filter!(UInt32Type, branchless_or_frozen_filter), + DataType::Float32 => filter!(Float32Type, branchless_or_frozen_filter), + DataType::Date32 => filter!(Date32Type, branchless_or_frozen_filter), DataType::Time32(unit) => match unit { - TimeUnit::Second => filter!(Time32SecondType), - TimeUnit::Millisecond => filter!(Time32MillisecondType), + TimeUnit::Second => { + filter!(Time32SecondType, branchless_or_frozen_filter) + } + TimeUnit::Millisecond => { + filter!(Time32MillisecondType, branchless_or_frozen_filter) + } _ => Ok(None), }, - DataType::Int64 => filter!(Int64Type), - DataType::UInt64 => filter!(UInt64Type), - DataType::Float64 => filter!(Float64Type), - DataType::Date64 => filter!(Date64Type), + DataType::Int64 => filter!(Int64Type, branchless_or_frozen_filter), + DataType::UInt64 => filter!(UInt64Type, branchless_or_frozen_filter), + DataType::Float64 => filter!(Float64Type, branchless_or_frozen_filter), + DataType::Date64 => filter!(Date64Type, branchless_or_frozen_filter), DataType::Time64(unit) => match unit { - TimeUnit::Microsecond => filter!(Time64MicrosecondType), - TimeUnit::Nanosecond => filter!(Time64NanosecondType), + TimeUnit::Microsecond => { + filter!(Time64MicrosecondType, branchless_or_frozen_filter) + } + TimeUnit::Nanosecond => { + filter!(Time64NanosecondType, branchless_or_frozen_filter) + } _ => Ok(None), }, DataType::Timestamp(unit, _) => match unit { - TimeUnit::Second => filter!(TimestampSecondType), - TimeUnit::Millisecond => filter!(TimestampMillisecondType), - TimeUnit::Microsecond => filter!(TimestampMicrosecondType), - TimeUnit::Nanosecond => filter!(TimestampNanosecondType), + TimeUnit::Second => { + filter!(TimestampSecondType, branchless_or_frozen_filter) + } + TimeUnit::Millisecond => { + filter!(TimestampMillisecondType, branchless_or_frozen_filter) + } + TimeUnit::Microsecond => { + filter!(TimestampMicrosecondType, branchless_or_frozen_filter) + } + TimeUnit::Nanosecond => { + filter!(TimestampNanosecondType, branchless_or_frozen_filter) + } }, DataType::Duration(unit) => match unit { - TimeUnit::Second => filter!(DurationSecondType), - TimeUnit::Millisecond => filter!(DurationMillisecondType), - TimeUnit::Microsecond => filter!(DurationMicrosecondType), - TimeUnit::Nanosecond => filter!(DurationNanosecondType), + TimeUnit::Second => { + filter!(DurationSecondType, branchless_or_frozen_filter) + } + TimeUnit::Millisecond => { + filter!(DurationMillisecondType, branchless_or_frozen_filter) + } + TimeUnit::Microsecond => { + filter!(DurationMicrosecondType, branchless_or_frozen_filter) + } + TimeUnit::Nanosecond => { + filter!(DurationNanosecondType, branchless_or_frozen_filter) + } }, - DataType::Decimal128(_, _) => filter!(Decimal128Type), + DataType::Decimal128(_, _) => { + filter!(Decimal128Type, branchless_filter) + } DataType::Interval(IntervalUnit::MonthDayNano) => { - filter!(IntervalMonthDayNanoType) + filter!(IntervalMonthDayNanoType, branchless_filter) } _ => 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 branchless_or_bitmap_filter( + in_array: &ArrayRef, + non_null_count: usize, +) -> Result> +where + T: BranchlessFilterType + BitmapFilterType, + BranchlessNative: Copy + PartialEq + Send + Sync, +{ + if let Some(filter) = branchless_filter::(in_array, non_null_count)? { + return Ok(Some(filter)); } + + Ok(Some(Arc::new(BitmapFilter::::try_new(in_array)?))) } -fn bitmap_filter(in_array: &ArrayRef) -> Result +fn branchless_or_frozen_filter( + in_array: &ArrayRef, + non_null_count: usize, +) -> Result> where - T: BitmapFilterType, + T: BranchlessFilterType, + BranchlessNative: Copy + Eq + Hash + Send + Sync, { - Ok(Arc::new(BitmapFilter::::try_new(in_array)?)) + if let Some(filter) = branchless_filter::(in_array, non_null_count)? { + return Ok(Some(filter)); + } + + primitive_frozen_filter::(in_array, non_null_count) +} + +fn primitive_frozen_filter( + in_array: &ArrayRef, + non_null_count: usize, +) -> Result> +where + T: BranchlessFilterType, + BranchlessNative: Copy + Eq + Hash + Send + Sync, +{ + if non_null_count <= T::MAX_LIST_LEN { + return Ok(None); + } + + Ok(Some(Arc::new(PrimitiveFrozenFilter::::try_new( + in_array, + )?))) } fn branchless_filter( @@ -149,7 +193,7 @@ where T: BranchlessFilterType, BranchlessNative: Copy + PartialEq + Send + Sync, { - // Larger lists use the standard filter. `try_new` checks the limit again. + // Larger lists use another filter strategy. `try_new` checks the limit again. if non_null_count > T::MAX_LIST_LEN { return Ok(None); } @@ -159,7 +203,7 @@ where #[cfg(test)] mod tests { - use arrow::array::UInt32Array; + use arrow::array::{Decimal128Array, UInt32Array}; use arrow::datatypes::UInt32Type; use super::super::branchless_filter::BranchlessFilterType; @@ -176,21 +220,56 @@ mod tests { let values = (0..max_len) .map(|value| Some(value as u32)) .collect::>(); - assert!(instantiate_branchless_filter(&uint32_array(values))?.is_some()); + assert!( + branchless_filter::(&uint32_array(values), max_len)?.is_some() + ); + + let values = (0..=max_len) + .map(|value| Some(value as u32)) + .collect::>(); + assert!( + branchless_filter::(&uint32_array(values), max_len + 1)? + .is_none() + ); + + Ok(()) + } + + #[test] + fn primitive_frozen_routing_starts_after_max_list_len() -> Result<()> { + let max_len = ::MAX_LIST_LEN; + + let values = (0..max_len) + .map(|value| Some(value as u32)) + .collect::>(); + assert!( + primitive_frozen_filter::(&uint32_array(values), max_len)? + .is_none() + ); let values = (0..=max_len) .map(|value| Some(value as u32)) .collect::>(); - assert!(instantiate_branchless_filter(&uint32_array(values))?.is_none()); + assert!( + primitive_frozen_filter::(&uint32_array(values), max_len + 1)? + .is_some() + ); Ok(()) } + #[test] + fn primitive_frozen_routing_excludes_128_bit_values() -> Result<()> { + let array: ArrayRef = Arc::new(Decimal128Array::from(vec![1, 2, 3, 4, 5])); + assert!(instantiate_primitive_filter(&array)?.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()); + assert!(branchless_filter::(&array, 0)?.is_some()); Ok(()) }