diff --git a/datafusion/physical-expr/benches/in_list_strategy.rs b/datafusion/physical-expr/benches/in_list_strategy.rs index c69af192b9cdd..791f1a684ce90 100644 --- a/datafusion/physical-expr/benches/in_list_strategy.rs +++ b/datafusion/physical-expr/benches/in_list_strategy.rs @@ -35,8 +35,8 @@ //! |------|-------|-----------------|-------------------| //! | Narrow integer cases | UInt8 | small value domain | 4, 16 | //! | Narrow integer cases | Int16, Float16 | larger value domain | 4, 64, 256 | -//! | 32-bit primitive cases | Int32, Float32 | small and large lists | 4, 32, 64, 256 | -//! | 64-bit primitive cases | Int64, TimestampNs | small and large lists | 4, 16, 32, 128 | +//! | 32-bit primitive cases | Int32, Float32 | small and large lists | 4, 32, 33, 64, 256, 1024, 10000 | +//! | 64-bit primitive cases | Int64, TimestampNs | small and large lists | 4, 16, 17, 32, 128, 1024, 10000 | //! | 128-bit interval cases | IntervalMonthDayNano | small lists | 4 | //! | Utf8 short-string cases | Utf8 | 8-byte strings | 4, 64, 256 | //! | Utf8 long-string cases | Utf8 | 24-byte strings | 4, 64, 256 | @@ -44,7 +44,7 @@ //! | 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/2/4/8/16) | fixed-width binary values | representative thresholds and large lists | use arrow::array::types::IntervalMonthDayNano; use arrow::array::*; @@ -166,11 +166,11 @@ fn random_string(rng: &mut StdRng, len: usize) -> String { fn strings_with_shared_prefix( rng: &mut StdRng, count: usize, - prefix_len: usize, + prefix: &str, + discriminator: char, ) -> Vec { - let prefix = random_string(rng, prefix_len); (0..count) - .map(|_| format!("{}{}", prefix, random_string(rng, 8))) // prefix + random 8-char suffix + .map(|_| format!("{prefix}{}{discriminator}", random_string(rng, 7))) .collect() } @@ -275,12 +275,14 @@ fn bench_string_shared_prefix( .wrapping_add(prefix_len as u64 * 0x4444); let mut rng = StdRng::seed_from_u64(seed); - // Generate IN list with a shared prefix. - let haystack = strings_with_shared_prefix(&mut rng, list_size, prefix_len); + // Use the same prefix and equal-length, disjoint suffixes for both pools so + // misses exercise length/prefix collisions rather than immediate rejection. + let prefix = random_string(&mut rng, prefix_len); + let haystack = strings_with_shared_prefix(&mut rng, list_size, &prefix, 'h'); // Generate non-matching strings with the same prefix to keep misses close // to the matching set. - let non_match_pool = strings_with_shared_prefix(&mut rng, 100, prefix_len); + let non_match_pool = strings_with_shared_prefix(&mut rng, 100, &prefix, 'm'); // Generate array with controlled match rate let values: A = (0..ARRAY_SIZE) @@ -313,6 +315,7 @@ fn bench_string_mixed_lengths( name: &str, list_size: usize, match_rate: f64, + inline_rate: f64, to_scalar: fn(String) -> ScalarValue, ) where A: Array + FromIterator> + 'static, @@ -320,12 +323,21 @@ fn bench_string_mixed_lengths( let seed = 0xABCD_EF01_u64.wrapping_add(list_size as u64 * 0x5555); let mut rng = StdRng::seed_from_u64(seed); - // Mixed lengths: some short (<= 12), some long (> 12) - let lengths = [4, 8, 12, 16, 20, 24]; + let inline_lengths = [4, 8, 12]; + let long_lengths = [16, 20, 24]; + let inline_count = ((list_size as f64 * inline_rate).round() as usize) + .max(1) + .min(list_size - 1); // Generate IN list with mixed lengths let haystack: Vec = (0..list_size) - .map(|_| { + .map(|idx| { + let inline = idx < inline_count; + let lengths = if inline { + &inline_lengths + } else { + &long_lengths + }; let len = *lengths.choose(&mut rng).unwrap(); random_string(&mut rng, len) }) @@ -337,6 +349,11 @@ fn bench_string_mixed_lengths( Some(if !haystack.is_empty() && rng.random_bool(match_rate) { haystack.choose(&mut rng).unwrap().clone() } else { + let lengths = if rng.random_bool(inline_rate) { + &inline_lengths + } else { + &long_lengths + }; let len = *lengths.choose(&mut rng).unwrap(); random_string(&mut rng, len) }) @@ -420,7 +437,7 @@ fn bench_narrow_integer(c: &mut Criterion) { fn bench_primitive(c: &mut Criterion) { // Int32: small and larger list sizes - for list_size in [4, 32, 64, 256] { + for list_size in [4, 32, 33, 64, 256, 1024, 10_000] { let list_case = if list_size <= 32 { "small_list" } else { @@ -442,7 +459,7 @@ fn bench_primitive(c: &mut Criterion) { } // Int64: small and larger list sizes - for list_size in [4, 16, 32, 128] { + for list_size in [4, 16, 17, 32, 128, 1024, 10_000] { let list_case = if list_size <= 16 { "small_list" } else { @@ -602,6 +619,7 @@ fn bench_utf8(c: &mut Criterion) { &format!("mixed_len/list={list_size}/match={match_pct}%"), list_size, match_pct as f64 / 100.0, + 0.5, to_scalar, ); } @@ -694,6 +712,23 @@ fn bench_utf8view(c: &mut Criterion) { &format!("mixed_len/list={list_size}/match={match_pct}%"), list_size, match_pct as f64 / 100.0, + 0.5, + to_scalar, + ); + } + } + + // Strongly skewed mixed lists exercise routing near the all-inline and + // all-long boundaries while retaining both representations. + for inline_pct in [2, 98] { + for match_pct in MATCH_RATES { + bench_string_mixed_lengths::( + c, + "utf8view", + &format!("mixed_len/inline={inline_pct}%/list=64/match={match_pct}%"), + 64, + match_pct as f64 / 100.0, + inline_pct as f64 / 100.0, to_scalar, ); } @@ -996,40 +1031,43 @@ 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]; - rng.fill(&mut buf[..]); - buf +/// Generates deterministic, disjoint hit and miss values. The high bit of the +/// final byte distinguishes misses, including for the one-byte domain. +fn fixed_binary_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 } -/// 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, ) { - let seed = 0xF1ED_B1A7_u64.wrapping_add(list_size as u64 * 0x6666); + 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)) + .map(|index| fixed_binary_value(width, index, false)) .collect(); - // Generate array with controlled match rate let values: Vec> = (0..ARRAY_SIZE) - .map(|_| { + .map(|index| { if !haystack.is_empty() && rng.random_bool(match_rate) { haystack.choose(&mut rng).unwrap().clone() } else { - random_fixed_binary_16(&mut rng) + fixed_binary_value(width, index, true) } }) .collect(); @@ -1040,28 +1078,40 @@ 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_rate * 100.0) as u32 + ), + ), &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), + (4, 4), + (4, 64), + (8, 4), + (8, 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 as f64 / 100.0); } } } diff --git a/datafusion/physical-expr/src/expressions/in_list.rs b/datafusion/physical-expr/src/expressions/in_list.rs index 874e149b58328..5de58fae6a699 100644 --- a/datafusion/physical-expr/src/expressions/in_list.rs +++ b/datafusion/physical-expr/src/expressions/in_list.rs @@ -38,6 +38,9 @@ use datafusion_expr::{ColumnarValue, expr_vec_fmt}; mod array_static_filter; mod branchless_filter; +mod byte_view_filter; +mod fixed_size_binary_filter; +mod frozen_set; mod primitive_filter; mod result; mod static_filter; @@ -215,15 +218,15 @@ impl InListExpr { expr, list, negated, - Some(instantiate_static_filter(array)?), + Some(instantiate_static_filter(array, &expr_data_type)?), )) } /// Create a new InList expression, using a static filter when possible. /// /// This validates data types and attempts to create a static filter for constant - /// list expressions. Uses specialized StaticFilter implementations for better - /// performance (e.g., Int32StaticFilter for Int32). + /// list expressions. Uses specialized `StaticFilter` implementations for better + /// performance (for example, primitive filters for fixed-width values). /// /// Returns an error if data types don't match. If the list contains non-constant /// expressions, falls back to dynamic evaluation at runtime. @@ -242,7 +245,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 }; @@ -2592,7 +2595,7 @@ mod tests { // Create IN list with Int32 literals: (100, 200, 300) let list = vec![lit(100i32), lit(200i32), lit(300i32)]; - // Create InListExpr via in_list() - this uses Int32StaticFilter for Int32 lists + // Create InListExpr via in_list() - this uses a specialized primitive filter let expr = in_list(col_a, list, &false, &schema)?; // Create dictionary-encoded batch with values [100, 200, 500] @@ -3576,6 +3579,49 @@ mod tests { )? ); + // Utf8View in_array, Utf8View and Dict(Utf8View) needles + let utf8view_in = + Arc::new(StringViewArray::from(vec!["a", "b", "c"])) as ArrayRef; + let utf8view_needle = + Arc::new(StringViewArray::from(vec!["a", "d", "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)? + ); + + // 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(fsb_needle), fsb_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..cb28251d20f49 100644 --- a/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs @@ -58,8 +58,9 @@ //! would be useful. Wider types have too many possible values for such a //! bitmap, so their limits are tuned separately. //! -//! Larger lists use the standard filter strategy, including bitmap filters for -//! one- and two-byte types. +//! Larger lists use another filter strategy: bitmap filters for one- and +//! two-byte types, frozen sets for supported four- and eight-byte types, and +//! the standard fallback for the remaining types. //! //! # What about nulls? //! @@ -71,7 +72,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 +245,17 @@ where check_values, }) } + + #[inline] + pub(super) fn contains_slice( + &self, + input_values: &[BranchlessNative], + nulls: Option<&NullBuffer>, + negated: bool, + ) -> BooleanArray { + 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 +284,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_slice(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..be0d4e0872079 --- /dev/null +++ b/datafusion/physical-expr/src/expressions/in_list/byte_view_filter.rs @@ -0,0 +1,678 @@ +// 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 Utf8View and BinaryView IN lists. + +use std::hash::BuildHasher; +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::hash_utils::RandomState; +use datafusion_common::{Result, exec_datafusion_err}; +use hashbrown::{DefaultHashBuilder, HashTable}; + +use super::branchless_filter::{BranchlessFilter, BranchlessFilterType}; +use super::frozen_set::{FrozenSet, FrozenSetHash}; +use super::result::{build_in_list_result, build_in_list_result_with_null_shortcircuit}; +use super::static_filter::{StaticFilter, handle_dictionary}; + +#[inline(always)] +fn view_len(view: u128) -> u32 { + view as u32 +} + +/// The low 64 bits of a long view contain its length and four-byte prefix. +#[inline(always)] +fn long_key(view: u128) -> u64 { + view as u64 +} + +#[inline(always)] +fn view_key(view: u128) -> u128 { + if view_len(view) <= MAX_INLINE_VIEW_LEN { + view + } else { + long_key(view) as u128 + } +} + +fn downcast_byte_view( + array: &dyn Array, +) -> Result<&GenericByteViewArray> { + array + .as_byte_view_opt::() + .ok_or_else(|| exec_datafusion_err!("Expected concrete {} array", T::DATA_TYPE)) +} + +/// Keyed folded-multiply hash specialized for Arrow's inline view encoding. +struct InlineViewHash { + key: [u64; 2], +} + +impl Default for InlineViewHash { + fn default() -> Self { + let state = DefaultHashBuilder::default(); + Self { + key: [ + BuildHasher::hash_one(&state, 0_u8), + BuildHasher::hash_one(&state, 1_u8), + ], + } + } +} + +impl FrozenSetHash for InlineViewHash { + #[inline(always)] + fn hash_one(&self, value: u128) -> u64 { + let lo = value as u64 ^ self.key[0]; + let hi = (value >> 64) as u64 ^ self.key[1]; + let product = lo as u128 * hi as u128; + product as u64 ^ (product >> 64) as u64 + } +} + +#[derive(Debug, PartialEq, Eq)] +enum ViewComposition { + Inline, + Long, + Mixed { + inline_count: usize, + long_count: usize, + }, +} + +fn view_composition(array: &ArrayRef) -> Result { + let array = downcast_byte_view::(array.as_ref())?; + let mut inline_count = 0; + let mut long_count = 0; + let mut visit = |idx: usize| { + if view_len(array.views()[idx]) <= MAX_INLINE_VIEW_LEN { + inline_count += 1; + } else { + long_count += 1; + } + }; + + match array.nulls() { + Some(nulls) => { + BitIndexIterator::new(nulls.validity(), nulls.offset(), nulls.len()) + .for_each(&mut visit); + } + None => (0..array.len()).for_each(&mut visit), + } + + Ok(match (inline_count, long_count) { + (inline_count, long_count) if inline_count > 0 && long_count > 0 => { + ViewComposition::Mixed { + inline_count, + long_count, + } + } + (0, long_count) if long_count > 0 => ViewComposition::Long, + _ => ViewComposition::Inline, + }) +} + +fn reinterpret_byte_view_as_decimal128( + array: &dyn Array, +) -> Result { + let array = downcast_byte_view::(array)?; + let views = array.views(); + let values = ScalarBuffer::::new(views.inner().clone(), 0, views.len()); + Ok(Arc::new(PrimitiveArray::::new( + values, + array.nulls().cloned(), + ))) +} + +fn make_byte_view_branchless_filter( + in_array: &ArrayRef, +) -> Result> { + let values = reinterpret_byte_view_as_decimal128::(in_array.as_ref())?; + + Ok(Arc::new(ByteViewBranchless:: { + inner: BranchlessFilter::::try_new(&values)?, + _marker: PhantomData, + })) +} + +struct ByteViewBranchless { + inner: BranchlessFilter, + _marker: PhantomData, +} + +struct InlineByteViewFilter { + set: FrozenSet, + null_count: usize, + _marker: PhantomData, +} + +impl InlineByteViewFilter { + fn try_new(in_array: &ArrayRef) -> Result { + if in_array.data_type() != &T::DATA_TYPE { + return Err(exec_datafusion_err!( + "InlineByteViewFilter: expected {} array, got {}", + T::DATA_TYPE, + in_array.data_type() + )); + } + + let array = downcast_byte_view::(in_array.as_ref())?; + let mut values = Vec::with_capacity(array.len() - array.null_count()); + match array.nulls() { + Some(nulls) => { + BitIndexIterator::new(nulls.validity(), nulls.offset(), nulls.len()) + .for_each(|idx| values.push(array.views()[idx])); + } + None => values.extend_from_slice(array.views()), + } + + Ok(Self { + set: FrozenSet::try_new_with_hasher(&values, InlineViewHash::default())?, + null_count: array.null_count(), + _marker: PhantomData, + }) + } +} + +impl StaticFilter for InlineByteViewFilter { + fn null_count(&self) -> usize { + self.null_count + } + + fn contains(&self, v: &dyn Array, negated: bool) -> Result { + handle_dictionary!(self, v, negated); + + if v.data_type() != &T::DATA_TYPE { + return Err(exec_datafusion_err!( + "InlineByteViewFilter: expected {} array, got {}", + T::DATA_TYPE, + v.data_type() + )); + } + + 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); + + if v.data_type() != &T::DATA_TYPE { + return Err(exec_datafusion_err!( + "ByteViewBranchless: expected {} array, got {}", + T::DATA_TYPE, + v.data_type() + )); + } + + let array = downcast_byte_view::(v)?; + let values: &[i128] = array.views().inner().typed_data(); + Ok(self.inner.contains_slice(values, array.nulls(), negated)) + } +} + +/// Two-stage filter for inline and mixed Utf8View and BinaryView arrays. +/// +/// Inline values are represented completely by their 128-bit view. For long +/// values, the length and prefix reject impossible matches before an exact +/// full-byte lookup. +struct ByteViewFilter { + in_array: ArrayRef, + view_keys: FrozenSet, + long_value_table: HashTable, + state: RandomState, + _marker: PhantomData, +} + +impl ByteViewFilter { + fn try_new(in_array: ArrayRef) -> Result { + if in_array.data_type() != &T::DATA_TYPE { + return Err(exec_datafusion_err!( + "ByteViewFilter: expected {} array, got {}", + T::DATA_TYPE, + in_array.data_type() + )); + } + + let array = downcast_byte_view::(in_array.as_ref())?; + let mut view_keys = Vec::new(); + let mut long_value_table = HashTable::new(); + let state = RandomState::default(); + + let mut insert = |idx: usize| { + let view = array.views()[idx]; + view_keys.push(view_key(view)); + if view_len(view) <= MAX_INLINE_VIEW_LEN { + return; + } + + // SAFETY: idx is produced from this array's bounds or validity bitmap. + let value: &[u8] = unsafe { array.value_unchecked(idx) }.as_ref(); + let hash = state.hash_one(value); + if long_value_table + .find(hash, |&stored_idx| { + // SAFETY: stored_idx was inserted from this array. + let stored: &[u8] = + unsafe { array.value_unchecked(stored_idx) }.as_ref(); + stored == value + }) + .is_none() + { + long_value_table.insert_unique(hash, idx, |&stored_idx| { + // SAFETY: stored_idx was inserted from this array. + let stored: &[u8] = + unsafe { array.value_unchecked(stored_idx) }.as_ref(); + state.hash_one(stored) + }); + } + }; + + match array.nulls() { + Some(nulls) => { + BitIndexIterator::new(nulls.validity(), nulls.offset(), nulls.len()) + .for_each(&mut insert); + } + None => (0..array.len()).for_each(&mut insert), + } + + view_keys.sort_unstable(); + view_keys.dedup(); + + Ok(Self { + in_array, + view_keys: FrozenSet::try_new_with_hasher( + &view_keys, + InlineViewHash::default(), + )?, + long_value_table, + state, + _marker: PhantomData, + }) + } +} + +impl StaticFilter for ByteViewFilter { + fn null_count(&self) -> usize { + self.in_array.null_count() + } + + fn contains(&self, v: &dyn Array, negated: bool) -> Result { + handle_dictionary!(self, v, negated); + + if v.data_type() != &T::DATA_TYPE { + return Err(exec_datafusion_err!( + "ByteViewFilter: expected {} array, got {}", + T::DATA_TYPE, + v.data_type() + )); + } + + let needles = downcast_byte_view::(v)?; + let haystack = downcast_byte_view::(self.in_array.as_ref())?; + Ok(build_in_list_result_with_null_shortcircuit( + needles.len(), + needles.nulls(), + self.in_array.null_count() > 0, + negated, + |i| { + let view = needles.views()[i]; + if !self.view_keys.contains(view_key(view)) { + return false; + } + if view_len(view) <= MAX_INLINE_VIEW_LEN { + return true; + } + + // SAFETY: i is in bounds and null indices are skipped. + let needle: &[u8] = unsafe { needles.value_unchecked(i) }.as_ref(); + let hash = self.state.hash_one(needle); + self.long_value_table + .find(hash, |&idx| { + // SAFETY: idx was inserted from self.in_array. + let value: &[u8] = + unsafe { haystack.value_unchecked(idx) }.as_ref(); + value == needle + }) + .is_some() + }, + )) + } +} + +fn make_byte_view_filter( + in_array: &ArrayRef, +) -> Result> { + Ok(Arc::new(ByteViewFilter::::try_new(Arc::clone( + in_array, + ))?)) +} + +fn make_inline_byte_view_filter( + in_array: &ArrayRef, +) -> Result> { + Ok(Arc::new(InlineByteViewFilter::::try_new(in_array)?)) +} + +fn instantiate_typed_byte_view_filter( + in_array: &ArrayRef, +) -> Result>> { + let non_null_count = in_array.len() - in_array.null_count(); + match view_composition::(in_array)? { + ViewComposition::Inline + if non_null_count + <= ::MAX_LIST_LEN => + { + make_byte_view_branchless_filter::(in_array).map(Some) + } + ViewComposition::Inline => make_inline_byte_view_filter::(in_array).map(Some), + ViewComposition::Mixed { + inline_count, + long_count, + } if inline_count <= long_count => make_byte_view_filter::(in_array).map(Some), + ViewComposition::Mixed { .. } => Ok(None), + ViewComposition::Long => Ok(None), + } +} + +/// Creates an optimized byte-view filter when its measured specialization wins. +/// +/// Inline lists use direct view comparisons or a frozen set. Long-dominant mixed +/// lists add length/prefix rejection and exact long-value verification. Other +/// compositions stay on the generic filter when this extra stage does not win. +pub(super) fn instantiate_byte_view_filter( + in_array: &ArrayRef, +) -> Result>> { + match in_array.data_type() { + DataType::Utf8View => { + instantiate_typed_byte_view_filter::(in_array) + } + DataType::BinaryView => { + instantiate_typed_byte_view_filter::(in_array) + } + _ => Ok(None), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{BinaryViewArray, DictionaryArray, Int8Array, StringViewArray}; + use arrow::datatypes::{BinaryViewType, 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 frozen_filter_handles_inline_slices() -> Result<()> { + let haystack: ArrayRef = Arc::new( + StringViewArray::from(vec![ + Some("outside"), + Some("a"), + Some("b"), + Some("c"), + Some("d"), + Some("e"), + Some("tail"), + ]) + .slice(1, 5), + ); + let filter = ByteViewFilter::::try_new(haystack)?; + let needles = + StringViewArray::from(vec![Some("outside"), Some("b"), Some("z"), Some("e")]) + .slice(1, 3); + + assert_contains(&filter, &needles, vec![Some(true), Some(false), Some(true)]) + } + + #[test] + fn inline_filter_handles_slices_nulls_and_not_in() -> 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 filter = InlineByteViewFilter::::try_new(&haystack)?; + 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 frozen_filter_verifies_long_prefix_collisions_and_nulls() -> Result<()> { + let haystack: ArrayRef = Arc::new( + StringViewArray::from(vec![ + Some("outside-long-value"), + Some("abcdefghijklmn1"), + None, + Some("short"), + Some("tail-long-value"), + ]) + .slice(1, 3), + ); + let filter = ByteViewFilter::::try_new(haystack)?; + let needles = StringViewArray::from(vec![ + Some("abcdefghijklmn1"), + Some("abcdefghijklmn2"), + Some("short"), + None, + ]); + + assert_contains(&filter, &needles, vec![Some(true), None, Some(true), None])?; + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![Some(false), None, Some(false), None]) + ); + Ok(()) + } + + #[test] + fn frozen_filter_handles_binary_views() -> Result<()> { + let haystack: ArrayRef = Arc::new(BinaryViewArray::from(vec![ + Some([0xff, 0x00].as_slice()), + Some(b"abcdefghijklmn1".as_slice()), + ])); + let filter = ByteViewFilter::::try_new(haystack)?; + let needles = BinaryViewArray::from(vec![ + Some([0xff, 0x00].as_slice()), + Some(b"abcdefghijklmn1".as_slice()), + Some(b"abcdefghijklmn2".as_slice()), + ]); + + assert_contains(&filter, &needles, vec![Some(true), Some(true), Some(false)]) + } + + #[test] + fn frozen_filter_handles_dictionary_needles() -> Result<()> { + let haystack: ArrayRef = Arc::new(StringViewArray::from(vec![ + Some("short"), + Some("abcdefghijklmn1"), + None, + ])); + let filter = ByteViewFilter::::try_new(haystack)?; + let keys = Int8Array::from(vec![Some(0), Some(1), None, Some(2)]); + let values = Arc::new(StringViewArray::from(vec![ + Some("short"), + Some("abcdefghijklmn1"), + Some("missing"), + ])); + let needles = DictionaryArray::try_new(keys, values)?; + + assert_contains(&filter, &needles, vec![Some(true), Some(true), None, None]) + } + + #[test] + fn branchless_filter_handles_utf8_and_binary_views() -> Result<()> { + let utf8_haystack: ArrayRef = + Arc::new(StringViewArray::from(vec![Some("one"), None, Some("two")])); + let utf8_filter = + make_byte_view_branchless_filter::(&utf8_haystack)?; + let utf8_needles = StringViewArray::from(vec![Some("two"), Some("three"), None]); + assert_contains(&*utf8_filter, &utf8_needles, vec![Some(true), None, None])?; + + let binary_haystack: ArrayRef = Arc::new(BinaryViewArray::from(vec![ + Some([0xff].as_slice()), + Some([0x00].as_slice()), + ])); + let binary_filter = + make_byte_view_branchless_filter::(&binary_haystack)?; + let binary_needles = + BinaryViewArray::from(vec![Some([0x00].as_slice()), Some([0x01].as_slice())]); + assert_contains( + &*binary_filter, + &binary_needles, + vec![Some(true), Some(false)], + ) + } + + #[test] + fn byte_view_filters_reject_other_types() -> Result<()> { + let haystack: ArrayRef = Arc::new(StringViewArray::from(vec![ + Some("short"), + Some("abcdefghijklmn1"), + ])); + let frozen = ByteViewFilter::::try_new(Arc::clone(&haystack))?; + let branchless_haystack: ArrayRef = + Arc::new(StringViewArray::from(vec![Some("short")])); + let branchless = + make_byte_view_branchless_filter::(&branchless_haystack)?; + let needles = BinaryViewArray::from(vec![Some(b"short".as_slice())]); + + assert!(frozen.contains(&needles, false).is_err()); + assert!(branchless.contains(&needles, false).is_err()); + Ok(()) + } + + #[test] + fn byte_view_routing_only_selects_measured_wins() -> Result<()> { + let inline_four: ArrayRef = Arc::new(StringViewArray::from(vec![ + Some("a"), + None, + Some("b"), + Some("c"), + Some("d"), + ])); + assert!(instantiate_byte_view_filter(&inline_four)?.is_some()); + + let inline_five: ArrayRef = Arc::new(StringViewArray::from(vec![ + Some("a"), + Some("b"), + Some("c"), + Some("d"), + Some("e"), + ])); + assert!(instantiate_byte_view_filter(&inline_five)?.is_some()); + + let all_null: ArrayRef = + Arc::new(StringViewArray::from(vec![None::<&str>, None::<&str>])); + assert!(instantiate_byte_view_filter(&all_null)?.is_some()); + + let all_long: ArrayRef = Arc::new(StringViewArray::from(vec![ + Some("abcdefghijklmn1"), + Some("abcdefghijklmn2"), + ])); + assert!(instantiate_byte_view_filter(&all_long)?.is_none()); + + let mixed: ArrayRef = Arc::new(StringViewArray::from(vec![ + Some("short"), + Some("abcdefghijklmn1"), + ])); + assert!(instantiate_byte_view_filter(&mixed)?.is_some()); + + let inline_dominant_mixed: ArrayRef = Arc::new(StringViewArray::from(vec![ + Some("a"), + Some("b"), + Some("c"), + Some("d"), + Some("e"), + Some("abcdefghijklmn1"), + ])); + assert!(instantiate_byte_view_filter(&inline_dominant_mixed)?.is_none()); + + let long_dominant_mixed: ArrayRef = Arc::new(StringViewArray::from(vec![ + Some("short"), + Some("abcdefghijklmn1"), + Some("abcdefghijklmn2"), + ])); + assert!(instantiate_byte_view_filter(&long_dominant_mixed)?.is_some()); + + let inline_binary: ArrayRef = Arc::new(BinaryViewArray::from(vec![ + Some(b"a".as_slice()), + Some(b"b".as_slice()), + Some(b"c".as_slice()), + Some(b"d".as_slice()), + Some(b"e".as_slice()), + ])); + assert!(instantiate_byte_view_filter(&inline_binary)?.is_some()); + + let mixed_binary: ArrayRef = Arc::new(BinaryViewArray::from(vec![ + Some([0xff].as_slice()), + Some(b"abcdefghijklmn1".as_slice()), + ])); + assert!(instantiate_byte_view_filter(&mixed_binary)?.is_some()); + + Ok(()) + } +} 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..d39317e0ee6d8 --- /dev/null +++ b/datafusion/physical-expr/src/expressions/in_list/fixed_size_binary_filter.rs @@ -0,0 +1,466 @@ +// 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. + +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::ScalarBuffer; +use arrow::datatypes::{ + ArrowPrimitiveType, DataType, Decimal128Type, UInt8Type, UInt16Type, UInt32Type, + UInt64Type, +}; +use datafusion_common::{Result, exec_datafusion_err}; + +use super::branchless_filter::{ + BranchlessFilter, BranchlessFilterType, BranchlessNative, +}; +use super::primitive_filter::{BitmapFilter, BitmapFilterType, PrimitiveFrozenFilter}; +use super::static_filter::{StaticFilter, handle_dictionary}; + +type StaticFilterRef = Arc; + +/// 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_fixed_size_binary( + array: &FixedSizeBinaryArray, +) -> Result> +where + T: ArrowPrimitiveType, +{ + let width = size_of::(); + if usize::try_from(array.value_length()).ok() != Some(width) { + return Err(exec_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 { + let mut values = Vec::with_capacity(array.len()); + for index in 0..array.len() { + // SAFETY: `FixedSizeBinaryArray` guarantees that `source` contains + // at least `array.len() * width` bytes. Arrow native values are + // trivially transmutable, and `read_unaligned` accepts this pointer. + let value = unsafe { + source + .as_ptr() + .add(index * width) + .cast::() + .read_unaligned() + }; + values.push(value); + } + ScalarBuffer::from(values) + }; + + Ok(PrimitiveArray::::new(values, array.nulls().cloned())) +} + +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: BranchlessFilterType, +{ + 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: ArrayRef = Arc::new(reinterpret_fixed_size_binary::(array)?); + self.inner.contains(primitive.as_ref(), negated) + } +} + +fn downcast_fixed_size_binary(in_array: &ArrayRef) -> Result<&FixedSizeBinaryArray> { + in_array.as_fixed_size_binary_opt().ok_or_else(|| { + exec_datafusion_err!( + "FixedSizeBinary filter: expected concrete fixed-size binary array, got {}", + in_array.data_type() + ) + }) +} + +fn bitmap_filter(in_array: &ArrayRef) -> Result +where + T: BitmapFilterType + BranchlessFilterType, +{ + let array = downcast_fixed_size_binary(in_array)?; + let primitive: ArrayRef = Arc::new(reinterpret_fixed_size_binary::(array)?); + let inner = Arc::new(BitmapFilter::::try_new(&primitive)?); + Ok(Arc::new(FixedSizeBinaryFilter::::new( + in_array.data_type().clone(), + inner, + ))) +} + +fn wide_filter(in_array: &ArrayRef) -> Result +where + T: BranchlessFilterType, + BranchlessNative: Copy + Eq + Hash + Send + Sync, +{ + let array = downcast_fixed_size_binary(in_array)?; + let primitive: ArrayRef = Arc::new(reinterpret_fixed_size_binary::(array)?); + let non_null_count = primitive.len() - primitive.null_count(); + let inner: StaticFilterRef = if non_null_count <= T::MAX_LIST_LEN { + Arc::new(BranchlessFilter::::try_new(&primitive)?) + } else { + Arc::new(PrimitiveFrozenFilter::::try_new(&primitive)?) + }; + Ok(Arc::new(FixedSizeBinaryFilter::::new( + in_array.data_type().clone(), + inner, + ))) +} + +pub(super) fn instantiate_fixed_size_binary_filter( + in_array: &ArrayRef, +) -> Result> { + let DataType::FixedSizeBinary(width) = in_array.data_type() else { + return Ok(None); + }; + if in_array.as_fixed_size_binary_opt().is_none() { + return Ok(None); + } + + let filter = match width { + 1 => bitmap_filter::(in_array)?, + 2 => bitmap_filter::(in_array)?, + 4 => wide_filter::(in_array)?, + 8 => wide_filter::(in_array)?, + 16 => wide_filter::(in_array)?, + _ => return Ok(None), + }; + Ok(Some(filter)) +} + +#[cfg(test)] +mod tests { + use std::any::Any; + + use arrow::array::{ArrayData, DictionaryArray, Int8Array}; + use arrow::buffer::{Buffer, NullBuffer}; + use arrow::datatypes::Int8Type; + + use super::*; + + #[derive(Clone, Debug)] + struct CustomFixedSizeBinaryArray(FixedSizeBinaryArray); + + // SAFETY: Every Array method delegates to the wrapped, valid Arrow array. + unsafe impl Array for CustomFixedSizeBinaryArray { + fn as_any(&self) -> &dyn Any { + self + } + + fn to_data(&self) -> ArrayData { + self.0.to_data() + } + + fn into_data(self) -> ArrayData { + self.0.into_data() + } + + fn data_type(&self) -> &DataType { + self.0.data_type() + } + + fn slice(&self, offset: usize, length: usize) -> ArrayRef { + Arc::new(Self(self.0.slice(offset, length))) + } + + fn len(&self) -> usize { + self.0.len() + } + + fn is_empty(&self) -> bool { + self.0.is_empty() + } + + fn offset(&self) -> usize { + self.0.offset() + } + + fn nulls(&self) -> Option<&NullBuffer> { + self.0.nulls() + } + + fn get_buffer_memory_size(&self) -> usize { + self.0.get_buffer_memory_size() + } + + fn get_array_memory_size(&self) -> usize { + self.0.get_array_memory_size() + } + } + + 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), + (2, 64), + (4, 32), + (4, 33), + (8, 16), + (8, 17), + (16, 4), + (16, 5), + ] { + let haystack = (0..list_len) + .map(|index| Some(value(width, index, false))) + .collect::>(); + let filter = make_filter(width, &haystack)?; + let needles = array( + width, + &[ + Some(value(width, list_len / 2, false)), + Some(value(width, list_len, true)), + 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 = 8; + let parent = array( + width, + &[ + Some(value(width, 1, false)), + None, + Some(value(width, 2, false)), + ], + ); + let in_array: ArrayRef = Arc::new(parent.slice(1, 2)); + let filter = instantiate_fixed_size_binary_filter(&in_array)?.unwrap(); + let needles = array( + width, + &[ + Some(value(width, 2, false)), + Some(value(width, 3, 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_empty_all_null_duplicate_and_sentinel_lists() -> Result<()> { + let empty = make_filter(4, &[])?; + let needle = array(4, &[Some(value(4, 0, false))]); + assert_eq!( + empty.contains(&needle, false)?, + BooleanArray::from(vec![Some(false)]) + ); + + let all_null = make_filter(4, &vec![None; 40])?; + assert_eq!( + all_null.contains(&needle, false)?, + BooleanArray::from(vec![None]) + ); + + // More than 32 values selects FrozenSet. Keep the first real member at + // zero and include duplicates to exercise its sentinel representation. + let mut haystack = vec![Some(value(4, 0, false)); 2]; + haystack.extend((1..32).map(|index| Some(value(4, index, false)))); + let filter = make_filter(4, &haystack)?; + let needles = array(4, &[Some(value(4, 0, false)), Some(value(4, 40, true))]); + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![Some(true), Some(false)]) + ); + 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]) + ); + Ok(()) + } + + #[test] + fn rejects_width_mismatch_and_unsupported_widths() -> Result<()> { + let filter = make_filter(4, &[Some(value(4, 1, false))])?; + let wrong_width = array(8, &[Some(value(8, 1, false))]); + assert!(filter.contains(&wrong_width, false).is_err()); + + let unsupported: ArrayRef = Arc::new(array(3, &[Some(vec![1, 2, 3])])); + assert!(instantiate_fixed_size_binary_filter(&unsupported)?.is_none()); + Ok(()) + } + + #[test] + fn custom_haystacks_fall_back_and_custom_needles_error() -> Result<()> { + let concrete = array(4, &[Some(value(4, 1, false))]); + let custom: ArrayRef = Arc::new(CustomFixedSizeBinaryArray(concrete.clone())); + assert!(instantiate_fixed_size_binary_filter(&custom)?.is_none()); + + let filter = make_filter(4, &[Some(value(4, 1, false))])?; + assert!(filter.contains(custom.as_ref(), false).is_err()); + Ok(()) + } + + #[test] + fn aligned_buffers_are_reinterpreted_zero_copy() -> 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_fixed_size_binary::(&array)?; + + assert_eq!(primitive.values().inner().as_ptr(), source_ptr); + Ok(()) + } + + #[test] + fn unaligned_haystacks_and_needles_are_copied_safely() -> Result<()> { + fn unaligned_array(width: i32, values: &[Vec]) -> 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, None) + } + + 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)); + let filter = instantiate_fixed_size_binary_filter(&haystack)?.unwrap(); + let needles = + unaligned_array(width, &[value(width, 3, false), value(width, 9, true)]); + + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![Some(true), Some(false)]) + ); + Ok(()) + } + + #[test] + fn preserves_nulls_when_copying_unaligned_values() -> Result<()> { + let width = 8; + let mut bytes = vec![0]; + bytes.extend(value(width, 1, false)); + bytes.extend(value(width, 2, false)); + let buffer = Buffer::from(bytes).slice(1); + let nulls = NullBuffer::from(vec![true, false]); + let array = FixedSizeBinaryArray::new(width, buffer, Some(nulls.clone())); + let primitive = reinterpret_fixed_size_binary::(&array)?; + + assert_eq!(primitive.nulls(), Some(&nulls)); + Ok(()) + } +} diff --git a/datafusion/physical-expr/src/expressions/in_list/frozen_set.rs b/datafusion/physical-expr/src/expressions/in_list/frozen_set.rs new file mode 100644 index 0000000000000..46fcc2a74c1b7 --- /dev/null +++ b/datafusion/physical-expr/src/expressions/in_list/frozen_set.rs @@ -0,0 +1,251 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Immutable set for fixed-width values. + +use std::hash::{BuildHasher, Hash}; + +use datafusion_common::{Result, exec_datafusion_err}; +use hashbrown::{DefaultHashBuilder, HashTable}; + +/// Computes the keyed hash used to place values in a [`FrozenSet`]. +pub(super) trait FrozenSetHash { + fn hash_one(&self, value: V) -> u64; +} + +impl FrozenSetHash for DefaultHashBuilder { + #[inline(always)] + fn hash_one(&self, value: V) -> u64 { + BuildHasher::hash_one(self, value) + } +} + +/// Immutable set optimized for repeated membership tests. +/// +/// Each hash bucket holds two values directly. Values that collide beyond those +/// two slots use a `HashTable`, which keeps the common lookup path to one hash +/// and two bucket comparisons without sacrificing collision handling. +/// +/// The first member is used as the empty-slot sentinel and handled before +/// lookup, so primary slots store `V` directly without an `Option` wrapper. +pub(super) struct FrozenSet { + hash_builder: H, + sentinel: Option, + buckets: Box<[[V; 2]]>, + overflowed: Box<[bool]>, + overflow: HashTable, +} + +impl FrozenSet +where + V: Copy + Eq + Hash, +{ + pub(super) fn try_new(values: &[V]) -> Result { + Self::try_new_with_hasher(values, DefaultHashBuilder::default()) + } +} + +impl FrozenSet +where + V: Copy + Eq, + H: FrozenSetHash, +{ + pub(super) fn try_new_with_hasher(values: &[V], hash_builder: H) -> Result { + let Some((&sentinel, values)) = values.split_first() else { + return Ok(Self { + hash_builder, + sentinel: None, + buckets: Box::default(), + overflowed: Box::default(), + overflow: HashTable::new(), + }); + }; + + if values.is_empty() { + return Ok(Self { + hash_builder, + sentinel: Some(sentinel), + buckets: Box::default(), + overflowed: Box::default(), + overflow: HashTable::new(), + }); + } + + let bucket_count = values + .len() + .checked_add(1) + .ok_or_else(|| exec_datafusion_err!("FrozenSet capacity overflow"))?; + bucket_count + .checked_mul(2) + .ok_or_else(|| exec_datafusion_err!("FrozenSet capacity overflow"))?; + let mut set = Self { + hash_builder, + sentinel: Some(sentinel), + buckets: vec![[sentinel; 2]; bucket_count].into_boxed_slice(), + overflowed: vec![false; bucket_count].into_boxed_slice(), + overflow: HashTable::with_capacity(bucket_count / 8), + }; + values.iter().copied().for_each(|value| set.insert(value)); + Ok(set) + } + + fn insert(&mut self, value: V) { + let sentinel = self.sentinel.expect("non-empty frozen set"); + if value == sentinel { + return; + } + + let hash = self.hash_builder.hash_one(value); + let bucket = reduce_hash(hash, self.buckets.len()); + let slots = &mut self.buckets[bucket]; + if slots[0] == sentinel { + slots[0] = value; + return; + } + if slots[0] == value { + return; + } + if slots[1] == sentinel { + slots[1] = value; + return; + } + if slots[1] == value { + return; + } + + let hash_builder = &self.hash_builder; + self.overflow + .entry( + hash, + |stored| *stored == value, + |stored| hash_builder.hash_one(*stored), + ) + .or_insert(value); + self.overflowed[bucket] = true; + } + + #[inline(always)] + pub(super) fn contains(&self, value: V) -> bool { + let Some(sentinel) = self.sentinel else { + return false; + }; + if value == sentinel { + return true; + } + if self.buckets.is_empty() { + return false; + } + + let hash = self.hash_builder.hash_one(value); + let bucket = reduce_hash(hash, self.buckets.len()); + // SAFETY: `reduce_hash` returns an index in `0..buckets.len()`. + let slots = unsafe { self.buckets.get_unchecked(bucket) }; + if (slots[0] == value) | (slots[1] == value) { + return true; + } + + // SAFETY: `overflowed` has exactly one entry per bucket. + (unsafe { *self.overflowed.get_unchecked(bucket) }) + && self + .overflow + .find(hash, |stored| *stored == value) + .is_some() + } +} + +#[inline(always)] +fn reduce_hash(hash: u64, len: usize) -> usize { + #[cfg(target_pointer_width = "64")] + return ((hash as u128 * len as u128) >> 64) as usize; + + #[cfg(target_pointer_width = "32")] + return (((hash as u32) as u64 * len as u64) >> 32) as usize; +} + +#[cfg(test)] +mod tests { + use super::*; + use std::hash::{Hash, Hasher}; + + #[derive(Clone, Copy, Debug, Eq, PartialEq)] + struct Colliding(u64); + + struct ConstantHash; + + impl FrozenSetHash for ConstantHash { + fn hash_one(&self, _value: u64) -> u64 { + 0 + } + } + + impl Hash for Colliding { + fn hash(&self, state: &mut H) { + 0_u8.hash(state); + } + } + + #[test] + fn handles_empty_duplicates_and_sentinel_member() { + let empty = FrozenSet::::try_new(&[]).unwrap(); + assert!(!empty.contains(0)); + + let singleton = FrozenSet::try_new(&[42]).unwrap(); + assert!(singleton.contains(42)); + assert!(!singleton.contains(7)); + + let values = [42, 7, 42, 9, 11, 7]; + let set = FrozenSet::try_new(&values).unwrap(); + for value in [42, 7, 9, 11] { + assert!(set.contains(value)); + } + for value in [0, 8, 10, 12] { + assert!(!set.contains(value)); + } + } + + #[test] + fn handles_many_values() { + let values = (0_u64..10_000).collect::>(); + let set = FrozenSet::try_new(&values).unwrap(); + assert!(values.iter().all(|&value| set.contains(value))); + assert!((10_000..20_000).all(|value| !set.contains(value))); + } + + #[test] + fn handles_u128_values() { + let values = [1_u128, (1_u128 << 64) | 1, (2_u128 << 96) | 7, u128::MAX]; + let set = FrozenSet::try_new(&values).unwrap(); + assert!(values.iter().all(|&value| set.contains(value))); + assert!(!set.contains(1_u128 << 96)); + } + + #[test] + fn handles_collisions() { + let values = (0..128).map(Colliding).collect::>(); + let set = FrozenSet::try_new(&values).unwrap(); + assert!(values.iter().all(|&value| set.contains(value))); + assert!((128..256).all(|value| !set.contains(Colliding(value)))); + } + + #[test] + fn handles_custom_hash_collisions() { + let values = (0_u64..128).collect::>(); + let set = FrozenSet::try_new_with_hasher(&values, ConstantHash).unwrap(); + assert!(values.iter().all(|&value| set.contains(value))); + assert!((128..256).all(|value| !set.contains(value))); + } +} diff --git a/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs b/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs index 8f8d9bad04afa..881cf628aedb2 100644 --- a/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs @@ -19,13 +19,15 @@ //! //! This module provides membership tests for Arrow primitive types. -use arrow::array::{Array, ArrayRef, AsArray, BooleanArray}; -use arrow::buffer::{BooleanBuffer, NullBuffer}; +use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, PrimitiveArray}; +use arrow::buffer::{NullBuffer, ScalarBuffer}; use arrow::datatypes::*; use arrow::util::bit_iterator::BitIndexIterator; -use datafusion_common::{HashSet, Result, exec_datafusion_err}; -use std::hash::{Hash, Hasher}; +use datafusion_common::{Result, exec_datafusion_err}; +use std::hash::Hash; +use super::branchless_filter::{BranchlessFilterType, BranchlessNative}; +use super::frozen_set::FrozenSet; use super::result::build_in_list_result; use super::static_filter::{StaticFilter, handle_dictionary}; @@ -222,215 +224,110 @@ where } } -/// Wrapper for f32 that implements Hash and Eq using bit comparison. -/// This treats NaN values as equal to each other when they have the same bit pattern. -#[derive(Clone, Copy)] -struct OrderedFloat32(f32); - -impl Hash for OrderedFloat32 { - fn hash(&self, state: &mut H) { - self.0.to_ne_bytes().hash(state); - } +fn primitive_values(array: &PrimitiveArray) -> ScalarBuffer> +where + T: BranchlessFilterType, +{ + let data = array.to_data(); + ScalarBuffer::>::new( + data.buffers()[0].clone(), + data.offset(), + data.len(), + ) } -impl PartialEq for OrderedFloat32 { - fn eq(&self, other: &Self) -> bool { - self.0.to_bits() == other.0.to_bits() - } +/// Frozen-set filter for larger fixed-width primitive `IN` lists. +pub(super) struct PrimitiveFrozenFilter +where + BranchlessNative: Copy + Eq + Hash, +{ + expected_data_type: DataType, + null_count: usize, + values: FrozenSet>, } -impl Eq for OrderedFloat32 {} +impl PrimitiveFrozenFilter +where + T: BranchlessFilterType, + BranchlessNative: Copy + Eq + Hash, +{ + pub(super) fn try_new(in_array: &ArrayRef) -> Result { + let in_array = in_array.as_primitive_opt::().ok_or_else(|| { + exec_datafusion_err!("PrimitiveFrozenFilter: expected {} array", T::DATA_TYPE) + })?; -impl From for OrderedFloat32 { - fn from(v: f32) -> Self { - Self(v) + let null_count = in_array.null_count(); + let values = primitive_values::(in_array); + let values = match in_array.nulls() { + None => FrozenSet::try_new(values.as_ref())?, + Some(nulls) => { + let values = + BitIndexIterator::new(nulls.validity(), nulls.offset(), nulls.len()) + .map(|i| values[i]) + .collect::>(); + FrozenSet::try_new(&values)? + } + }; + Ok(Self { + expected_data_type: in_array.data_type().clone(), + null_count, + values, + }) } -} - -/// Wrapper for f64 that implements Hash and Eq using bit comparison. -/// This treats NaN values as equal to each other when they have the same bit pattern. -#[derive(Clone, Copy)] -struct OrderedFloat64(f64); -impl Hash for OrderedFloat64 { - fn hash(&self, state: &mut H) { - self.0.to_ne_bytes().hash(state); - } -} - -impl PartialEq for OrderedFloat64 { - fn eq(&self, other: &Self) -> bool { - self.0.to_bits() == other.0.to_bits() + #[inline] + fn contains_slice( + &self, + values: &[BranchlessNative], + nulls: Option<&NullBuffer>, + negated: bool, + ) -> BooleanArray { + build_in_list_result(values.len(), nulls, self.null_count > 0, negated, |i| { + // SAFETY: `build_in_list_result` invokes this closure for + // indices in `0..values.len()`. + let needle = unsafe { *values.get_unchecked(i) }; + self.values.contains(needle) + }) } } -impl Eq for OrderedFloat64 {} - -impl From for OrderedFloat64 { - fn from(v: f64) -> Self { - Self(v) +impl StaticFilter for PrimitiveFrozenFilter +where + T: BranchlessFilterType, + BranchlessNative: Copy + Eq + Hash + Send + Sync, +{ + fn null_count(&self) -> usize { + self.null_count } -} - -// Macro to generate specialized StaticFilter implementations for primitive types -macro_rules! primitive_static_filter { - ($Name:ident, $ArrowType:ty) => { - primitive_static_filter!( - $Name, - $ArrowType, - <$ArrowType as ArrowPrimitiveType>::Native, - |v| v - ); - }; - ($Name:ident, $ArrowType:ty, $SetValueType:ty, $to_set_value:expr) => { - pub(super) struct $Name { - null_count: usize, - values: HashSet<$SetValueType>, - } - - impl $Name { - pub(super) fn try_new(in_array: &ArrayRef) -> Result { - let in_array = - in_array.as_primitive_opt::<$ArrowType>().ok_or_else(|| { - exec_datafusion_err!( - "Failed to downcast an array to a '{}' array", - stringify!($ArrowType) - ) - })?; - - let mut values = HashSet::with_capacity(in_array.len()); - let null_count = in_array.null_count(); - - for v in in_array.iter().flatten() { - values.insert(($to_set_value)(v)); - } - Ok(Self { null_count, values }) - } - } - - impl StaticFilter for $Name { - fn null_count(&self) -> usize { - self.null_count - } + fn contains(&self, v: &dyn Array, negated: bool) -> Result { + handle_dictionary!(self, v, negated); - fn contains(&self, v: &dyn Array, negated: bool) -> Result { - handle_dictionary!(self, v, negated); - - let v = v.as_primitive_opt::<$ArrowType>().ok_or_else(|| { - exec_datafusion_err!( - "Failed to downcast an array to a '{}' array", - stringify!($ArrowType) - ) - })?; - - let haystack_has_nulls = self.null_count > 0; - let needle_values = v.values(); - let needle_nulls = v.nulls(); - let needle_has_nulls = v.null_count() > 0; - - // Truth table for `value [NOT] IN (set)` with SQL three-valued logic: - // ("-" means the value doesn't affect the result) - // - // | needle_null | haystack_null | negated | in set? | result | - // |-------------|---------------|---------|---------|--------| - // | true | - | false | - | null | - // | true | - | true | - | null | - // | false | true | false | yes | true | - // | false | true | false | no | null | - // | false | true | true | yes | false | - // | false | true | true | no | null | - // | false | false | false | yes | true | - // | false | false | false | no | false | - // | false | false | true | yes | false | - // | false | false | true | no | true | - - // Compute the "contains" result using collect_bool (fast batched approach) - // This ignores nulls - we handle them separately - let contains_buffer = if negated { - BooleanBuffer::collect_bool(needle_values.len(), |i| { - !self.values.contains(&($to_set_value)(needle_values[i])) - }) - } else { - BooleanBuffer::collect_bool(needle_values.len(), |i| { - self.values.contains(&($to_set_value)(needle_values[i])) - }) - }; - - // Compute the null mask - // Output is null when: - // 1. needle value is null, OR - // 2. needle value is not in set AND haystack has nulls - let result_nulls = match (needle_has_nulls, haystack_has_nulls) { - (false, false) => { - // No nulls anywhere - None - } - (true, false) => { - // Only needle has nulls - just use needle's null mask - needle_nulls.cloned() - } - (false, true) => { - // Only haystack has nulls - result is null when value not in set - // Valid (not null) when original "in set" is true - // For NOT IN: contains_buffer = !original, so validity = !contains_buffer - let validity = if negated { - !&contains_buffer - } else { - contains_buffer.clone() - }; - Some(NullBuffer::new(validity)) - } - (true, true) => { - // Both have nulls - combine needle nulls with haystack-induced nulls - let needle_validity = - needle_nulls.map(|n| n.inner().clone()).unwrap_or_else( - || BooleanBuffer::new_set(needle_values.len()), - ); - - // Valid when original "in set" is true (see above) - let haystack_validity = if negated { - !&contains_buffer - } else { - contains_buffer.clone() - }; - - // Combined validity: valid only where both are valid - let combined_validity = &needle_validity & &haystack_validity; - Some(NullBuffer::new(combined_validity)) - } - }; - - Ok(BooleanArray::new(contains_buffer, result_nulls)) - } + if !PrimitiveArray::::is_compatible(v.data_type()) { + return Err(exec_datafusion_err!( + "PrimitiveFrozenFilter: expected {} array, got {}", + self.expected_data_type, + v.data_type() + )); } - }; -} - -primitive_static_filter!(Int32StaticFilter, Int32Type); -primitive_static_filter!(Int64StaticFilter, Int64Type); -primitive_static_filter!(UInt32StaticFilter, UInt32Type); -primitive_static_filter!(UInt64StaticFilter, UInt64Type); -// Macro to generate specialized StaticFilter implementations for float types -// Floats require a wrapper type (OrderedFloat*) to implement Hash/Eq due to NaN semantics -macro_rules! float_static_filter { - ($Name:ident, $ArrowType:ty, $OrderedType:ty) => { - primitive_static_filter!($Name, $ArrowType, $OrderedType, <$OrderedType>::from); - }; + let v = v.as_primitive_opt::().ok_or_else(|| { + exec_datafusion_err!("PrimitiveFrozenFilter: expected {} array", T::DATA_TYPE) + })?; + let values = primitive_values::(v); + Ok(self.contains_slice(values.as_ref(), v.nulls(), negated)) + } } -// Generate specialized filters for float types using ordered wrappers -float_static_filter!(Float32StaticFilter, Float32Type, OrderedFloat32); -float_static_filter!(Float64StaticFilter, Float64Type, OrderedFloat64); - #[cfg(test)] mod tests { use super::*; use std::sync::Arc; use arrow::array::{ - DictionaryArray, Float16Array, Int8Array, Int16Array, UInt8Array, UInt16Array, + DictionaryArray, Float16Array, Float32Array, Int8Array, Int16Array, + TimestampMillisecondArray, TimestampNanosecondArray, UInt8Array, UInt16Array, + UInt32Array, }; use half::f16; @@ -584,4 +481,72 @@ mod tests { Ok(()) } + + #[test] + fn primitive_frozen_filter_handles_slices_and_nulls() -> Result<()> { + let haystack: ArrayRef = Arc::new( + UInt32Array::from(vec![ + Some(999), + Some(10), + None, + Some(20), + Some(10), + Some(30), + ]) + .slice(1, 5), + ); + let filter = PrimitiveFrozenFilter::::try_new(&haystack)?; + let needles = + UInt32Array::from(vec![Some(0), Some(10), Some(11), Some(30), None]) + .slice(1, 4); + + assert_contains(&filter, &needles, vec![Some(true), None, Some(true), None])?; + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![Some(false), None, Some(false), None]) + ); + + Ok(()) + } + + #[test] + fn primitive_frozen_filter_floats_use_bit_equality() -> Result<()> { + let nan_a = f32::from_bits(0x7fc0_0001); + let nan_b = f32::from_bits(0x7fc0_0002); + let haystack: ArrayRef = + Arc::new(Float32Array::from(vec![Some(-0.0), Some(nan_a)])); + let filter = PrimitiveFrozenFilter::::try_new(&haystack)?; + let needles = + Float32Array::from(vec![Some(0.0), Some(-0.0), Some(nan_a), Some(nan_b)]); + + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![Some(false), Some(true), Some(true), Some(false)]) + ); + + Ok(()) + } + + #[test] + fn primitive_frozen_filter_timestamp_uses_physical_compatibility() -> Result<()> { + let haystack: ArrayRef = Arc::new( + TimestampNanosecondArray::from(vec![Some(1), Some(3)]).with_timezone("UTC"), + ); + let filter = + PrimitiveFrozenFilter::::try_new(&haystack)?; + + let different_timezone = TimestampNanosecondArray::from(vec![Some(1), Some(2)]) + .with_timezone("Europe/Paris"); + assert_contains(&filter, &different_timezone, vec![Some(true), Some(false)])?; + + let different_unit = TimestampMillisecondArray::from(vec![Some(1)]); + let err = filter + .contains(&different_unit, false) + .unwrap_err() + .to_string(); + assert!(err.contains("Timestamp(ns"), "{err}"); + assert!(err.contains("Timestamp(ms"), "{err}"); + + Ok(()) + } } diff --git a/datafusion/physical-expr/src/expressions/in_list/result.rs b/datafusion/physical-expr/src/expressions/in_list/result.rs index 3ebdbfe19f743..2e6f41c241138 100644 --- a/datafusion/physical-expr/src/expressions/in_list/result.rs +++ b/datafusion/physical-expr/src/expressions/in_list/result.rs @@ -58,6 +58,34 @@ where build_result_from_contains(needle_nulls, haystack_has_nulls, negated, contains_buf) } +/// Builds an IN-list result without evaluating null needle values. +#[inline] +pub(crate) fn build_in_list_result_with_null_shortcircuit( + len: usize, + needle_nulls: Option<&NullBuffer>, + haystack_has_nulls: bool, + negated: bool, + mut contains: C, +) -> BooleanArray +where + C: FnMut(usize) -> bool, +{ + if needle_nulls.is_none_or(|nulls| nulls.null_count() == 0) { + return build_in_list_result( + len, + needle_nulls, + haystack_has_nulls, + negated, + contains, + ); + } + + let needle_nulls = needle_nulls.expect("null count was checked above"); + build_in_list_result(len, Some(needle_nulls), haystack_has_nulls, negated, |i| { + needle_nulls.is_valid(i) && contains(i) + }) +} + /// Builds a BooleanArray result from a pre-computed contains buffer. /// /// This version does not assume contains_buf is pre-masked at null positions. diff --git a/datafusion/physical-expr/src/expressions/in_list/strategy.rs b/datafusion/physical-expr/src/expressions/in_list/strategy.rs index d5ca8154a92f6..a741d25233d2c 100644 --- a/datafusion/physical-expr/src/expressions/in_list/strategy.rs +++ b/datafusion/physical-expr/src/expressions/in_list/strategy.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use std::sync::Arc; +use std::{hash::Hash, sync::Arc}; use arrow::array::ArrayRef; use arrow::compute::cast; @@ -34,111 +34,177 @@ 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 let Some(filter) = instantiate_branchless_filter(&in_array)? { + // These filters inspect their Arrow physical representation directly. + if dictionary_value_type(expr_data_type) == in_array.data_type() { + if let Some(filter) = instantiate_fixed_size_binary_filter(&in_array)? { + return Ok(filter); + } + if let Some(filter) = instantiate_byte_view_filter(&in_array)? { + return Ok(filter); + } + } + + if let Some(filter) = instantiate_primitive_filter(&in_array)? { return Ok(filter); } - instantiate_standard_filter(in_array) + Ok(Arc::new(ArrayStaticFilter::try_new(in_array)?)) +} + +fn 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 - // falling through to the generic ArrayStaticFilter. + // specialized primitive filters are used instead of falling through to + // the generic ArrayStaticFilter. match in_array.data_type() { DataType::Dictionary(_, value_type) => Ok(cast(&in_array, value_type.as_ref())?), _ => Ok(in_array), } } -fn instantiate_branchless_filter(in_array: &ArrayRef) -> Result> { +fn instantiate_primitive_filter(in_array: &ArrayRef) -> Result> { let non_null_count = in_array.len() - in_array.null_count(); macro_rules! filter { - ($arrow_type:ty) => { - branchless_filter::<$arrow_type>(in_array, non_null_count) + ($arrow_type:ty, $strategy:ident) => { + $strategy::<$arrow_type>(in_array, non_null_count) }; } match in_array.data_type() { - DataType::Int8 => filter!(Int8Type), - DataType::UInt8 => filter!(UInt8Type), - DataType::Int16 => filter!(Int16Type), - DataType::UInt16 => filter!(UInt16Type), - DataType::Float16 => filter!(Float16Type), - DataType::Int32 => filter!(Int32Type), - DataType::UInt32 => filter!(UInt32Type), - DataType::Float32 => filter!(Float32Type), - DataType::Date32 => filter!(Date32Type), + DataType::Int8 => filter!(Int8Type, branchless_or_bitmap_filter), + DataType::UInt8 => filter!(UInt8Type, branchless_or_bitmap_filter), + DataType::Int16 => filter!(Int16Type, branchless_or_bitmap_filter), + DataType::UInt16 => filter!(UInt16Type, branchless_or_bitmap_filter), + DataType::Float16 => filter!(Float16Type, branchless_or_bitmap_filter), + DataType::Int32 => filter!(Int32Type, branchless_or_frozen_filter), + DataType::UInt32 => filter!(UInt32Type, branchless_or_frozen_filter), + DataType::Float32 => filter!(Float32Type, branchless_or_frozen_filter), + DataType::Date32 => filter!(Date32Type, branchless_or_frozen_filter), DataType::Time32(unit) => match unit { - TimeUnit::Second => filter!(Time32SecondType), - TimeUnit::Millisecond => filter!(Time32MillisecondType), + TimeUnit::Second => { + filter!(Time32SecondType, branchless_or_frozen_filter) + } + TimeUnit::Millisecond => { + filter!(Time32MillisecondType, branchless_or_frozen_filter) + } _ => Ok(None), }, - DataType::Int64 => filter!(Int64Type), - DataType::UInt64 => filter!(UInt64Type), - DataType::Float64 => filter!(Float64Type), - DataType::Date64 => filter!(Date64Type), + DataType::Int64 => filter!(Int64Type, branchless_or_frozen_filter), + DataType::UInt64 => filter!(UInt64Type, branchless_or_frozen_filter), + DataType::Float64 => filter!(Float64Type, branchless_or_frozen_filter), + DataType::Date64 => filter!(Date64Type, branchless_or_frozen_filter), DataType::Time64(unit) => match unit { - TimeUnit::Microsecond => filter!(Time64MicrosecondType), - TimeUnit::Nanosecond => filter!(Time64NanosecondType), + TimeUnit::Microsecond => { + filter!(Time64MicrosecondType, branchless_or_frozen_filter) + } + TimeUnit::Nanosecond => { + filter!(Time64NanosecondType, branchless_or_frozen_filter) + } _ => Ok(None), }, DataType::Timestamp(unit, _) => match unit { - TimeUnit::Second => filter!(TimestampSecondType), - TimeUnit::Millisecond => filter!(TimestampMillisecondType), - TimeUnit::Microsecond => filter!(TimestampMicrosecondType), - TimeUnit::Nanosecond => filter!(TimestampNanosecondType), + TimeUnit::Second => { + filter!(TimestampSecondType, branchless_or_frozen_filter) + } + TimeUnit::Millisecond => { + filter!(TimestampMillisecondType, branchless_or_frozen_filter) + } + TimeUnit::Microsecond => { + filter!(TimestampMicrosecondType, branchless_or_frozen_filter) + } + TimeUnit::Nanosecond => { + filter!(TimestampNanosecondType, branchless_or_frozen_filter) + } }, DataType::Duration(unit) => match unit { - TimeUnit::Second => filter!(DurationSecondType), - TimeUnit::Millisecond => filter!(DurationMillisecondType), - TimeUnit::Microsecond => filter!(DurationMicrosecondType), - TimeUnit::Nanosecond => filter!(DurationNanosecondType), + TimeUnit::Second => { + filter!(DurationSecondType, branchless_or_frozen_filter) + } + TimeUnit::Millisecond => { + filter!(DurationMillisecondType, branchless_or_frozen_filter) + } + TimeUnit::Microsecond => { + filter!(DurationMicrosecondType, branchless_or_frozen_filter) + } + TimeUnit::Nanosecond => { + filter!(DurationNanosecondType, branchless_or_frozen_filter) + } }, - DataType::Decimal128(_, _) => filter!(Decimal128Type), + DataType::Decimal128(_, _) => { + filter!(Decimal128Type, branchless_filter) + } DataType::Interval(IntervalUnit::MonthDayNano) => { - filter!(IntervalMonthDayNanoType) + filter!(IntervalMonthDayNanoType, branchless_filter) } _ => Ok(None), } } -fn instantiate_standard_filter(in_array: ArrayRef) -> Result { - match in_array.data_type() { - DataType::Int8 => bitmap_filter::(&in_array), - DataType::UInt8 => bitmap_filter::(&in_array), - DataType::Int16 => bitmap_filter::(&in_array), - DataType::UInt16 => bitmap_filter::(&in_array), - DataType::Float16 => bitmap_filter::(&in_array), - DataType::Int32 => Ok(Arc::new(Int32StaticFilter::try_new(&in_array)?)), - DataType::Int64 => Ok(Arc::new(Int64StaticFilter::try_new(&in_array)?)), - DataType::UInt32 => Ok(Arc::new(UInt32StaticFilter::try_new(&in_array)?)), - DataType::UInt64 => Ok(Arc::new(UInt64StaticFilter::try_new(&in_array)?)), - // Float primitive types (use ordered wrappers for Hash/Eq) - DataType::Float32 => Ok(Arc::new(Float32StaticFilter::try_new(&in_array)?)), - DataType::Float64 => Ok(Arc::new(Float64StaticFilter::try_new(&in_array)?)), - _ => { - // Fall through to generic implementation for unsupported types - // (Struct, etc.). - Ok(Arc::new(ArrayStaticFilter::try_new(in_array)?)) - } +fn branchless_or_bitmap_filter( + in_array: &ArrayRef, + non_null_count: usize, +) -> Result> +where + T: BranchlessFilterType + BitmapFilterType, + BranchlessNative: Copy + PartialEq + Send + Sync, +{ + if let Some(filter) = branchless_filter::(in_array, non_null_count)? { + return Ok(Some(filter)); + } + + Ok(Some(Arc::new(BitmapFilter::::try_new(in_array)?))) +} + +fn branchless_or_frozen_filter( + in_array: &ArrayRef, + non_null_count: usize, +) -> Result> +where + T: BranchlessFilterType, + BranchlessNative: Copy + Eq + Hash + Send + Sync, +{ + if let Some(filter) = branchless_filter::(in_array, non_null_count)? { + return Ok(Some(filter)); } + + primitive_frozen_filter::(in_array, non_null_count) } -fn bitmap_filter(in_array: &ArrayRef) -> Result +fn primitive_frozen_filter( + in_array: &ArrayRef, + non_null_count: usize, +) -> Result> where - T: BitmapFilterType, + T: BranchlessFilterType, + BranchlessNative: Copy + Eq + Hash + Send + Sync, { - Ok(Arc::new(BitmapFilter::::try_new(in_array)?)) + if non_null_count <= T::MAX_LIST_LEN { + return Ok(None); + } + + Ok(Some(Arc::new(PrimitiveFrozenFilter::::try_new( + in_array, + )?))) } fn branchless_filter( @@ -149,7 +215,7 @@ where T: BranchlessFilterType, BranchlessNative: Copy + PartialEq + Send + Sync, { - // Larger lists use the standard filter. `try_new` checks the limit again. + // Larger lists use another filter strategy. `try_new` checks the limit again. if non_null_count > T::MAX_LIST_LEN { return Ok(None); } @@ -159,7 +225,7 @@ where #[cfg(test)] mod tests { - use arrow::array::UInt32Array; + use arrow::array::{Decimal128Array, UInt32Array}; use arrow::datatypes::UInt32Type; use super::super::branchless_filter::BranchlessFilterType; @@ -176,13 +242,48 @@ mod tests { let values = (0..max_len) .map(|value| Some(value as u32)) .collect::>(); - assert!(instantiate_branchless_filter(&uint32_array(values))?.is_some()); + assert!( + branchless_filter::(&uint32_array(values), max_len)?.is_some() + ); + + let values = (0..=max_len) + .map(|value| Some(value as u32)) + .collect::>(); + assert!( + branchless_filter::(&uint32_array(values), max_len + 1)? + .is_none() + ); + + Ok(()) + } + + #[test] + fn primitive_frozen_routing_starts_after_max_list_len() -> Result<()> { + let max_len = ::MAX_LIST_LEN; + + let values = (0..max_len) + .map(|value| Some(value as u32)) + .collect::>(); + assert!( + primitive_frozen_filter::(&uint32_array(values), max_len)? + .is_none() + ); let values = (0..=max_len) .map(|value| Some(value as u32)) .collect::>(); - assert!(instantiate_branchless_filter(&uint32_array(values))?.is_none()); + assert!( + primitive_frozen_filter::(&uint32_array(values), max_len + 1)? + .is_some() + ); + + Ok(()) + } + #[test] + fn primitive_frozen_routing_excludes_128_bit_values() -> Result<()> { + let array: ArrayRef = Arc::new(Decimal128Array::from(vec![1, 2, 3, 4, 5])); + assert!(instantiate_primitive_filter(&array)?.is_none()); Ok(()) } @@ -190,8 +291,29 @@ mod tests { fn branchless_routing_handles_zero_non_null_values() -> Result<()> { let array = uint32_array(vec![None; 3]); - assert!(instantiate_branchless_filter(&array)?.is_some()); + assert!(branchless_filter::(&array, 0)?.is_some()); Ok(()) } + + #[test] + fn physical_type_gate_recursively_unwraps_dictionaries() { + let value_type = DataType::FixedSizeBinary(16); + let expr_type = DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::Dictionary( + Box::new(DataType::Int16), + Box::new(value_type.clone()), + )), + ); + + assert_eq!(dictionary_value_type(&expr_type), &value_type); + assert_ne!( + dictionary_value_type(&DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::FixedSizeBinary(8)), + )), + &value_type + ); + } }