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); } } } 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); }