From 2fba1c8974154890d269d232f7152dc64197d101 Mon Sep 17 00:00:00 2001 From: Geoffrey Claude Date: Wed, 5 Aug 2026 13:58:36 +0200 Subject: [PATCH 1/4] Add FixedSizeBinary IN LIST benchmarks --- .../physical-expr/benches/in_list_strategy.rs | 69 ++++++++++++------- 1 file changed, 45 insertions(+), 24 deletions(-) diff --git a/datafusion/physical-expr/benches/in_list_strategy.rs b/datafusion/physical-expr/benches/in_list_strategy.rs index c69af192b9cdd..5038f6d298cb3 100644 --- a/datafusion/physical-expr/benches/in_list_strategy.rs +++ b/datafusion/physical-expr/benches/in_list_strategy.rs @@ -44,14 +44,14 @@ //! | Utf8View length-12 cases | Utf8View | 12-byte strings | 16, 64 | //! | Utf8View long-string cases | Utf8View | 24-byte strings | 4, 16, 64, 256 | //! | 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 | +//! | Fixed-size binary cases | FixedSizeBinary(1), FixedSizeBinary(2), FixedSizeBinary(16) | fixed-width binary values | 16 (1 byte), 64 (2 bytes), 4/64/256/10000 (16 bytes) | use arrow::array::types::IntervalMonthDayNano; use arrow::array::*; use arrow::datatypes::{Field, Int32Type, IntervalMonthDayNanoType, Schema}; use arrow::record_batch::RecordBatch; use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; -use datafusion_common::ScalarValue; +use datafusion_common::{HashSet, ScalarValue}; use datafusion_physical_expr::expressions::{col, in_list, lit}; use half::f16; use rand::distr::Alphanumeric; @@ -996,32 +996,48 @@ fn bench_nulls(c: &mut Criterion) { } // ============================================================================= -// FIXED SIZE BINARY BENCHMARKS (FixedSizeBinary<16>, e.g. UUIDs) +// FIXED SIZE BINARY BENCHMARKS // ============================================================================= -/// Generates a random 16-byte value (UUID-sized). -fn random_fixed_binary_16(rng: &mut StdRng) -> Vec { - let mut buf = vec![0u8; 16]; +fn random_fixed_binary(rng: &mut StdRng, width: i32) -> Vec { + let mut buf = vec![0u8; width as usize]; rng.fill(&mut buf[..]); buf } -/// Benchmarks FixedSizeBinary(16) IN list evaluation. /// FixedSizeBinary doesn't use the generic numeric helpers since its array /// construction differs from primitive types. fn bench_fixed_size_binary_inner( c: &mut Criterion, - name: &str, + width: i32, list_size: usize, - match_rate: f64, + match_pct: u32, ) { - let seed = 0xF1ED_B1A7_u64.wrapping_add(list_size as u64 * 0x6666); + assert!(match_pct <= 100); + if let Some(domain_size) = match width { + 1 => Some(1_usize << 8), + 2 => Some(1_usize << 16), + _ => None, + } { + // The input generator needs at least one value outside the list. + assert!(list_size < domain_size); + } + let match_rate = f64::from(match_pct) / 100.0; + + let seed = 0xF1ED_B1A7_u64 + .wrapping_add(list_size as u64 * 0x6666) + .wrapping_add(width as u64 * 0x7777); let mut rng = StdRng::seed_from_u64(seed); - // Generate IN list values (16-byte each) - let haystack: Vec> = (0..list_size) - .map(|_| random_fixed_binary_16(&mut rng)) - .collect(); + // Keep the number of distinct values equal to the configured list size. + let mut haystack_set = HashSet::with_capacity(list_size); + let mut haystack = Vec::with_capacity(list_size); + while haystack.len() < list_size { + let value = random_fixed_binary(&mut rng, width); + if haystack_set.insert(value.clone()) { + haystack.push(value); + } + } // Generate array with controlled match rate let values: Vec> = (0..ARRAY_SIZE) @@ -1029,7 +1045,12 @@ fn bench_fixed_size_binary_inner( if !haystack.is_empty() && rng.random_bool(match_rate) { haystack.choose(&mut rng).unwrap().clone() } else { - random_fixed_binary_16(&mut rng) + loop { + let value = random_fixed_binary(&mut rng, width); + if !haystack_set.contains(&value) { + break value; + } + } } }) .collect(); @@ -1040,28 +1061,28 @@ fn bench_fixed_size_binary_inner( let schema = Schema::new(vec![Field::new("a", array.data_type().clone(), true)]); let exprs: Vec<_> = haystack .iter() - .map(|v| lit(ScalarValue::FixedSizeBinary(16, Some(v.clone())))) + .map(|v| lit(ScalarValue::FixedSizeBinary(width, Some(v.clone())))) .collect(); let expr = in_list(col("a", &schema).unwrap(), exprs, &false, &schema).unwrap(); let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(array) as ArrayRef]) .unwrap(); c.bench_with_input( - BenchmarkId::new("fixed_size_binary", name), + BenchmarkId::new( + "fixed_size_binary", + format!("fsb{width}/list={list_size}/match={match_pct}%"), + ), &batch, |b, batch| b.iter(|| expr.evaluate(batch).unwrap()), ); } fn bench_fixed_size_binary(c: &mut Criterion) { - for list_size in [4, 64, 256, 10000] { + for (width, list_size) in + [(1, 16), (2, 64), (16, 4), (16, 64), (16, 256), (16, 10000)] + { for match_pct in MATCH_RATES { - bench_fixed_size_binary_inner( - c, - &format!("fsb16/list={list_size}/match={match_pct}%"), - list_size, - match_pct as f64 / 100.0, - ); + bench_fixed_size_binary_inner(c, width, list_size, match_pct); } } } From 89d1595753e5f7722837e00cc43c7ba932328fe4 Mon Sep 17 00:00:00 2001 From: Geoffrey Claude Date: Wed, 5 Aug 2026 13:59:03 +0200 Subject: [PATCH 2/4] Optimize IN LIST for fixed-size binary arrays --- .../physical-expr/src/expressions/in_list.rs | 34 ++ .../in_list/fixed_size_binary_filter.rs | 467 ++++++++++++++++++ .../src/expressions/in_list/strategy.rs | 5 + 3 files changed, 506 insertions(+) create mode 100644 datafusion/physical-expr/src/expressions/in_list/fixed_size_binary_filter.rs diff --git a/datafusion/physical-expr/src/expressions/in_list.rs b/datafusion/physical-expr/src/expressions/in_list.rs index 874e149b58328..b00dfd98c03d1 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 fixed_size_binary_filter; mod primitive_filter; mod result; mod static_filter; @@ -3548,6 +3549,39 @@ mod tests { ); } + // FixedSizeBinary in_array, FixedSizeBinary and Dictionary needles + let fsb_in = Arc::new(FixedSizeBinaryArray::try_from_iter( + [ + [1, 2, 3, 4].as_slice(), + [5, 6, 7, 8].as_slice(), + [9, 10, 11, 12].as_slice(), + ] + .into_iter(), + )?) as ArrayRef; + let fsb_needle = Arc::new(FixedSizeBinaryArray::try_from_iter( + [ + [1, 2, 3, 4].as_slice(), + [13, 14, 15, 16].as_slice(), + [5, 6, 7, 8].as_slice(), + ] + .into_iter(), + )?) as ArrayRef; + assert_eq!( + expected, + eval_in_list_from_array(Arc::clone(&fsb_needle), Arc::clone(&fsb_in))? + ); + assert_eq!( + expected, + eval_in_list_from_array( + wrap_in_dict(Arc::clone(&fsb_needle)), + Arc::clone(&fsb_in), + )? + ); + assert_eq!( + expected, + eval_in_list_from_array(wrap_in_dict(fsb_needle), wrap_in_dict(fsb_in))? + ); + // Utf8 (falls through to ArrayStaticFilter) let utf8_in = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef; let utf8_needle = Arc::new(StringArray::from(vec!["a", "d", "b"])) as ArrayRef; diff --git a/datafusion/physical-expr/src/expressions/in_list/fixed_size_binary_filter.rs b/datafusion/physical-expr/src/expressions/in_list/fixed_size_binary_filter.rs new file mode 100644 index 0000000000000..38da1353212fa --- /dev/null +++ b/datafusion/physical-expr/src/expressions/in_list/fixed_size_binary_filter.rs @@ -0,0 +1,467 @@ +// 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. + +//! Optimized filters for fixed-size binary `IN` lists. +//! +//! Supported widths use an Arrow primitive type with the same in-memory size: +//! +//! | Width | Primitive type | Branchless through | Larger lists | +//! |------:|----------------|-------------------:|--------------| +//! | 1 | `UInt8` | 16 values | bitmap | +//! | 2 | `UInt16` | 8 values | bitmap | +//! | 4 | `UInt32` | 32 values | hash set | +//! | 8 | `UInt64` | 16 values | hash set | +//! | 16 | `Decimal128` | 4 values | hash set | +//! +//! The limits count non-null list values. +//! +//! The list and input bytes are read the same way, so each primitive value is +//! an exact key for comparison, bitmap lookup, or hashing. No arithmetic, +//! ordering, or decimal operations are used. +//! +//! Aligned Arrow buffers are reused without copying. Unaligned buffers are +//! copied into aligned primitive storage. + +use std::hash::Hash; +use std::marker::PhantomData; +use std::mem::{align_of, size_of}; +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, AsArray, BooleanArray, FixedSizeBinaryArray, PrimitiveArray, +}; +use arrow::buffer::{Buffer, ScalarBuffer}; +use arrow::datatypes::{ + ArrowPrimitiveType, DataType, Decimal128Type, UInt8Type, UInt16Type, UInt32Type, + UInt64Type, +}; +use datafusion_common::{HashSet, Result, exec_datafusion_err, internal_datafusion_err}; + +use super::branchless_filter::{ + BranchlessFilter, BranchlessFilterType, BranchlessNative, +}; +use super::primitive_filter::{BitmapFilter, BitmapFilterType}; +use super::result::build_in_list_result; +use super::static_filter::{StaticFilter, handle_dictionary}; + +type StaticFilterRef = Arc; + +fn use_branchless(count: usize) -> bool { + count <= T::MAX_LIST_LEN +} + +/// Reinterpret fixed-size binary values as same-width primitive values. +/// +/// Arrow buffers are normally sufficiently aligned, making this zero-copy. A +/// valid Arrow array can still be constructed from a sliced, unaligned buffer; +/// in that case, copy each value into aligned primitive storage. +fn as_primitive(array: &FixedSizeBinaryArray) -> Result> +where + T: ArrowPrimitiveType, +{ + let width = size_of::(); + if usize::try_from(array.value_length()).ok() != Some(width) { + return Err(internal_datafusion_err!( + "FixedSizeBinary filter: expected {width}-byte values, got {}", + array.value_length() + )); + } + + let source = array.values(); + let values = if source.as_ptr().align_offset(align_of::()) == 0 { + ScalarBuffer::new(source.clone(), 0, array.len()) + } else { + ScalarBuffer::new(Buffer::from(source.as_slice()), 0, array.len()) + }; + + Ok(PrimitiveArray::::new(values, array.nulls().cloned())) +} + +/// Generic hash-set membership for the 4-, 8-, and 16-byte representations. +/// +/// The standard primitive filters are concrete per Arrow type. This local +/// generic form lets the three supported widths share one implementation. +struct HashSetFilter { + null_count: usize, + values: HashSet, + _marker: PhantomData, +} + +impl HashSetFilter +where + T: ArrowPrimitiveType, + T::Native: Copy + Eq + Hash, +{ + fn new(in_array: &PrimitiveArray) -> Self { + let mut values = HashSet::with_capacity(in_array.len()); + for value in in_array.iter().flatten() { + values.insert(value); + } + + Self { + null_count: in_array.null_count(), + values, + _marker: PhantomData, + } + } +} + +impl StaticFilter for HashSetFilter +where + T: ArrowPrimitiveType + Send + Sync + 'static, + T::Native: Copy + Eq + Hash + Send + Sync, +{ + 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::().ok_or_else(|| { + internal_datafusion_err!("HashSetFilter: expected {} array", T::DATA_TYPE) + })?; + let input_values = v.values(); + Ok(build_in_list_result( + v.len(), + v.nulls(), + self.null_count > 0, + negated, + #[inline(always)] + |index| self.values.contains(&input_values[index]), + )) + } +} + +/// Adapts a primitive filter to concrete, same-width `FixedSizeBinary` arrays. +struct FixedSizeBinaryFilter { + data_type: DataType, + inner: StaticFilterRef, + _marker: PhantomData, +} + +impl FixedSizeBinaryFilter { + fn new(data_type: DataType, inner: StaticFilterRef) -> Self { + Self { + data_type, + inner, + _marker: PhantomData, + } + } +} + +impl StaticFilter for FixedSizeBinaryFilter +where + T: ArrowPrimitiveType + Send + Sync + 'static, +{ + fn null_count(&self) -> usize { + self.inner.null_count() + } + + fn contains(&self, v: &dyn Array, negated: bool) -> Result { + handle_dictionary!(self, v, negated); + + if v.data_type() != &self.data_type { + return Err(exec_datafusion_err!( + "FixedSizeBinary filter: expected {} array, got {}", + self.data_type, + v.data_type() + )); + } + let array = v.as_fixed_size_binary_opt().ok_or_else(|| { + exec_datafusion_err!( + "FixedSizeBinary filter: expected concrete {} array", + self.data_type + ) + })?; + let primitive = as_primitive::(array)?; + self.inner.contains(&primitive, negated) + } +} + +fn branchless_or_bitmap( + array: &FixedSizeBinaryArray, + count: usize, +) -> Result +where + T: BranchlessFilterType + BitmapFilterType, + BranchlessNative: Copy + Eq + Send + Sync, +{ + let primitive: ArrayRef = Arc::new(as_primitive::(array)?); + let inner: StaticFilterRef = if use_branchless::(count) { + Arc::new(BranchlessFilter::::try_new(&primitive)?) + } else { + Arc::new(BitmapFilter::::try_new(&primitive)?) + }; + Ok(Arc::new(FixedSizeBinaryFilter::::new( + array.data_type().clone(), + inner, + ))) +} + +fn branchless_or_hash_set( + array: &FixedSizeBinaryArray, + count: usize, +) -> Result +where + T: BranchlessFilterType, + T::Native: Copy + Eq + Hash + Send + Sync, + BranchlessNative: Copy + Eq + Send + Sync, +{ + let primitive = as_primitive::(array)?; + let inner: StaticFilterRef = if use_branchless::(count) { + let primitive: ArrayRef = Arc::new(primitive); + Arc::new(BranchlessFilter::::try_new(&primitive)?) + } else { + Arc::new(HashSetFilter::::new(&primitive)) + }; + Ok(Arc::new(FixedSizeBinaryFilter::::new( + array.data_type().clone(), + inner, + ))) +} + +/// Creates an optimized filter for supported concrete `FixedSizeBinary` arrays. +pub(super) fn instantiate_fixed_size_binary_filter( + in_array: &ArrayRef, +) -> Result> { + let DataType::FixedSizeBinary(width) = in_array.data_type() else { + return Ok(None); + }; + let Some(array) = in_array.as_fixed_size_binary_opt() else { + return Ok(None); + }; + + let count = array.len() - array.null_count(); + + let filter = match width { + 1 => branchless_or_bitmap::(array, count)?, + 2 => branchless_or_bitmap::(array, count)?, + 4 => branchless_or_hash_set::(array, count)?, + 8 => branchless_or_hash_set::(array, count)?, + 16 => branchless_or_hash_set::(array, count)?, + _ => return Ok(None), + }; + Ok(Some(filter)) +} + +#[cfg(test)] +mod tests { + use arrow::array::{DictionaryArray, Int8Array, StringArray}; + use arrow::buffer::{Buffer, NullBuffer}; + use arrow::datatypes::Int8Type; + + use super::*; + + fn value(width: i32, index: usize, miss: bool) -> Vec { + let mut value = (index as u128).to_le_bytes()[..width as usize].to_vec(); + let last = value.last_mut().unwrap(); + if miss { + *last |= 0x80; + } else { + *last &= 0x7f; + } + value + } + + fn array(width: i32, values: &[Option>]) -> FixedSizeBinaryArray { + FixedSizeBinaryArray::try_from_sparse_iter_with_size( + values.iter().map(|value| value.as_deref()), + width, + ) + .unwrap() + } + + fn make_filter(width: i32, values: &[Option>]) -> Result { + let in_array: ArrayRef = Arc::new(array(width, values)); + Ok(instantiate_fixed_size_binary_filter(&in_array)?.unwrap()) + } + + #[test] + fn filters_supported_widths_across_strategy_thresholds() -> Result<()> { + for (width, list_len) in [ + (1, 16), + (1, 17), + (2, 8), + (2, 9), + (4, 32), + (4, 33), + (8, 16), + (8, 17), + (16, 4), + (16, 5), + ] { + let mut hit = vec![0x80; width as usize]; + hit[width as usize - 1] = 0xff; + let mut miss = hit.clone(); + miss[width as usize - 1] ^= 1; + + let mut haystack = (0..list_len - 1) + .map(|index| Some(value(width, index, false))) + .collect::>(); + haystack.push(Some(hit.clone())); + let filter = make_filter(width, &haystack)?; + let needles = array(width, &[Some(hit), Some(miss), None]); + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![Some(true), Some(false), None]), + "width={width}, list_len={list_len}" + ); + } + Ok(()) + } + + #[test] + fn handles_slices_nulls_and_not_in() -> Result<()> { + let width = 16; + let parent = array( + width, + &[ + Some(value(width, 0, false)), + Some(value(width, 1, false)), + None, + Some(value(width, 2, false)), + Some(value(width, 3, false)), + Some(value(width, 4, false)), + Some(value(width, 5, false)), + Some(value(width, 6, false)), + ], + ); + // Five non-null values select the hash-set path. + let in_array: ArrayRef = Arc::new(parent.slice(1, 6)); + let filter = instantiate_fixed_size_binary_filter(&in_array)?.unwrap(); + let needles = array( + width, + &[ + Some(value(width, 2, false)), + Some(value(width, 7, false)), + None, + ], + ); + + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![Some(true), None, None]) + ); + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![Some(false), None, None]) + ); + Ok(()) + } + + #[test] + fn handles_dictionary_needles() -> Result<()> { + let filter = make_filter(4, &[Some(value(4, 7, false))])?; + let dictionary_values: ArrayRef = Arc::new(array( + 4, + &[Some(value(4, 7, false)), Some(value(4, 8, false))], + )); + let keys = Int8Array::from(vec![Some(0), Some(1), None]); + let needles = + DictionaryArray::::try_new(keys, dictionary_values).unwrap(); + + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![Some(true), Some(false), None]) + ); + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![Some(false), Some(true), None]) + ); + Ok(()) + } + + #[test] + fn rejects_unsupported_arrays() -> Result<()> { + let filter = make_filter(4, &[Some(value(4, 1, false))])?; + let wrong_width = array(8, &[Some(value(8, 1, false))]); + let error = filter + .contains(&wrong_width, false) + .unwrap_err() + .to_string(); + assert!( + error.contains("expected FixedSizeBinary(4) array, got FixedSizeBinary(8)"), + "{error}" + ); + + let wrong_type = StringArray::from(vec!["one"]); + let error = filter.contains(&wrong_type, false).unwrap_err().to_string(); + assert!( + error.contains("expected FixedSizeBinary(4) array, got Utf8"), + "{error}" + ); + + for width in [0, 3, 5, 15, 17] { + let unsupported: ArrayRef = + Arc::new(FixedSizeBinaryArray::new_null(width, 1)); + assert!( + instantiate_fixed_size_binary_filter(&unsupported)?.is_none(), + "width={width}" + ); + } + + Ok(()) + } + + fn unaligned_array( + width: i32, + values: &[Vec], + nulls: Option, + ) -> FixedSizeBinaryArray { + let mut bytes = vec![0]; + bytes.extend(values.iter().flatten()); + let buffer = Buffer::from(bytes).slice(1); + assert_ne!( + buffer.as_ptr().align_offset(width as usize), + 0, + "test buffer must be unaligned" + ); + FixedSizeBinaryArray::new(width, buffer, nulls) + } + + #[test] + fn handles_aligned_and_unaligned_buffers() -> Result<()> { + let buffer = Buffer::from_vec(vec![1_u64, 2, 3]); + let source_ptr = buffer.as_ptr(); + let array = FixedSizeBinaryArray::new(8, buffer, None); + let primitive = as_primitive::(&array)?; + assert_eq!(primitive.values().inner().as_ptr(), source_ptr); + + let width = 16; + let haystack_values = (0..5) + .map(|index| value(width, index, false)) + .collect::>(); + let haystack: ArrayRef = Arc::new(unaligned_array(width, &haystack_values, None)); + let needles = unaligned_array( + width, + &[ + value(width, 3, false), + value(width, 8, true), + value(width, 9, true), + ], + Some(NullBuffer::from(vec![true, false, true])), + ); + let filter = instantiate_fixed_size_binary_filter(&haystack)?.unwrap(); + + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![Some(true), 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 d5ca8154a92f6..69319b2f1451b 100644 --- a/datafusion/physical-expr/src/expressions/in_list/strategy.rs +++ b/datafusion/physical-expr/src/expressions/in_list/strategy.rs @@ -34,6 +34,7 @@ use super::array_static_filter::ArrayStaticFilter; use super::branchless_filter::{ BranchlessFilter, BranchlessFilterType, BranchlessNative, }; +use super::fixed_size_binary_filter::instantiate_fixed_size_binary_filter; use super::primitive_filter::*; use super::static_filter::StaticFilter; @@ -42,6 +43,10 @@ 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_fixed_size_binary_filter(&in_array)? { + return Ok(filter); + } + if let Some(filter) = instantiate_branchless_filter(&in_array)? { return Ok(filter); } From b0443b6bb0411d05d8ab5eae0926599ee502a37d Mon Sep 17 00:00:00 2001 From: Geoffrey Claude Date: Tue, 4 Aug 2026 18:09:49 +0200 Subject: [PATCH 3/4] Optimize IN LIST for byte view arrays with standard hash sets --- .../physical-expr/src/expressions/in_list.rs | 22 +- .../expressions/in_list/branchless_filter.rs | 25 +- .../expressions/in_list/byte_view_filter.rs | 309 ++++++++++++++++++ .../src/expressions/in_list/strategy.rs | 36 +- 4 files changed, 380 insertions(+), 12 deletions(-) create mode 100644 datafusion/physical-expr/src/expressions/in_list/byte_view_filter.rs diff --git a/datafusion/physical-expr/src/expressions/in_list.rs b/datafusion/physical-expr/src/expressions/in_list.rs index b00dfd98c03d1..c3362c1a9a8c9 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 byte_view_filter; mod fixed_size_binary_filter; mod primitive_filter; mod result; @@ -216,7 +217,7 @@ impl InListExpr { expr, list, negated, - Some(instantiate_static_filter(array)?), + Some(instantiate_static_filter(array, &expr_data_type)?), )) } @@ -243,7 +244,7 @@ impl InListExpr { // Try to create a static filter if all list expressions are constants let static_filter = match try_evaluate_constant_list(&list, schema)? { - Some(in_array) => Some(instantiate_static_filter(in_array)?), + Some(in_array) => Some(instantiate_static_filter(in_array, &expr_data_type)?), None => None, // Non-constant expressions, fall back to dynamic evaluation }; @@ -3610,6 +3611,23 @@ mod tests { )? ); + // Utf8View in_array, Utf8View and Dict(Utf8View) needles + let utf8view_in = + Arc::new(StringViewArray::from(vec!["a", "b", "c", "d", "e"])) as ArrayRef; + let utf8view_needle = + Arc::new(StringViewArray::from(vec!["a", "missing", "b"])) as ArrayRef; + assert_eq!( + expected, + eval_in_list_from_array( + Arc::clone(&utf8view_needle), + Arc::clone(&utf8view_in), + )? + ); + assert_eq!( + expected, + eval_in_list_from_array(wrap_in_dict(utf8view_needle), utf8view_in)? + ); + // Struct in_array, Struct needle: multi-column join let struct_fields = Fields::from(vec![ Field::new("c0", DataType::Utf8, true), 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..b9c5e10c138af 100644 --- a/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs @@ -71,7 +71,7 @@ use std::mem::size_of; use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, PrimitiveArray}; -use arrow::buffer::{BooleanBuffer, ScalarBuffer}; +use arrow::buffer::{BooleanBuffer, NullBuffer, ScalarBuffer}; use arrow::datatypes::*; use arrow::util::bit_iterator::BitIndexIterator; use datafusion_common::{Result, exec_datafusion_err, internal_datafusion_err}; @@ -244,6 +244,20 @@ where check_values, }) } + + /// Evaluates values that the caller has already checked use `T`'s physical + /// representation. `nulls`, when present, must cover the same slice. + #[inline] + pub(super) fn contains_raw_values( + &self, + input_values: &[BranchlessNative], + nulls: Option<&NullBuffer>, + negated: bool, + ) -> BooleanArray { + debug_assert!(nulls.is_none_or(|nulls| nulls.len() == input_values.len())); + let matches = (self.check_values)(self.in_list_values.as_ref(), input_values); + build_result_from_contains(nulls, self.null_count > 0, negated, matches) + } } impl StaticFilter for BranchlessFilter @@ -272,14 +286,7 @@ where exec_datafusion_err!("BranchlessFilter: expected {} array", T::DATA_TYPE) })?; let input_values = branchless_values::(v); - let matches = - (self.check_values)(self.in_list_values.as_ref(), input_values.as_ref()); - Ok(build_result_from_contains( - v.nulls(), - self.null_count > 0, - negated, - matches, - )) + Ok(self.contains_raw_values(input_values.as_ref(), v.nulls(), negated)) } } diff --git a/datafusion/physical-expr/src/expressions/in_list/byte_view_filter.rs b/datafusion/physical-expr/src/expressions/in_list/byte_view_filter.rs new file mode 100644 index 0000000000000..8a660c53c72b4 --- /dev/null +++ b/datafusion/physical-expr/src/expressions/in_list/byte_view_filter.rs @@ -0,0 +1,309 @@ +// 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. + +//! Filters for `Utf8View` and `BinaryView` `IN` lists whose non-null values are +//! at most 12 bytes. +//! +//! Arrow stores such values directly in a `u128` view: the low 32 bits contain +//! the length and the remaining bits contain zero-padded bytes. Equality of the +//! views is therefore equality of the values. Long views contain a prefix and +//! buffer location instead, so a list containing one uses the generic filter. +//! +//! Lists with up to four non-null values use the existing 128-bit branchless +//! filter. Larger lists use a `HashSet`. `Decimal128Type` only lets the +//! branchless filter compare the same 128 bits; no decimal operations are used. +//! A long input value cannot match an inline list value because its length is +//! part of the view. + +use std::marker::PhantomData; +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, AsArray, BooleanArray, GenericByteViewArray, MAX_INLINE_VIEW_LEN, + PrimitiveArray, +}; +use arrow::buffer::ScalarBuffer; +use arrow::datatypes::{ + BinaryViewType, ByteViewType, DataType, Decimal128Type, StringViewType, +}; +use arrow::util::bit_iterator::BitIndexIterator; +use datafusion_common::{HashSet, Result, exec_datafusion_err}; + +use super::branchless_filter::{BranchlessFilter, BranchlessFilterType}; +use super::result::build_in_list_result; +use super::static_filter::{StaticFilter, handle_dictionary}; + +#[inline(always)] +fn view_len(view: u128) -> u32 { + view as u32 +} + +fn downcast_byte_view( + array: &dyn Array, +) -> Result<&GenericByteViewArray> { + array.as_byte_view_opt::().ok_or_else(|| { + exec_datafusion_err!( + "Expected concrete {} array, got {}", + T::DATA_TYPE, + array.data_type() + ) + }) +} + +/// A byte-view array whose non-null values all use Arrow's inline layout. +struct InlineViews<'a, T: ByteViewType> { + array: &'a GenericByteViewArray, +} + +impl<'a, T: ByteViewType> InlineViews<'a, T> { + fn try_new(array: &'a ArrayRef) -> Result> { + let array = downcast_byte_view::(array.as_ref())?; + Ok(all_inline(array).then_some(Self { array })) + } +} + +fn all_inline(array: &GenericByteViewArray) -> bool { + let is_inline = |idx: usize| view_len(array.views()[idx]) <= MAX_INLINE_VIEW_LEN; + match array.nulls() { + Some(nulls) => { + BitIndexIterator::new(nulls.validity(), nulls.offset(), nulls.len()) + .all(is_inline) + } + None => (0..array.len()).all(is_inline), + } +} + +fn as_decimal128(inline: &InlineViews<'_, T>) -> ArrayRef { + let array = inline.array; + let views = array.views(); + // `views.inner()` is already sliced to the array's offset. + let values = ScalarBuffer::::new(views.inner().clone(), 0, views.len()); + Arc::new(PrimitiveArray::::new( + values, + array.nulls().cloned(), + )) +} + +fn branchless_filter( + inline: &InlineViews<'_, T>, +) -> Result> { + let values = as_decimal128(inline); + + Ok(Arc::new(ByteViewBranchless:: { + inner: BranchlessFilter::::try_new(&values)?, + _marker: PhantomData, + })) +} + +struct ByteViewBranchless { + inner: BranchlessFilter, + _marker: PhantomData, +} + +/// Exact set membership for inline views. +/// +/// Arrow requires unused inline bytes to be zero, so equal values have equal +/// `u128` views. +struct ByteViewHashSet { + set: HashSet, + null_count: usize, + _marker: PhantomData, +} + +impl ByteViewHashSet { + fn new(inline: &InlineViews<'_, T>) -> Self { + let array = inline.array; + let mut set = HashSet::with_capacity(array.len() - array.null_count()); + match array.nulls() { + Some(nulls) => { + BitIndexIterator::new(nulls.validity(), nulls.offset(), nulls.len()) + .for_each(|idx| { + set.insert(array.views()[idx]); + }); + } + None => set.extend(array.views().iter().copied()), + } + + Self { + set, + null_count: array.null_count(), + _marker: PhantomData, + } + } +} + +impl StaticFilter for ByteViewHashSet { + fn null_count(&self) -> usize { + self.null_count + } + + fn contains(&self, v: &dyn Array, negated: bool) -> Result { + handle_dictionary!(self, v, negated); + let array = downcast_byte_view::(v)?; + let views = array.views(); + Ok(build_in_list_result( + array.len(), + array.nulls(), + self.null_count > 0, + negated, + |idx| { + // SAFETY: `build_in_list_result` visits indices in `0..array.len()`. + self.set.contains(unsafe { views.get_unchecked(idx) }) + }, + )) + } +} + +impl StaticFilter for ByteViewBranchless { + fn null_count(&self) -> usize { + self.inner.null_count() + } + + fn contains(&self, v: &dyn Array, negated: bool) -> Result { + handle_dictionary!(self, v, negated); + let array = downcast_byte_view::(v)?; + let values: &[i128] = array.views().inner().typed_data(); + Ok(self + .inner + .contains_raw_values(values, array.nulls(), negated)) + } +} + +fn use_branchless(count: usize) -> bool { + count <= ::MAX_LIST_LEN +} + +fn instantiate_typed_filter( + in_array: &ArrayRef, +) -> Result>> { + let Some(inline) = InlineViews::::try_new(in_array)? else { + return Ok(None); + }; + + let count = inline.array.len() - inline.array.null_count(); + if use_branchless(count) { + branchless_filter(&inline).map(Some) + } else { + Ok(Some(Arc::new(ByteViewHashSet::new(&inline)))) + } +} + +/// Returns a filter when every non-null byte view is inline. +pub(super) fn instantiate_byte_view_filter( + in_array: &ArrayRef, +) -> Result>> { + match in_array.data_type() { + DataType::Utf8View => instantiate_typed_filter::(in_array), + DataType::BinaryView => instantiate_typed_filter::(in_array), + _ => Ok(None), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{BinaryViewArray, StringViewArray}; + use arrow::datatypes::StringViewType; + + fn assert_contains( + filter: &dyn StaticFilter, + needles: &dyn Array, + expected: Vec>, + ) -> Result<()> { + assert_eq!( + filter.contains(needles, false)?, + BooleanArray::from(expected) + ); + Ok(()) + } + + #[test] + fn hash_set_nulls_and_slices() -> Result<()> { + let haystack: ArrayRef = Arc::new( + StringViewArray::from(vec![ + Some("outside"), + Some("a"), + Some("b"), + None, + Some("c"), + Some("d"), + Some("e"), + Some("tail"), + ]) + .slice(1, 6), + ); + let inline = InlineViews::::try_new(&haystack)? + .expect("all non-null values are inline"); + let filter = ByteViewHashSet::new(&inline); + let needles = + StringViewArray::from(vec![Some("b"), Some("missing"), None, Some("e")]); + + 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 routing_and_boundaries() -> Result<()> { + assert!(use_branchless(4)); + assert!(!use_branchless(5)); + + let inline = "abcdefghijkl"; + let long = "abcdefghijklm"; + let values = StringViewArray::from(vec![ + "outside", inline, "one", "two", "three", "four", "tail", + ]); + let needles = StringViewArray::from(vec![inline, long, "missing"]); + + // Covers the branchless/hash-set cutoff and a non-zero view offset. + for haystack in [ + Arc::new(values.slice(1, 4)) as ArrayRef, + Arc::new(values.slice(1, 5)) as ArrayRef, + ] { + let filter = instantiate_byte_view_filter(&haystack)?.unwrap(); + assert_contains( + &*filter, + &needles, + vec![Some(true), Some(false), Some(false)], + )?; + } + + let mixed: ArrayRef = Arc::new(StringViewArray::from(vec!["short", long])); + assert!(instantiate_byte_view_filter(&mixed)?.is_none()); + + let all_null: ArrayRef = Arc::new(StringViewArray::from(vec![None::<&str>])); + let filter = instantiate_byte_view_filter(&all_null)?.unwrap(); + let needle = StringViewArray::from(vec!["missing"]); + assert_eq!(filter.contains(&needle, false)?, BooleanArray::new_null(1)); + assert_eq!(filter.contains(&needle, true)?, BooleanArray::new_null(1)); + + let binary: ArrayRef = Arc::new(BinaryViewArray::from(vec![ + b"a".as_slice(), + b"b".as_slice(), + b"c".as_slice(), + b"d".as_slice(), + b"e".as_slice(), + ])); + let filter = instantiate_byte_view_filter(&binary)?.unwrap(); + let needles = BinaryViewArray::from(vec![b"a".as_slice(), b"z".as_slice()]); + assert_contains(&*filter, &needles, vec![Some(true), 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 69319b2f1451b..ac456dbbdfa20 100644 --- a/datafusion/physical-expr/src/expressions/in_list/strategy.rs +++ b/datafusion/physical-expr/src/expressions/in_list/strategy.rs @@ -34,15 +34,25 @@ use super::array_static_filter::ArrayStaticFilter; use super::branchless_filter::{ BranchlessFilter, BranchlessFilterType, BranchlessNative, }; +use super::byte_view_filter::instantiate_byte_view_filter; use super::fixed_size_binary_filter::instantiate_fixed_size_binary_filter; use super::primitive_filter::*; use super::static_filter::StaticFilter; type StaticFilterRef = Arc; -pub(super) fn instantiate_static_filter(in_array: ArrayRef) -> Result { +pub(super) fn instantiate_static_filter( + in_array: ArrayRef, + expr_data_type: &DataType, +) -> Result { let in_array = flatten_dictionary_haystack(in_array)?; + if view_types_match(expr_data_type, in_array.data_type()) + && let Some(filter) = instantiate_byte_view_filter(&in_array)? + { + return Ok(filter); + } + if let Some(filter) = instantiate_fixed_size_binary_filter(&in_array)? { return Ok(filter); } @@ -54,6 +64,20 @@ pub(super) fn instantiate_static_filter(in_array: ArrayRef) -> Result bool { + matches!(list_type, DataType::Utf8View | DataType::BinaryView) + && dictionary_value_type(expr_type) == list_type +} + +fn dictionary_value_type(mut data_type: &DataType) -> &DataType { + while let DataType::Dictionary(_, value_type) = data_type { + data_type = value_type; + } + data_type +} + 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 @@ -174,6 +198,16 @@ mod tests { Arc::new(UInt32Array::from(values)) } + #[test] + fn byte_view_routing_requires_same_physical_type() { + let dict_view = + DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8View)); + + assert!(view_types_match(&DataType::Utf8View, &DataType::Utf8View)); + assert!(view_types_match(&dict_view, &DataType::Utf8View)); + assert!(!view_types_match(&DataType::Utf8, &DataType::Utf8View)); + } + #[test] fn branchless_routing_respects_max_list_len() -> Result<()> { let max_len = ::MAX_LIST_LEN; From 689cae2a4239b2a96795a0cfec5679bdb97228cf Mon Sep 17 00:00:00 2001 From: Geoffrey Claude Date: Sat, 8 Aug 2026 10:45:20 +0200 Subject: [PATCH 4/4] Optimize large integer IN LIST filters --- .../physical-expr/src/expressions/in_list.rs | 1 + .../src/expressions/in_list/integer_set.rs | 165 ++++++++++++ .../expressions/in_list/primitive_filter.rs | 239 +++++++----------- .../src/expressions/in_list/strategy.rs | 1 - 4 files changed, 256 insertions(+), 150 deletions(-) create mode 100644 datafusion/physical-expr/src/expressions/in_list/integer_set.rs diff --git a/datafusion/physical-expr/src/expressions/in_list.rs b/datafusion/physical-expr/src/expressions/in_list.rs index c3362c1a9a8c9..745c9f68c4921 100644 --- a/datafusion/physical-expr/src/expressions/in_list.rs +++ b/datafusion/physical-expr/src/expressions/in_list.rs @@ -40,6 +40,7 @@ mod array_static_filter; mod branchless_filter; mod byte_view_filter; mod fixed_size_binary_filter; +mod integer_set; mod primitive_filter; mod result; mod static_filter; diff --git a/datafusion/physical-expr/src/expressions/in_list/integer_set.rs b/datafusion/physical-expr/src/expressions/in_list/integer_set.rs new file mode 100644 index 0000000000000..7aceb9bb542df --- /dev/null +++ b/datafusion/physical-expr/src/expressions/in_list/integer_set.rs @@ -0,0 +1,165 @@ +// 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. + +//! Integer membership with a byte index and hash fallback. + +use std::{hash::Hash, mem::size_of}; + +use arrow::buffer::BooleanBuffer; +use datafusion_common::HashSet; + +const BUCKETS: usize = 256; +const SLOTS: usize = 4; + +pub(super) trait IntegerKey: Copy + Eq + Hash { + fn byte(self, index: usize) -> usize; +} + +macro_rules! integer_key { + ($($t:ty),+ $(,)?) => { + $( + impl IntegerKey for $t { + #[inline(always)] + fn byte(self, index: usize) -> usize { + ((self as u64 >> (index * 8)) & 0xff) as usize + } + } + )+ + }; +} + +integer_key!(i32, u32, i64, u64); + +pub(super) enum IntegerSet { + Indexed { + byte: usize, + buckets: Box<[[V; SLOTS]; BUCKETS]>, + }, + Hash(HashSet), +} + +impl IntegerSet { + pub(super) fn new(input: impl IntoIterator) -> Self { + let input = input.into_iter(); + let mut set = HashSet::with_capacity(input.size_hint().1.unwrap_or(0)); + set.extend(input); + if !set.is_empty() && set.len() <= BUCKETS * SLOTS { + for byte in 0..size_of::() { + let mut counts = [0_u16; BUCKETS]; + for &value in &set { + counts[value.byte(byte)] += 1; + } + if counts.iter().all(|&count| count <= SLOTS as u16) { + let sentinel = *set.iter().next().expect("set is not empty"); + // Padding cannot match a value from another bucket. + let mut buckets = Box::new([[sentinel; SLOTS]; BUCKETS]); + let mut next = [0_u8; BUCKETS]; + for &value in &set { + let bucket = value.byte(byte); + buckets[bucket][next[bucket] as usize] = value; + next[bucket] += 1; + } + return Self::Indexed { byte, buckets }; + } + } + } + Self::Hash(set) + } + + pub(super) fn contains_values(&self, values: &[V]) -> BooleanBuffer { + match self { + Self::Indexed { byte, buckets } => { + BooleanBuffer::collect_bool(values.len(), |i| { + let value = values[i]; + let [a, b, c, d] = buckets[value.byte(*byte)]; + // Compare all four values so the compiler can use SIMD. + (value == a) | (value == b) | (value == c) | (value == d) + }) + } + Self::Hash(set) => { + BooleanBuffer::collect_bool(values.len(), |i| set.contains(&values[i])) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn assert_matches(values: &[V], needles: &[V]) { + let expected = values + .iter() + .copied() + .collect::>(); + let actual = IntegerSet::new(values.to_vec()).contains_values(needles); + for (i, value) in needles.iter().enumerate() { + assert_eq!(actual.value(i), expected.contains(value), "value {value:?}"); + } + } + + #[test] + fn matches_standard_set() { + let mut state = 1_u64; + for len in [0, 1, 7, 32, 129, 256, 257, 1025] { + let values = (0..len) + .map(|_| { + state = state + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1); + state % 2049 + }) + .collect::>(); + let needles = (0..2200).collect::>(); + assert_matches(&values, &needles); + } + } + + #[test] + fn handles_bounds_and_duplicates() { + assert_matches( + &[i32::MIN, -1, 0, 1, i32::MAX, 0], + &[i32::MIN, i32::MIN + 1, -1, 0, 1, i32::MAX], + ); + assert_matches( + &[i64::MIN, -1, 0, 1, i64::MAX, 0], + &[i64::MIN, i64::MIN + 1, -1, 0, 1, i64::MAX], + ); + assert_matches(&[u32::MIN, 1, u32::MAX, 1], &[u32::MIN, 1, 2, u32::MAX]); + assert_matches(&[u64::MIN, 1, u64::MAX, 1], &[u64::MIN, 1, 2, u64::MAX]); + } + + #[test] + fn selects_index_or_hash() { + assert!(matches!( + IntegerSet::new(vec![0_u32, 256, 512, 768, 1024]), + IntegerSet::Indexed { byte: 1, .. } + )); + + let mut values = Vec::with_capacity(625); + for a in 0..5 { + for b in 0..5 { + for c in 0..5 { + for d in 0..5 { + values.push(a | b << 8 | c << 16 | d << 24); + } + } + } + } + assert!(matches!(IntegerSet::new(values), IntegerSet::Hash(_))); + } +} 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..2847725affabd 100644 --- a/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs @@ -20,13 +20,14 @@ //! This module provides membership tests for Arrow primitive types. use arrow::array::{Array, ArrayRef, AsArray, BooleanArray}; -use arrow::buffer::{BooleanBuffer, NullBuffer}; +use arrow::buffer::BooleanBuffer; use arrow::datatypes::*; use arrow::util::bit_iterator::BitIndexIterator; use datafusion_common::{HashSet, Result, exec_datafusion_err}; -use std::hash::{Hash, Hasher}; +use std::iter::FromIterator; -use super::result::build_in_list_result; +use super::integer_set::IntegerSet; +use super::result::{build_in_list_result, build_result_from_contains}; use super::static_filter::{StaticFilter, handle_dictionary}; /// Storage for the bits used by [`BitmapFilter`]. @@ -222,70 +223,29 @@ 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); - } -} - -impl PartialEq for OrderedFloat32 { - fn eq(&self, other: &Self) -> bool { - self.0.to_bits() == other.0.to_bits() - } -} - -impl Eq for OrderedFloat32 {} - -impl From for OrderedFloat32 { - fn from(v: f32) -> Self { - Self(v) - } -} - -/// 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() - } -} - -impl Eq for OrderedFloat64 {} - -impl From for OrderedFloat64 { - fn from(v: f64) -> Self { - Self(v) - } -} - // 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 + |v| v, + IntegerSet<<$ArrowType as ArrowPrimitiveType>::Native>, + IntegerSet::new, + IntegerSet::contains_values ); }; - ($Name:ident, $ArrowType:ty, $SetValueType:ty, $to_set_value:expr) => { + ( + $Name:ident, + $ArrowType:ty, + $to_set_value:expr, + $SetType:ty, + $new_set:expr, + $contains_values:expr + ) => { pub(super) struct $Name { null_count: usize, - values: HashSet<$SetValueType>, + values: $SetType, } impl $Name { @@ -298,12 +258,8 @@ macro_rules! primitive_static_filter { ) })?; - 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)); - } + let values = ($new_set)(in_array.iter().flatten().map($to_set_value)); Ok(Self { null_count, values }) } @@ -324,84 +280,14 @@ macro_rules! primitive_static_filter { ) })?; - 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)) + let contains = ($contains_values)(&self.values, needle_values); + Ok(build_result_from_contains( + v.nulls(), + self.null_count > 0, + negated, + contains, + )) } } }; @@ -412,17 +298,36 @@ 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); - }; +#[inline] +fn float32_contains(values: &HashSet<[u8; 4]>, needles: &[f32]) -> BooleanBuffer { + BooleanBuffer::collect_bool(needles.len(), |i| { + values.contains(&needles[i].to_ne_bytes()) + }) +} + +#[inline] +fn float64_contains(values: &HashSet<[u8; 8]>, needles: &[f64]) -> BooleanBuffer { + BooleanBuffer::collect_bool(needles.len(), |i| { + values.contains(&needles[i].to_ne_bytes()) + }) } -// Generate specialized filters for float types using ordered wrappers -float_static_filter!(Float32StaticFilter, Float32Type, OrderedFloat32); -float_static_filter!(Float64StaticFilter, Float64Type, OrderedFloat64); +primitive_static_filter!( + Float32StaticFilter, + Float32Type, + f32::to_ne_bytes, + HashSet<[u8; 4]>, + HashSet::from_iter, + float32_contains +); +primitive_static_filter!( + Float64StaticFilter, + Float64Type, + f64::to_ne_bytes, + HashSet<[u8; 8]>, + HashSet::from_iter, + float64_contains +); #[cfg(test)] mod tests { @@ -430,7 +335,8 @@ mod tests { use std::sync::Arc; use arrow::array::{ - DictionaryArray, Float16Array, Int8Array, Int16Array, UInt8Array, UInt16Array, + DictionaryArray, Float16Array, Float32Array, Int8Array, Int16Array, UInt8Array, + UInt16Array, UInt32Array, }; use half::f16; @@ -584,4 +490,39 @@ mod tests { Ok(()) } + + #[test] + fn static_filter_handles_slices_nulls_and_dictionaries() -> Result<()> { + let mut values = (0..36).map(Some).collect::>(); + values[2] = None; + let haystack: ArrayRef = Arc::new(UInt32Array::from(values).slice(1, 34)); + let filter = UInt32StaticFilter::try_new(&haystack)?; + + let keys = Int8Array::from(vec![Some(0), Some(1), None, Some(2)]); + let values = Arc::new(UInt32Array::from(vec![1, 35, 34])); + let needles = DictionaryArray::try_new(keys, values)?; + 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 static_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 = Float32StaticFilter::try_new(&haystack)?; + let needles = + Float32Array::from(vec![Some(0.0), Some(-0.0), Some(nan_a), Some(nan_b)]); + + assert_contains( + &filter, + &needles, + vec![Some(false), Some(true), Some(true), Some(false)], + ) + } } diff --git a/datafusion/physical-expr/src/expressions/in_list/strategy.rs b/datafusion/physical-expr/src/expressions/in_list/strategy.rs index ac456dbbdfa20..93dbedd08a603 100644 --- a/datafusion/physical-expr/src/expressions/in_list/strategy.rs +++ b/datafusion/physical-expr/src/expressions/in_list/strategy.rs @@ -152,7 +152,6 @@ fn instantiate_standard_filter(in_array: ArrayRef) -> Result { 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)?)), _ => {