diff --git a/datafusion/physical-expr/benches/in_list_strategy.rs b/datafusion/physical-expr/benches/in_list_strategy.rs index c69af192b9cdd..29432cbba246c 100644 --- a/datafusion/physical-expr/benches/in_list_strategy.rs +++ b/datafusion/physical-expr/benches/in_list_strategy.rs @@ -15,10 +15,10 @@ // specific language governing permissions and limitations // under the License. -//! Focused benchmarks for `InList` cases. +//! Benchmarks for static `IN LIST` filters. //! -//! This benchmark file adds targeted coverage for representative `IN LIST` -//! workloads with controlled parameters: +//! The cases control match rate and list size across several value types and +//! string layouts: //! //! - **Controlled match rates**: Exercises both hit-heavy and miss-heavy paths //! - **List size scaling**: Measures behavior across small and large `IN` lists @@ -27,7 +27,7 @@ //! - **Shared-prefix strings**: Adds collision-heavy string cases where values //! only differ late in the string //! - **Mixed-length strings**: Covers inputs that combine short and long values -//! - **Null handling**: Includes representative `NULL` and `NOT IN` cases +//! - **Null handling**: Covers `NULL` and `NOT IN` cases //! //! # Case Coverage //! @@ -38,24 +38,27 @@ //! | 32-bit primitive cases | Int32, Float32 | small and large lists | 4, 32, 64, 256 | //! | 64-bit primitive cases | Int64, TimestampNs | small and large lists | 4, 16, 32, 128 | //! | 128-bit interval cases | IntervalMonthDayNano | small lists | 4 | +//! | Decimal128 cases | Decimal128 | larger lists | 5, 64 | //! | Utf8 short-string cases | Utf8 | 8-byte strings | 4, 64, 256 | //! | Utf8 long-string cases | Utf8 | 24-byte strings | 4, 64, 256 | //! | Utf8View short-string cases | Utf8View | 8-byte strings | 4, 16, 64, 256 | //! | 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) | aligned values and unaligned 16-byte inputs | 16 (1 byte), 64 (2 bytes), 4/64/256/10000 (16 bytes) | use arrow::array::types::IntervalMonthDayNano; use arrow::array::*; +use arrow::buffer::Buffer; 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; use rand::prelude::*; +use std::mem::align_of; use std::sync::Arc; const ARRAY_SIZE: usize = 8192; @@ -463,6 +466,23 @@ fn bench_primitive(c: &mut Criterion) { } } + // Decimal128: benchmark the first hash-set list size (5) and a larger list (64). + for list_size in [5, 64] { + for match_pct in MATCH_RATES { + bench_numeric::( + c, + "primitive", + &format!("decimal128/large_list/list={list_size}/match={match_pct}%"), + &NumericBenchConfig::new( + list_size, + match_pct as f64 / 100.0, + |rng| i128::from(rng.random::()), + |v| ScalarValue::Decimal128(Some(v), 38, 10), + ), + ); + } + } + // NOT IN benchmark: test negated path bench_numeric::( c, @@ -852,7 +872,7 @@ fn bench_dictionary(c: &mut Criterion) { // NULL HANDLING BENCHMARKS // ============================================================================= // -// Tests representative null-containing inputs across primitive and string cases. +// Null-containing primitive and string cases. fn bench_nulls(c: &mut Criterion) { // ========================================================================= @@ -996,32 +1016,71 @@ 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. +fn unaligned_fixed_size_binary_16(values: &[Vec]) -> FixedSizeBinaryArray { + const WIDTH: usize = 16; + let alignment = align_of::(); + let payload_len = values.len() * WIDTH; + let mut bytes = vec![0_u8; payload_len + alignment]; + let offset = usize::from(bytes.as_ptr().align_offset(alignment) == 0); + for (target, value) in bytes[offset..offset + payload_len] + .chunks_exact_mut(WIDTH) + .zip(values.iter()) + { + assert_eq!(value.len(), WIDTH); + target.copy_from_slice(value); + } + let buffer = Buffer::from(bytes).slice_with_length(offset, payload_len); + assert_ne!( + buffer.as_ptr().align_offset(alignment), + 0, + "benchmark input must be unaligned" + ); + FixedSizeBinaryArray::new(WIDTH as i32, buffer, None) +} + /// 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, + unaligned_input: bool, ) { - 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,41 +1088,61 @@ 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(); + drop(haystack_set); - let refs: Vec<&[u8]> = values.iter().map(|v| v.as_slice()).collect(); - let array = FixedSizeBinaryArray::try_from_iter(refs.into_iter()).unwrap(); + let array = if unaligned_input { + assert_eq!(width, 16); + unaligned_fixed_size_binary_16(&values) + } else { + FixedSizeBinaryArray::try_from_iter(values.iter().map(Vec::as_slice)).unwrap() + }; 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", { + let name = format!("fsb{width}/list={list_size}/match={match_pct}%"); + if unaligned_input { + format!("{name}/input=unaligned") + } else { + name + } + }), &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, false); } } + + // At 16 bytes per value, an unaligned 8,192-row input copies 128 KiB per + // evaluation. List size 64 exercises the larger-list hash-set path. + for match_pct in MATCH_RATES { + bench_fixed_size_binary_inner(c, 16, 64, match_pct, true); + } } // ============================================================================= diff --git a/datafusion/physical-expr/src/expressions/in_list.rs b/datafusion/physical-expr/src/expressions/in_list.rs index 874e149b58328..154decfd8bb89 100644 --- a/datafusion/physical-expr/src/expressions/in_list.rs +++ b/datafusion/physical-expr/src/expressions/in_list.rs @@ -38,12 +38,13 @@ 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; mod strategy; -use static_filter::StaticFilter; +use static_filter::StaticFilterRef; use strategy::instantiate_static_filter; /// InList @@ -51,7 +52,7 @@ pub struct InListExpr { expr: Arc, list: Vec>, negated: bool, - static_filter: Option>, + static_filter: Option, } impl Debug for InListExpr { @@ -148,7 +149,7 @@ impl InListExpr { expr: Arc, list: Vec>, negated: bool, - static_filter: Option>, + static_filter: Option, ) -> Self { Self { expr, @@ -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 branchless, bitmap, or hash-set filters + /// when the list's physical representation supports them. /// /// Returns an error if data types don't match. If the list contains non-constant /// expressions, falls back to dynamic evaluation at runtime. @@ -2592,7 +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(), which selects a primitive static filter. let expr = in_list(col_a, list, &false, &schema)?; // Create dictionary-encoded batch with values [100, 200, 500] @@ -3548,6 +3549,38 @@ 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))? + ); + // The dictionary does not reference its second value, so that value + // must not become a member of the flattened list. + let dict_fsb_in = Arc::new(DictionaryArray::new( + Int32Array::from(vec![0, 2]), + Arc::clone(&fsb_in), + )); + assert_eq!( + BooleanArray::from(vec![Some(true), Some(false), Some(false)]), + eval_in_list_from_array(wrap_in_dict(fsb_needle), 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/branchless_filter.rs b/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs index cd0cbd0de59a8..8539dc7956dc4 100644 --- a/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs @@ -42,7 +42,7 @@ //! different NaN values. [`BranchlessFilterType`] defines these safe, //! same-sized mappings and checks their sizes at compile time. //! -//! The fast path is intentionally limited to short lists: +//! The fast path is limited to short lists: //! //! - 16 values for 1-byte types //! - 8 values for 2-byte types diff --git a/datafusion/physical-expr/src/expressions/in_list/fixed_size_binary_filter.rs b/datafusion/physical-expr/src/expressions/in_list/fixed_size_binary_filter.rs new file mode 100644 index 0000000000000..ab8376b06763b --- /dev/null +++ b/datafusion/physical-expr/src/expressions/in_list/fixed_size_binary_filter.rs @@ -0,0 +1,368 @@ +// 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 representation with the same +//! in-memory size: +//! +//! | Width | Primitive representation | +//! |------:|--------------------------| +//! | 1 | `UInt8` | +//! | 2 | `UInt16` | +//! | 4 | `UInt32` | +//! | 8 | `UInt64` | +//! | 16 | `Decimal128` | +//! +//! The shared primitive selector applies the native primitive branchless cutoffs +//! and chooses the bitmap or hash-set fallback. +//! +//! 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. +//! +//! Reinterpreting an aligned Arrow buffer is zero-copy. An unaligned buffer is +//! copied into aligned primitive storage before filter construction or probing. + +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::{Result, exec_datafusion_err, internal_datafusion_err}; + +use super::primitive_filter::instantiate_primitive_filter; +use super::static_filter::{StaticFilter, StaticFilterRef, handle_dictionary}; + +/// 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 reinterpret_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())) +} + +/// Adapts a primitive filter to concrete, same-width `FixedSizeBinary` arrays. +struct FixedSizeBinaryFilter { + data_type: DataType, + inner: StaticFilterRef, + _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 = reinterpret_as_primitive::(array)?; + self.inner.contains(&primitive, negated) + } +} + +fn instantiate_for_primitive(array: &FixedSizeBinaryArray) -> Result +where + T: ArrowPrimitiveType + Send + Sync + 'static, +{ + let primitive: ArrayRef = Arc::new(reinterpret_as_primitive::(array)?); + let inner = instantiate_primitive_filter(&primitive)?.ok_or_else(|| { + internal_datafusion_err!( + "FixedSizeBinary filter: no primitive filter for {}", + primitive.data_type() + ) + })?; + Ok(Arc::new(FixedSizeBinaryFilter:: { + data_type: array.data_type().clone(), + inner, + _marker: PhantomData, + })) +} + +/// 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 filter = match width { + 1 => instantiate_for_primitive::(array)?, + 2 => instantiate_for_primitive::(array)?, + 4 => instantiate_for_primitive::(array)?, + 8 => instantiate_for_primitive::(array)?, + 16 => instantiate_for_primitive::(array)?, + _ => 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, 0, false)), + Some(value(width, 6, false)), + Some(value(width, 7, false)), + None, + ], + ); + + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![Some(true), None, None, None, None]) + ); + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![Some(false), None, None, 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(align_of::()), + 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 = reinterpret_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/primitive_filter.rs b/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs index 8f8d9bad04afa..bfafae31ed622 100644 --- a/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs @@ -20,21 +20,133 @@ //! This module provides membership tests for Arrow primitive types. use arrow::array::{Array, ArrayRef, AsArray, BooleanArray}; -use arrow::buffer::{BooleanBuffer, NullBuffer}; use arrow::datatypes::*; use arrow::util::bit_iterator::BitIndexIterator; use datafusion_common::{HashSet, Result, exec_datafusion_err}; use std::hash::{Hash, Hasher}; +use std::marker::PhantomData; +use std::sync::Arc; +use super::branchless_filter::{BranchlessFilter, BranchlessFilterType}; use super::result::build_in_list_result; -use super::static_filter::{StaticFilter, handle_dictionary}; +use super::static_filter::{StaticFilter, StaticFilterRef, handle_dictionary}; + +/// Selects an optimized filter for a primitive representation. +/// +/// Supported short lists use a branchless filter. Larger supported lists use a +/// bitmap or hash-set filter. Returns `None` for primitive types without a +/// specialized filter. Representation adapters call this after conversion and +/// use the same cutoffs and fallback policy as native primitive arrays. +pub(super) fn instantiate_primitive_filter( + in_array: &ArrayRef, +) -> Result> { + if let Some(filter) = instantiate_branchless_filter(in_array)? { + return Ok(Some(filter)); + } + + let filter: StaticFilterRef = match in_array.data_type() { + DataType::Int8 => Arc::new(BitmapFilter::::try_new(in_array)?), + DataType::UInt8 => Arc::new(BitmapFilter::::try_new(in_array)?), + DataType::Int16 => Arc::new(BitmapFilter::::try_new(in_array)?), + DataType::UInt16 => Arc::new(BitmapFilter::::try_new(in_array)?), + DataType::Float16 => Arc::new(BitmapFilter::::try_new(in_array)?), + DataType::Int32 => { + Arc::new(PrimitiveHashSetFilter::::try_new(in_array)?) + } + DataType::Int64 => { + Arc::new(PrimitiveHashSetFilter::::try_new(in_array)?) + } + DataType::UInt32 => { + Arc::new(PrimitiveHashSetFilter::::try_new(in_array)?) + } + DataType::UInt64 => { + Arc::new(PrimitiveHashSetFilter::::try_new(in_array)?) + } + DataType::Decimal128(_, _) => { + Arc::new(PrimitiveHashSetFilter::::try_new(in_array)?) + } + // Float primitive types use ordered wrapper keys for Hash/Eq. + DataType::Float32 => Arc::new(PrimitiveHashSetFilter::< + Float32Type, + OrderedFloat32, + >::try_new(in_array)?), + DataType::Float64 => Arc::new(PrimitiveHashSetFilter::< + Float64Type, + OrderedFloat64, + >::try_new(in_array)?), + _ => return Ok(None), + }; + + Ok(Some(filter)) +} + +fn instantiate_branchless_filter(in_array: &ArrayRef) -> Result> { + let non_null_count = in_array.len() - in_array.null_count(); + + macro_rules! branchless { + ($arrow_type:ty) => {{ + // Larger lists use the standard primitive filter. `try_new` + // checks the limit again when called directly. + if non_null_count > <$arrow_type as BranchlessFilterType>::MAX_LIST_LEN { + Ok(None) + } else { + let filter: StaticFilterRef = + Arc::new(BranchlessFilter::<$arrow_type>::try_new(in_array)?); + Ok(Some(filter)) + } + }}; + } + + match in_array.data_type() { + DataType::Int8 => branchless!(Int8Type), + DataType::UInt8 => branchless!(UInt8Type), + DataType::Int16 => branchless!(Int16Type), + DataType::UInt16 => branchless!(UInt16Type), + DataType::Float16 => branchless!(Float16Type), + DataType::Int32 => branchless!(Int32Type), + DataType::UInt32 => branchless!(UInt32Type), + DataType::Float32 => branchless!(Float32Type), + DataType::Date32 => branchless!(Date32Type), + DataType::Time32(unit) => match unit { + TimeUnit::Second => branchless!(Time32SecondType), + TimeUnit::Millisecond => branchless!(Time32MillisecondType), + _ => Ok(None), + }, + DataType::Int64 => branchless!(Int64Type), + DataType::UInt64 => branchless!(UInt64Type), + DataType::Float64 => branchless!(Float64Type), + DataType::Date64 => branchless!(Date64Type), + DataType::Time64(unit) => match unit { + TimeUnit::Microsecond => branchless!(Time64MicrosecondType), + TimeUnit::Nanosecond => branchless!(Time64NanosecondType), + _ => Ok(None), + }, + DataType::Timestamp(unit, _) => match unit { + TimeUnit::Second => branchless!(TimestampSecondType), + TimeUnit::Millisecond => branchless!(TimestampMillisecondType), + TimeUnit::Microsecond => branchless!(TimestampMicrosecondType), + TimeUnit::Nanosecond => branchless!(TimestampNanosecondType), + }, + DataType::Duration(unit) => match unit { + TimeUnit::Second => branchless!(DurationSecondType), + TimeUnit::Millisecond => branchless!(DurationMillisecondType), + TimeUnit::Microsecond => branchless!(DurationMicrosecondType), + TimeUnit::Nanosecond => branchless!(DurationNanosecondType), + }, + DataType::Decimal128(_, _) => branchless!(Decimal128Type), + DataType::Interval(IntervalUnit::MonthDayNano) => { + branchless!(IntervalMonthDayNanoType) + } + _ => Ok(None), + } +} /// Storage for the bits used by [`BitmapFilter`]. /// /// `BitmapFilter` represents an `IN` list with one bit for each possible /// value, so membership checks become direct bit tests. This trait lets the /// same filter code use different storage sizes for different integer widths. -pub(super) trait BitmapStorage: Send + Sync { +trait BitmapStorage: Send + Sync { fn new_zeroed() -> Self; fn set_bit(&mut self, index: usize); fn get_bit(&self, index: usize) -> bool; @@ -81,9 +193,7 @@ impl BitmapStorage for Box<[u64; 1024]> { /// supplies the bitmap storage size and maps values to their bit-pattern index /// for the primitive domains that are small enough to represent with one bit /// per possible value. -pub(super) trait BitmapFilterType: - ArrowPrimitiveType + Send + Sync + 'static -{ +trait BitmapFilterType: ArrowPrimitiveType + Send + Sync + 'static { type Storage: BitmapStorage; /// Returns the index in the bitmap to check for this value. @@ -151,7 +261,7 @@ impl BitmapFilterType for Float16Type { /// the bit selected by each value. Evaluating input values checks the same bit /// position. Null handling and `NOT IN` inversion are handled by /// `build_in_list_result`. -pub(super) struct BitmapFilter { +struct BitmapFilter { null_count: usize, bits: T::Storage, } @@ -160,7 +270,7 @@ impl BitmapFilter where T: BitmapFilterType, { - pub(super) fn try_new(in_array: &ArrayRef) -> Result { + fn try_new(in_array: &ArrayRef) -> Result { let prim_array = in_array.as_primitive_opt::().ok_or_else(|| { exec_datafusion_err!("BitmapFilter: expected {} array", T::DATA_TYPE) })?; @@ -272,157 +382,77 @@ impl From for OrderedFloat64 { } } -// Macro to generate specialized StaticFilter implementations for primitive types -macro_rules! primitive_static_filter { - ($Name:ident, $ArrowType:ty) => { - primitive_static_filter!( - $Name, - $ArrowType, - <$ArrowType as ArrowPrimitiveType>::Native, - |v| v - ); - }; - ($Name:ident, $ArrowType:ty, $SetValueType:ty, $to_set_value:expr) => { - pub(super) struct $Name { - null_count: usize, - values: HashSet<$SetValueType>, - } - - impl $Name { - pub(super) fn try_new(in_array: &ArrayRef) -> Result { - let in_array = - in_array.as_primitive_opt::<$ArrowType>().ok_or_else(|| { - exec_datafusion_err!( - "Failed to downcast an array to a '{}' array", - stringify!($ArrowType) - ) - })?; - - let mut values = HashSet::with_capacity(in_array.len()); - let null_count = in_array.null_count(); - - for v in in_array.iter().flatten() { - values.insert(($to_set_value)(v)); - } +/// Hash-set membership for primitive types. +/// +/// `K` defaults to the Arrow type's native value. Floats use an ordered wrapper +/// key because their native values do not implement [`Eq`] and [`Hash`]. +struct PrimitiveHashSetFilter< + T: ArrowPrimitiveType, + K = ::Native, +> { + null_count: usize, + values: HashSet, + _marker: PhantomData, +} - Ok(Self { null_count, values }) - } +impl PrimitiveHashSetFilter +where + T: ArrowPrimitiveType, + T::Native: Copy, + K: From + Eq + Hash, +{ + fn try_new(in_array: &ArrayRef) -> Result { + let in_array = in_array.as_primitive_opt::().ok_or_else(|| { + exec_datafusion_err!( + "PrimitiveHashSetFilter: expected {} array", + T::DATA_TYPE + ) + })?; + let mut values = HashSet::with_capacity(in_array.len() - in_array.null_count()); + for value in in_array.iter().flatten() { + values.insert(K::from(value)); } - impl StaticFilter for $Name { - fn null_count(&self) -> usize { - self.null_count - } - - fn contains(&self, v: &dyn Array, negated: bool) -> Result { - handle_dictionary!(self, v, negated); - - let v = v.as_primitive_opt::<$ArrowType>().ok_or_else(|| { - exec_datafusion_err!( - "Failed to downcast an array to a '{}' array", - stringify!($ArrowType) - ) - })?; - - let haystack_has_nulls = self.null_count > 0; - let needle_values = v.values(); - let needle_nulls = v.nulls(); - let needle_has_nulls = v.null_count() > 0; - - // Truth table for `value [NOT] IN (set)` with SQL three-valued logic: - // ("-" means the value doesn't affect the result) - // - // | needle_null | haystack_null | negated | in set? | result | - // |-------------|---------------|---------|---------|--------| - // | true | - | false | - | null | - // | true | - | true | - | null | - // | false | true | false | yes | true | - // | false | true | false | no | null | - // | false | true | true | yes | false | - // | false | true | true | no | null | - // | false | false | false | yes | true | - // | false | false | false | no | false | - // | false | false | true | yes | false | - // | false | false | true | no | true | - - // Compute the "contains" result using collect_bool (fast batched approach) - // This ignores nulls - we handle them separately - let contains_buffer = if negated { - BooleanBuffer::collect_bool(needle_values.len(), |i| { - !self.values.contains(&($to_set_value)(needle_values[i])) - }) - } else { - BooleanBuffer::collect_bool(needle_values.len(), |i| { - self.values.contains(&($to_set_value)(needle_values[i])) - }) - }; - - // Compute the null mask - // Output is null when: - // 1. needle value is null, OR - // 2. needle value is not in set AND haystack has nulls - let result_nulls = match (needle_has_nulls, haystack_has_nulls) { - (false, false) => { - // No nulls anywhere - None - } - (true, false) => { - // Only needle has nulls - just use needle's null mask - needle_nulls.cloned() - } - (false, true) => { - // Only haystack has nulls - result is null when value not in set - // Valid (not null) when original "in set" is true - // For NOT IN: contains_buffer = !original, so validity = !contains_buffer - let validity = if negated { - !&contains_buffer - } else { - contains_buffer.clone() - }; - Some(NullBuffer::new(validity)) - } - (true, true) => { - // Both have nulls - combine needle nulls with haystack-induced nulls - let needle_validity = - needle_nulls.map(|n| n.inner().clone()).unwrap_or_else( - || BooleanBuffer::new_set(needle_values.len()), - ); - - // Valid when original "in set" is true (see above) - let haystack_validity = if negated { - !&contains_buffer - } else { - contains_buffer.clone() - }; - - // Combined validity: valid only where both are valid - let combined_validity = &needle_validity & &haystack_validity; - Some(NullBuffer::new(combined_validity)) - } - }; - - Ok(BooleanArray::new(contains_buffer, result_nulls)) - } - } - }; + Ok(Self { + null_count: in_array.null_count(), + values, + _marker: PhantomData, + }) + } } -primitive_static_filter!(Int32StaticFilter, Int32Type); -primitive_static_filter!(Int64StaticFilter, Int64Type); -primitive_static_filter!(UInt32StaticFilter, UInt32Type); -primitive_static_filter!(UInt64StaticFilter, UInt64Type); +impl StaticFilter for PrimitiveHashSetFilter +where + T: ArrowPrimitiveType + Send + Sync + 'static, + T::Native: Copy + Send + Sync, + K: From + Eq + Hash + Send + Sync + 'static, +{ + fn null_count(&self) -> usize { + self.null_count + } -// Macro to generate specialized StaticFilter implementations for float types -// Floats require a wrapper type (OrderedFloat*) to implement Hash/Eq due to NaN semantics -macro_rules! float_static_filter { - ($Name:ident, $ArrowType:ty, $OrderedType:ty) => { - primitive_static_filter!($Name, $ArrowType, $OrderedType, <$OrderedType>::from); - }; -} + fn contains(&self, v: &dyn Array, negated: bool) -> Result { + handle_dictionary!(self, v, negated); -// Generate specialized filters for float types using ordered wrappers -float_static_filter!(Float32StaticFilter, Float32Type, OrderedFloat32); -float_static_filter!(Float64StaticFilter, Float64Type, OrderedFloat64); + let v = v.as_primitive_opt::().ok_or_else(|| { + exec_datafusion_err!( + "PrimitiveHashSetFilter: expected {} array", + T::DATA_TYPE + ) + })?; + let input_values = v.values(); + Ok(build_in_list_result( + v.len(), + v.nulls(), + self.null_count > 0, + negated, + |index| { + let key = K::from(input_values[index]); + self.values.contains(&key) + }, + )) + } +} #[cfg(test)] mod tests { @@ -430,10 +460,15 @@ mod tests { use std::sync::Arc; use arrow::array::{ - DictionaryArray, Float16Array, Int8Array, Int16Array, UInt8Array, UInt16Array, + DictionaryArray, Float16Array, Float32Array, Float64Array, Int8Array, Int16Array, + UInt8Array, UInt16Array, UInt32Array, }; use half::f16; + fn uint32_array(values: Vec>) -> ArrayRef { + Arc::new(UInt32Array::from(values)) + } + fn assert_contains( filter: &dyn StaticFilter, needles: &dyn Array, @@ -446,6 +481,60 @@ mod tests { Ok(()) } + #[test] + fn branchless_routing_respects_max_list_len() -> Result<()> { + let max_len = ::MAX_LIST_LEN; + + let values = (0..max_len) + .map(|value| Some(value as u32)) + .collect::>(); + assert!(instantiate_branchless_filter(&uint32_array(values))?.is_some()); + + let values = (0..=max_len) + .map(|value| Some(value as u32)) + .collect::>(); + assert!(instantiate_branchless_filter(&uint32_array(values))?.is_none()); + + Ok(()) + } + + #[test] + fn branchless_routing_handles_zero_non_null_values() -> Result<()> { + let array = uint32_array(vec![None; 3]); + + assert!(instantiate_branchless_filter(&array)?.is_some()); + + Ok(()) + } + + #[test] + fn primitive_hash_filter_handles_float_keys() -> Result<()> { + let nan32 = f32::NAN; + let other_nan32 = f32::from_bits(nan32.to_bits() + 1); + let haystack: ArrayRef = Arc::new(Float32Array::from(vec![0.0, nan32])); + let filter = + PrimitiveHashSetFilter::::try_new(&haystack)?; + let needles = Float32Array::from(vec![ + Some(0.0), + Some(-0.0), + Some(nan32), + Some(other_nan32), + None, + ]); + assert_contains( + &filter, + &needles, + vec![Some(true), Some(false), Some(true), Some(false), None], + )?; + + let nan64 = f64::NAN; + let haystack: ArrayRef = Arc::new(Float64Array::from(vec![1.0, nan64])); + let filter = + PrimitiveHashSetFilter::::try_new(&haystack)?; + let needles = Float64Array::from(vec![Some(1.0), Some(nan64), Some(2.0)]); + assert_contains(&filter, &needles, vec![Some(true), Some(true), Some(false)]) + } + #[test] fn bitmap_filter_u8_handles_nulls() -> Result<()> { let haystack: ArrayRef = Arc::new(UInt8Array::from(vec![Some(1), None, Some(3)])); diff --git a/datafusion/physical-expr/src/expressions/in_list/result.rs b/datafusion/physical-expr/src/expressions/in_list/result.rs index 3ebdbfe19f743..1048963b63500 100644 --- a/datafusion/physical-expr/src/expressions/in_list/result.rs +++ b/datafusion/physical-expr/src/expressions/in_list/result.rs @@ -24,15 +24,21 @@ use arrow::array::BooleanArray; use arrow::buffer::{BooleanBuffer, NullBuffer}; -// Truth table for (needle_nulls, haystack_has_nulls, negated): -// (Some, true, false) => values: valid & contains, nulls: valid & contains -// (None, true, false) => values: contains, nulls: contains -// (Some, true, true) => values: valid & !contains, nulls: valid & contains -// (None, true, true) => values: !contains, nulls: contains -// (Some, false, false) => values: valid & contains, nulls: valid -// (Some, false, true) => values: valid & !contains, nulls: valid -// (None, false, false) => values: contains, nulls: none -// (None, false, true) => values: !contains, nulls: none +// Truth table for `value [NOT] IN (set)` with SQL three-valued logic: +// ("-" means the value does not affect the result) +// +// | needle null | set has null | negated | found in set | result | +// |-------------|--------------|---------|--------------|--------| +// | true | - | false | - | null | +// | true | - | true | - | null | +// | false | true | false | true | true | +// | false | true | false | false | null | +// | false | true | true | true | false | +// | false | true | true | false | null | +// | false | false | false | true | true | +// | false | false | false | false | false | +// | false | false | true | true | false | +// | false | false | true | false | true | /// Builds a BooleanArray result for IN list operations. /// @@ -44,7 +50,7 @@ use arrow::buffer::{BooleanBuffer, NullBuffer}; /// This version computes contains for all positions, including nulls, then applies /// null masking via bitmap operations. #[inline] -pub(crate) fn build_in_list_result( +pub(super) fn build_in_list_result( len: usize, needle_nulls: Option<&NullBuffer>, haystack_has_nulls: bool, @@ -63,7 +69,7 @@ where /// This version does not assume contains_buf is pre-masked at null positions. /// It handles nulls using bitmap operations. #[inline] -pub(crate) fn build_result_from_contains( +pub(super) fn build_result_from_contains( needle_nulls: Option<&NullBuffer>, haystack_has_nulls: bool, negated: bool, diff --git a/datafusion/physical-expr/src/expressions/in_list/static_filter.rs b/datafusion/physical-expr/src/expressions/in_list/static_filter.rs index 3c964d4183474..e74ea9131e63b 100644 --- a/datafusion/physical-expr/src/expressions/in_list/static_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/static_filter.rs @@ -15,9 +15,13 @@ // specific language governing permissions and limitations // under the License. +use std::sync::Arc; + use arrow::array::{Array, BooleanArray}; use datafusion_common::Result; +pub(super) type StaticFilterRef = Arc; + /// Trait for InList static filters. /// /// Static filters store a pre-computed set of values (the haystack) and check diff --git a/datafusion/physical-expr/src/expressions/in_list/strategy.rs b/datafusion/physical-expr/src/expressions/in_list/strategy.rs index d5ca8154a92f6..98bd698507844 100644 --- a/datafusion/physical-expr/src/expressions/in_list/strategy.rs +++ b/datafusion/physical-expr/src/expressions/in_list/strategy.rs @@ -19,179 +19,34 @@ use std::sync::Arc; use arrow::array::ArrayRef; use arrow::compute::cast; -use arrow::datatypes::{ - DataType, Date32Type, Date64Type, Decimal128Type, DurationMicrosecondType, - DurationMillisecondType, DurationNanosecondType, DurationSecondType, Float16Type, - Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, - IntervalMonthDayNanoType, IntervalUnit, Time32MillisecondType, Time32SecondType, - Time64MicrosecondType, Time64NanosecondType, TimeUnit, TimestampMicrosecondType, - TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, UInt8Type, - UInt16Type, UInt32Type, UInt64Type, -}; +use arrow::datatypes::DataType; use datafusion_common::Result; use super::array_static_filter::ArrayStaticFilter; -use super::branchless_filter::{ - BranchlessFilter, BranchlessFilterType, BranchlessNative, -}; -use super::primitive_filter::*; -use super::static_filter::StaticFilter; - -type StaticFilterRef = Arc; +use super::fixed_size_binary_filter::instantiate_fixed_size_binary_filter; +use super::primitive_filter::instantiate_primitive_filter; +use super::static_filter::StaticFilterRef; pub(super) fn instantiate_static_filter(in_array: ArrayRef) -> Result { let in_array = flatten_dictionary_haystack(in_array)?; - if let Some(filter) = instantiate_branchless_filter(&in_array)? { + if let Some(filter) = instantiate_fixed_size_binary_filter(&in_array)? { return Ok(filter); } - instantiate_standard_filter(in_array) + if let Some(filter) = instantiate_primitive_filter(&in_array)? { + return Ok(filter); + } + + Ok(Arc::new(ArrayStaticFilter::try_new(in_array)?)) } fn flatten_dictionary_haystack(in_array: ArrayRef) -> Result { // Flatten dictionary-encoded haystacks to their value type so that - // specialized filters (e.g. Int32StaticFilter) are used instead of - // falling through to the generic ArrayStaticFilter. + // specialized primitive filters are used instead of falling through to the + // generic ArrayStaticFilter. match in_array.data_type() { DataType::Dictionary(_, value_type) => Ok(cast(&in_array, value_type.as_ref())?), _ => Ok(in_array), } } - -fn instantiate_branchless_filter(in_array: &ArrayRef) -> Result> { - let non_null_count = in_array.len() - in_array.null_count(); - - macro_rules! filter { - ($arrow_type:ty) => { - branchless_filter::<$arrow_type>(in_array, non_null_count) - }; - } - - match in_array.data_type() { - DataType::Int8 => filter!(Int8Type), - DataType::UInt8 => filter!(UInt8Type), - DataType::Int16 => filter!(Int16Type), - DataType::UInt16 => filter!(UInt16Type), - DataType::Float16 => filter!(Float16Type), - DataType::Int32 => filter!(Int32Type), - DataType::UInt32 => filter!(UInt32Type), - DataType::Float32 => filter!(Float32Type), - DataType::Date32 => filter!(Date32Type), - DataType::Time32(unit) => match unit { - TimeUnit::Second => filter!(Time32SecondType), - TimeUnit::Millisecond => filter!(Time32MillisecondType), - _ => Ok(None), - }, - DataType::Int64 => filter!(Int64Type), - DataType::UInt64 => filter!(UInt64Type), - DataType::Float64 => filter!(Float64Type), - DataType::Date64 => filter!(Date64Type), - DataType::Time64(unit) => match unit { - TimeUnit::Microsecond => filter!(Time64MicrosecondType), - TimeUnit::Nanosecond => filter!(Time64NanosecondType), - _ => Ok(None), - }, - DataType::Timestamp(unit, _) => match unit { - TimeUnit::Second => filter!(TimestampSecondType), - TimeUnit::Millisecond => filter!(TimestampMillisecondType), - TimeUnit::Microsecond => filter!(TimestampMicrosecondType), - TimeUnit::Nanosecond => filter!(TimestampNanosecondType), - }, - DataType::Duration(unit) => match unit { - TimeUnit::Second => filter!(DurationSecondType), - TimeUnit::Millisecond => filter!(DurationMillisecondType), - TimeUnit::Microsecond => filter!(DurationMicrosecondType), - TimeUnit::Nanosecond => filter!(DurationNanosecondType), - }, - DataType::Decimal128(_, _) => filter!(Decimal128Type), - DataType::Interval(IntervalUnit::MonthDayNano) => { - filter!(IntervalMonthDayNanoType) - } - _ => Ok(None), - } -} - -fn instantiate_standard_filter(in_array: ArrayRef) -> Result { - match in_array.data_type() { - DataType::Int8 => bitmap_filter::(&in_array), - DataType::UInt8 => bitmap_filter::(&in_array), - DataType::Int16 => bitmap_filter::(&in_array), - DataType::UInt16 => bitmap_filter::(&in_array), - DataType::Float16 => bitmap_filter::(&in_array), - DataType::Int32 => Ok(Arc::new(Int32StaticFilter::try_new(&in_array)?)), - DataType::Int64 => Ok(Arc::new(Int64StaticFilter::try_new(&in_array)?)), - DataType::UInt32 => Ok(Arc::new(UInt32StaticFilter::try_new(&in_array)?)), - DataType::UInt64 => Ok(Arc::new(UInt64StaticFilter::try_new(&in_array)?)), - // Float primitive types (use ordered wrappers for Hash/Eq) - DataType::Float32 => Ok(Arc::new(Float32StaticFilter::try_new(&in_array)?)), - DataType::Float64 => Ok(Arc::new(Float64StaticFilter::try_new(&in_array)?)), - _ => { - // Fall through to generic implementation for unsupported types - // (Struct, etc.). - Ok(Arc::new(ArrayStaticFilter::try_new(in_array)?)) - } - } -} - -fn bitmap_filter(in_array: &ArrayRef) -> Result -where - T: BitmapFilterType, -{ - Ok(Arc::new(BitmapFilter::::try_new(in_array)?)) -} - -fn branchless_filter( - in_array: &ArrayRef, - non_null_count: usize, -) -> Result> -where - T: BranchlessFilterType, - BranchlessNative: Copy + PartialEq + Send + Sync, -{ - // Larger lists use the standard filter. `try_new` checks the limit again. - if non_null_count > T::MAX_LIST_LEN { - return Ok(None); - } - - Ok(Some(Arc::new(BranchlessFilter::::try_new(in_array)?))) -} - -#[cfg(test)] -mod tests { - use arrow::array::UInt32Array; - use arrow::datatypes::UInt32Type; - - use super::super::branchless_filter::BranchlessFilterType; - use super::*; - - fn uint32_array(values: Vec>) -> ArrayRef { - Arc::new(UInt32Array::from(values)) - } - - #[test] - fn branchless_routing_respects_max_list_len() -> Result<()> { - let max_len = ::MAX_LIST_LEN; - - let values = (0..max_len) - .map(|value| Some(value as u32)) - .collect::>(); - assert!(instantiate_branchless_filter(&uint32_array(values))?.is_some()); - - let values = (0..=max_len) - .map(|value| Some(value as u32)) - .collect::>(); - assert!(instantiate_branchless_filter(&uint32_array(values))?.is_none()); - - Ok(()) - } - - #[test] - fn branchless_routing_handles_zero_non_null_values() -> Result<()> { - let array = uint32_array(vec![None; 3]); - - assert!(instantiate_branchless_filter(&array)?.is_some()); - - Ok(()) - } -}