diff --git a/datafusion/physical-expr/benches/in_list_strategy.rs b/datafusion/physical-expr/benches/in_list_strategy.rs index 762c2c97a1115..24ca0ae17658e 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 //! @@ -45,14 +45,20 @@ //! | 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 | +//! | Dictionary dispatch cases | Int32, Dictionary | per-batch dispatch overhead | 4, 64 | +//! | Fixed-size binary direct-comparison case | FixedSizeBinary(1) | direct-comparison cutoff | 16 | +//! | Fixed-size binary direct-comparison case | FixedSizeBinary(16) | direct-comparison cutoff | 4 | +//! | Fixed-size binary bitmap case | FixedSizeBinary(2) | bitmap lookup | 64 | +//! | Fixed-size binary hash-set cases | FixedSizeBinary(16) | hash-set scaling | 64, 256, 10000 | +//! | Fixed-size binary unaligned case | FixedSizeBinary(16) | per-evaluation alignment copy | 64 | use arrow::array::types::IntervalMonthDayNano; use arrow::array::*; +use arrow::buffer::{Buffer, MutableBuffer}; 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; @@ -866,11 +872,65 @@ fn bench_dictionary(c: &mut Criterion) { bench_dict_string(c, "utf8_short/dict=500/list=20", 500, 20, 10); } +/// Measures the fixed per-batch cost of handling dictionary and plain inputs. +/// Small batches make dispatch overhead visible, while the standard 8,192-row +/// batch shows whether it matters in the usual vectorized case. +fn bench_dictionary_dispatch(c: &mut Criterion) { + for list_size in [4_usize, 64] { + let strategy = if list_size == 4 { + "branchless" + } else { + "hash_set" + }; + let haystack = (0..list_size as i32).collect::>(); + + for batch_size in [1, 8, 64, ARRAY_SIZE] { + let keys = (0..batch_size) + .map(|index| (index % 8) as i32) + .collect::>(); + let values = keys.iter().map(|key| key * 2).collect::>(); + + for dictionary in [false, true] { + let array: ArrayRef = if dictionary { + Arc::new( + DictionaryArray::::try_new( + Int32Array::from(keys.clone()), + Arc::new(Int32Array::from_iter_values((0..8).map(|v| v * 2))), + ) + .unwrap(), + ) + } else { + Arc::new(Int32Array::from(values.clone())) + }; + + let schema = + Schema::new(vec![Field::new("a", array.data_type().clone(), false)]); + let exprs = haystack.iter().map(|value| lit(*value)).collect(); + let expr = + in_list(col("a", &schema).unwrap(), exprs, &false, &schema).unwrap(); + let batch = RecordBatch::try_new(Arc::new(schema), vec![array]).unwrap(); + let encoding = if dictionary { "dictionary" } else { "plain" }; + + c.bench_with_input( + BenchmarkId::new( + "dictionary_dispatch", + format!( + "i32/{strategy}/{encoding}/list={list_size}/batch={batch_size}" + ), + ), + &batch, + |b, batch| b.iter(|| expr.evaluate(batch).unwrap()), + ); + } + } + } +} + // ============================================================================= // 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) { // ========================================================================= @@ -1014,72 +1074,180 @@ fn bench_nulls(c: &mut Criterion) { } // ============================================================================= -// FIXED SIZE BINARY BENCHMARKS (FixedSizeBinary<16>, e.g. UUIDs) +// FIXED SIZE BINARY BENCHMARKS // ============================================================================= -/// Generates a random 16-byte value (UUID-sized). -fn random_fixed_binary_16(rng: &mut StdRng) -> Vec { - let mut buf = vec![0u8; 16]; +fn random_fixed_binary(rng: &mut StdRng, width: i32) -> Vec { + let mut buf = vec![0u8; width as usize]; rng.fill(&mut buf[..]); buf } -/// Benchmarks FixedSizeBinary(16) IN list evaluation. -/// FixedSizeBinary doesn't use the generic numeric helpers since its array -/// construction differs from primitive types. -fn bench_fixed_size_binary_inner( - c: &mut Criterion, - name: &str, +#[derive(Clone, Copy)] +enum InputLayout { + Aligned, + UnalignedI128, +} + +#[derive(Clone, Copy)] +struct FixedSizeBinaryBenchConfig { + width: i32, + list_size: usize, + input_layout: InputLayout, +} + +impl FixedSizeBinaryBenchConfig { + const fn aligned(width: i32, list_size: usize) -> Self { + Self { + width, + list_size, + input_layout: InputLayout::Aligned, + } + } + + const fn unaligned_i128(list_size: usize) -> Self { + Self { + width: 16, + list_size, + input_layout: InputLayout::UnalignedI128, + } + } +} + +const FIXED_SIZE_BINARY_CASES: [FixedSizeBinaryBenchConfig; 7] = [ + FixedSizeBinaryBenchConfig::aligned(1, 16), + FixedSizeBinaryBenchConfig::aligned(2, 64), + FixedSizeBinaryBenchConfig::aligned(16, 4), + FixedSizeBinaryBenchConfig::aligned(16, 64), + FixedSizeBinaryBenchConfig::aligned(16, 256), + FixedSizeBinaryBenchConfig::aligned(16, 10000), + // 8,192 rows at 16 bytes each copy 128 KiB per evaluation. + FixedSizeBinaryBenchConfig::unaligned_i128(64), +]; + +fn generate_fixed_size_binary_data( + rng: &mut StdRng, + width: i32, list_size: usize, match_rate: f64, -) { - let seed = 0xF1ED_B1A7_u64.wrapping_add(list_size as u64 * 0x6666); - let mut rng = StdRng::seed_from_u64(seed); +) -> (Vec>, Vec>) { + if let Some(domain_size) = match width { + 1 => Some(1_usize << 8), + 2 => Some(1_usize << 16), + _ => None, + } { + // The value generator needs at least one value outside the haystack. + assert!(list_size < domain_size); + } - // 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 haystack 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(rng, width); + if haystack_set.insert(value.clone()) { + haystack.push(value); + } + } - // Generate array with controlled match rate - let values: Vec> = (0..ARRAY_SIZE) + // Generate values with the configured match rate. + let values = (0..ARRAY_SIZE) .map(|_| { if !haystack.is_empty() && rng.random_bool(match_rate) { - haystack.choose(&mut rng).unwrap().clone() + haystack.choose(rng).unwrap().clone() } else { - random_fixed_binary_16(&mut rng) + loop { + let value = random_fixed_binary(rng, width); + if !haystack_set.contains(&value) { + break value; + } + } } }) .collect(); - let refs: Vec<&[u8]> = values.iter().map(|v| v.as_slice()).collect(); - let array = FixedSizeBinaryArray::try_from_iter(refs.into_iter()).unwrap(); + (haystack, values) +} + +fn unaligned_fixed_size_binary_16(values: &[Vec]) -> FixedSizeBinaryArray { + const WIDTH: usize = 16; + let payload_len = values.len() * WIDTH; + let mut bytes = MutableBuffer::with_capacity(payload_len + 1); + bytes.push(0_u8); + for value in values { + assert_eq!(value.len(), WIDTH); + bytes.extend_from_slice(value); + } + + // MutableBuffer starts at an Arrow-aligned address. Fixed-size binary + // values only require byte alignment, so slicing off this padding byte + // creates a valid Arrow buffer that models unaligned external input. + let buffer = Buffer::from(bytes).slice(1); + assert!( + !buffer.as_ptr().cast::().is_aligned(), + "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, + config: FixedSizeBinaryBenchConfig, + match_pct: u32, +) { + assert!(match_pct <= 100); + let match_rate = f64::from(match_pct) / 100.0; + + let seed = 0xF1ED_B1A7_u64 + .wrapping_add(config.list_size as u64 * 0x6666) + .wrapping_add(config.width as u64 * 0x7777); + let mut rng = StdRng::seed_from_u64(seed); + + let (haystack, values) = generate_fixed_size_binary_data( + &mut rng, + config.width, + config.list_size, + match_rate, + ); + + let array = match config.input_layout { + InputLayout::Aligned => { + FixedSizeBinaryArray::try_from_iter(values.iter().map(Vec::as_slice)).unwrap() + } + InputLayout::UnalignedI128 => unaligned_fixed_size_binary_16(&values), + }; 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(config.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{}/list={}/match={match_pct}%", + config.width, config.list_size + ); + match config.input_layout { + InputLayout::Aligned => name, + InputLayout::UnalignedI128 => format!("{name}/input=unaligned"), + } + }), &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 config in FIXED_SIZE_BINARY_CASES { 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, config, match_pct); } } } @@ -1091,7 +1259,7 @@ fn bench_fixed_size_binary(c: &mut Criterion) { criterion_group! { name = benches; config = Criterion::default(); - targets = bench_narrow_integer, bench_primitive, bench_f32, bench_timestamp_ns, bench_interval_month_day_nano, bench_utf8, bench_utf8view, bench_dictionary, bench_nulls, bench_fixed_size_binary + targets = bench_narrow_integer, bench_primitive, bench_f32, bench_timestamp_ns, bench_interval_month_day_nano, bench_utf8, bench_utf8view, bench_dictionary, bench_dictionary_dispatch, bench_nulls, bench_fixed_size_binary } criterion_main!(benches); diff --git a/datafusion/physical-expr/src/expressions/in_list.rs b/datafusion/physical-expr/src/expressions/in_list.rs index 0fb978cd0bafe..3c260190b8208 100644 --- a/datafusion/physical-expr/src/expressions/in_list.rs +++ b/datafusion/physical-expr/src/expressions/in_list.rs @@ -38,6 +38,8 @@ use datafusion_expr::{ColumnarValue, expr_vec_fmt}; mod array_static_filter; mod branchless_filter; +mod dictionary_filter; +mod fixed_size_binary_filter; mod primitive_filter; mod result; mod static_filter; @@ -215,7 +217,7 @@ impl InListExpr { expr, list, negated, - Some(instantiate_static_filter(array)?), + Some(instantiate_static_filter(array, &expr_data_type)?), )) } @@ -242,7 +244,7 @@ impl InListExpr { // Try to create a static filter if all list expressions are constants let static_filter = match try_evaluate_constant_list(&list, schema)? { - Some(in_array) => Some(instantiate_static_filter(in_array)?), + Some(in_array) => Some(instantiate_static_filter(in_array, &expr_data_type)?), None => None, // Non-constant expressions, fall back to dynamic evaluation }; @@ -1264,6 +1266,31 @@ mod tests { Ok(()) } + #[test] + fn in_list_nested_dictionary_scalar() -> Result<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let batch = RecordBatch::try_new( + Arc::new(schema.clone()), + vec![Arc::new(Int32Array::from(vec![0, 0, 0]))], + )?; + let needle = lit(ScalarValue::Dictionary( + Box::new(DataType::Int16), + Box::new(ScalarValue::Dictionary( + Box::new(DataType::Int8), + Box::new(ScalarValue::Int32(Some(2))), + )), + )); + + let expr = in_list(needle, vec![lit(1_i32), lit(2_i32)], &false, &schema)?; + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + assert_eq!( + as_boolean_array(&result), + &BooleanArray::from(vec![true, true, true]) + ); + + Ok(()) + } + #[test] fn in_list_utf8_with_dict_types() -> Result<()> { fn dict_lit(key_type: DataType, value: &str) -> Arc { @@ -3548,6 +3575,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/array_static_filter.rs b/datafusion/physical-expr/src/expressions/in_list/array_static_filter.rs index 75e92dbcc59b4..c2d3b4728274e 100644 --- a/datafusion/physical-expr/src/expressions/in_list/array_static_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/array_static_filter.rs @@ -15,12 +15,9 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{ - Array, ArrayRef, BooleanArray, downcast_array, downcast_dictionary_array, - make_comparator, -}; +use arrow::array::{Array, ArrayRef, BooleanArray, make_comparator}; use arrow::buffer::{BooleanBuffer, NullBuffer}; -use arrow::compute::{SortOptions, take}; +use arrow::compute::SortOptions; use arrow::datatypes::DataType; use arrow::util::bit_iterator::BitIndexIterator; use datafusion_common::Result; @@ -141,22 +138,6 @@ impl StaticFilter for ArrayStaticFilter { )); } - // Unwrap dictionary-encoded needles when the value type matches - // in_array, evaluating against the dictionary values and mapping - // back via keys. - downcast_dictionary_array! { - v => { - // Only unwrap when the haystack (in_array) type matches - // the dictionary value type - if v.values().data_type() == self.in_array.data_type() { - let values_contains = self.contains(v.values().as_ref(), negated)?; - let result = take(&values_contains, v.keys(), None)?; - return Ok(downcast_array(result.as_ref())); - } - } - _ => {} - } - self.find_needles_in_haystack(v, negated) } } 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 8539dc7956dc4..b4d4c10fd7e7c 100644 --- a/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs @@ -77,7 +77,7 @@ use arrow::util::bit_iterator::BitIndexIterator; use datafusion_common::{Result, exec_datafusion_err, internal_datafusion_err}; use super::result::build_result_from_contains; -use super::static_filter::{StaticFilter, handle_dictionary}; +use super::static_filter::StaticFilter; pub(super) type BranchlessNative = <::CompareType as ArrowPrimitiveType>::Native; @@ -256,8 +256,6 @@ where } fn contains(&self, v: &dyn Array, negated: bool) -> Result { - handle_dictionary!(self, v, negated); - // Arrow compatibility ignores timestamp timezone and decimal precision/scale // while still requiring the same primitive representation. if !PrimitiveArray::::is_compatible(v.data_type()) { diff --git a/datafusion/physical-expr/src/expressions/in_list/dictionary_filter.rs b/datafusion/physical-expr/src/expressions/in_list/dictionary_filter.rs new file mode 100644 index 0000000000000..01de6f28d640e --- /dev/null +++ b/datafusion/physical-expr/src/expressions/in_list/dictionary_filter.rs @@ -0,0 +1,161 @@ +// 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. + +use arrow::array::{Array, AsArray, BooleanArray}; +use arrow::compute::take; +use arrow::datatypes::DataType; +use datafusion_common::Result; + +use super::static_filter::{StaticFilter, StaticFilterRef}; + +/// Adds dictionary-encoded needle support to a filter that expects plain arrays. +/// This wrapper is only used when the input expression returns dictionaries. +pub(super) struct DictionaryFilter { + inner: StaticFilterRef, +} + +impl DictionaryFilter { + pub(super) fn new(inner: StaticFilterRef) -> Self { + Self { inner } + } +} + +impl StaticFilter for DictionaryFilter { + fn null_count(&self) -> usize { + self.inner.null_count() + } + + fn contains(&self, needles: &dyn Array, negated: bool) -> Result { + let Some(dictionary) = needles.as_any_dictionary_opt() else { + return self.inner.contains(needles, negated); + }; + let values = dictionary.values(); + let values_contains = if matches!(values.data_type(), DataType::Dictionary(_, _)) + { + self.contains(values.as_ref(), negated)? + } else { + self.inner.contains(values.as_ref(), negated)? + }; + let result = take(&values_contains, dictionary.keys(), None)?; + Ok(result.as_boolean().clone()) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::array::{ + ArrayRef, Decimal128Array, DictionaryArray, Int8Array, Int16Array, Int32Array, + Int64Array, TimestampNanosecondArray, UInt8Array, UInt16Array, UInt32Array, + UInt64Array, + }; + + use super::super::strategy::instantiate_static_filter; + use super::*; + + #[test] + fn dictionary_needles_support_all_key_types() -> Result<()> { + let values: ArrayRef = Arc::new(Int32Array::from(vec![1, 2])); + let expected = BooleanArray::from(vec![Some(true), Some(false), None]); + + macro_rules! check_keys { + ($keys:expr) => {{ + let needles = DictionaryArray::try_new($keys, Arc::clone(&values))?; + let filter = instantiate_static_filter( + Arc::new(Int32Array::from(vec![1, 3])), + needles.data_type(), + )?; + assert_eq!(filter.contains(&needles, false)?, expected); + }}; + } + + check_keys!(Int8Array::from(vec![Some(0), Some(1), None])); + check_keys!(Int16Array::from(vec![Some(0), Some(1), None])); + check_keys!(Int32Array::from(vec![Some(0), Some(1), None])); + check_keys!(Int64Array::from(vec![Some(0), Some(1), None])); + check_keys!(UInt8Array::from(vec![Some(0), Some(1), None])); + check_keys!(UInt16Array::from(vec![Some(0), Some(1), None])); + check_keys!(UInt32Array::from(vec![Some(0), Some(1), None])); + check_keys!(UInt64Array::from(vec![Some(0), Some(1), None])); + + Ok(()) + } + + #[test] + fn nested_dictionary_needles_preserve_nulls() -> Result<()> { + let values: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), Some(2), None])); + let inner: ArrayRef = Arc::new(DictionaryArray::try_new( + Int8Array::from(vec![0, 1, 2]), + values, + )?); + let needles = DictionaryArray::try_new( + Int16Array::from(vec![Some(0), Some(1), Some(2), None]), + inner, + )?; + let filter = instantiate_static_filter( + Arc::new(Int32Array::from(vec![1, 3])), + needles.data_type(), + )?; + + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![Some(true), Some(false), None, None]) + ); + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![Some(false), Some(true), None, None]) + ); + + Ok(()) + } + + #[test] + fn dictionary_timestamp_needles_keep_timezone_compatibility() -> Result<()> { + let timestamps: ArrayRef = + Arc::new(TimestampNanosecondArray::from(vec![1, 3]).with_timezone("UTC")); + let values: ArrayRef = Arc::new( + TimestampNanosecondArray::from(vec![1, 2]).with_timezone("Europe/Paris"), + ); + let needles = DictionaryArray::try_new(Int8Array::from(vec![0, 1]), values)?; + let filter = instantiate_static_filter(timestamps, needles.data_type())?; + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![true, false]) + ); + + Ok(()) + } + + #[test] + fn dictionary_decimal_needles_keep_precision_scale_compatibility() -> Result<()> { + // Five list values use the hash-set filter. + let decimals: ArrayRef = Arc::new( + Decimal128Array::from(vec![1, 3, 5, 7, 9]).with_precision_and_scale(10, 2)?, + ); + let values: ArrayRef = + Arc::new(Decimal128Array::from(vec![1, 2]).with_precision_and_scale(11, 3)?); + let needles = DictionaryArray::try_new(Int8Array::from(vec![0, 1]), values)?; + let filter = instantiate_static_filter(decimals, needles.data_type())?; + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![true, false]) + ); + + 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..ddad5a70994ac --- /dev/null +++ b/datafusion/physical-expr/src/expressions/in_list/fixed_size_binary_filter.rs @@ -0,0 +1,370 @@ +// 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::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}; + +/// 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 array.value_size() != width { + return Err(internal_datafusion_err!( + "FixedSizeBinary filter: expected {width}-byte values, got {}", + array.value_size() + )); + } + + let source = array.values(); + let values = if source.as_ptr().cast::().is_aligned() { + ScalarBuffer::new(source.clone(), 0, array.len()) + } else { + // `Buffer::from(&[u8])` copies into Arrow-aligned storage. + 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 { + 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, MutableBuffer, NullBuffer}; + use arrow::datatypes::Int8Type; + + use super::super::dictionary_filter::DictionaryFilter; + 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 = DictionaryFilter::new(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_i128_array( + values: &[Vec], + nulls: Option, + ) -> FixedSizeBinaryArray { + let width = size_of::(); + let mut bytes = MutableBuffer::with_capacity(1 + width * values.len()); + bytes.push(0_u8); + for value in values { + assert_eq!(value.len(), width); + bytes.extend_from_slice(value); + } + let buffer = Buffer::from(bytes).slice(1); + assert!( + !buffer.as_ptr().cast::().is_aligned(), + "test buffer must be unaligned" + ); + FixedSizeBinaryArray::new(width as i32, 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_i128_array(&haystack_values, None)); + let needles = unaligned_i128_array( + &[ + 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 bfafae31ed622..ceb9bd525b965 100644 --- a/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs @@ -19,17 +19,18 @@ //! //! This module provides membership tests for Arrow primitive types. +use std::hash::{Hash, Hasher}; +use std::marker::PhantomData; +use std::sync::Arc; + use arrow::array::{Array, ArrayRef, AsArray, BooleanArray}; 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, StaticFilterRef, handle_dictionary}; +use super::static_filter::{StaticFilter, StaticFilterRef}; /// Selects an optimized filter for a primitive representation. /// @@ -311,7 +312,6 @@ where } fn contains(&self, v: &dyn Array, negated: bool) -> Result { - handle_dictionary!(self, v, negated); let v = v.as_primitive_opt::().ok_or_else(|| { exec_datafusion_err!("BitmapFilter: expected {} array", T::DATA_TYPE) })?; @@ -432,8 +432,6 @@ where } fn contains(&self, v: &dyn Array, negated: bool) -> Result { - handle_dictionary!(self, v, negated); - let v = v.as_primitive_opt::().ok_or_else(|| { exec_datafusion_err!( "PrimitiveHashSetFilter: expected {} array", @@ -465,6 +463,8 @@ mod tests { }; use half::f16; + use super::super::dictionary_filter::DictionaryFilter; + fn uint32_array(values: Vec>) -> ArrayRef { Arc::new(UInt32Array::from(values)) } @@ -553,13 +553,19 @@ mod tests { #[test] fn bitmap_filter_u8_handles_dictionary_needles() -> Result<()> { let haystack: ArrayRef = Arc::new(UInt8Array::from(vec![Some(1), None, Some(3)])); - let filter = BitmapFilter::::try_new(&haystack)?; + let inner: StaticFilterRef = + Arc::new(BitmapFilter::::try_new(&haystack)?); + let filter = DictionaryFilter::new(inner); let keys = Int8Array::from(vec![Some(0), Some(1), None, Some(2)]); let values = Arc::new(UInt8Array::from(vec![Some(1), Some(2), Some(3)])); let needles = DictionaryArray::try_new(keys, values)?; - assert_contains(&filter, &needles, vec![Some(true), None, None, Some(true)]) + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![Some(true), None, None, Some(true)]) + ); + Ok(()) } #[test] 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 e74ea9131e63b..255d677c94b2a 100644 --- a/datafusion/physical-expr/src/expressions/in_list/static_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/static_filter.rs @@ -20,39 +20,19 @@ use std::sync::Arc; use arrow::array::{Array, BooleanArray}; use datafusion_common::Result; -pub(super) type StaticFilterRef = Arc; +pub(super) type StaticFilterRef = Arc; /// Trait for InList static filters. /// /// Static filters store a pre-computed set of values (the haystack) and check /// whether needle values are contained in that set. The haystack is always /// represented in its non-dictionary (value) type. Dictionary haystacks are -/// flattened via `cast()` before construction. +/// flattened before construction. /// -/// Dictionary-encoded needles are unwrapped inside `contains()` and -/// evaluated against the dictionary's values. -pub(super) trait StaticFilter { +/// Dictionary-encoded needles are unwrapped before the concrete filter is called. +pub(super) trait StaticFilter: Send + Sync { fn null_count(&self) -> usize; - /// Checks if values in `v` (needle) are contained in this filter's - /// haystack. `v` may be dictionary-encoded, in which case the - /// implementation unwraps the dictionary and operates on its values. - fn contains(&self, v: &dyn Array, negated: bool) -> Result; + /// Checks if non-dictionary needles are contained in this filter's haystack. + fn contains(&self, needles: &dyn Array, negated: bool) -> Result; } - -/// Evaluate dictionary-encoded needles by applying a filter to dictionary -/// values and remapping the result through the keys. -macro_rules! handle_dictionary { - ($self:ident, $v:ident, $negated:ident) => { - arrow::array::downcast_dictionary_array! { - $v => { - let values_contains = $self.contains($v.values().as_ref(), $negated)?; - let result = arrow::compute::take(&values_contains, $v.keys(), None)?; - return Ok(arrow::array::downcast_array(result.as_ref())) - } - _ => {} - } - }; -} - -pub(super) use handle_dictionary; diff --git a/datafusion/physical-expr/src/expressions/in_list/strategy.rs b/datafusion/physical-expr/src/expressions/in_list/strategy.rs index d008217ee19fa..093752cce04f2 100644 --- a/datafusion/physical-expr/src/expressions/in_list/strategy.rs +++ b/datafusion/physical-expr/src/expressions/in_list/strategy.rs @@ -17,31 +17,96 @@ use std::sync::Arc; -use arrow::array::ArrayRef; -use arrow::compute::cast; +use arrow::array::{ArrayRef, AsArray}; +use arrow::compute::take; use arrow::datatypes::DataType; use datafusion_common::Result; use super::array_static_filter::ArrayStaticFilter; +use super::dictionary_filter::DictionaryFilter; +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 { +pub(super) fn instantiate_static_filter( + in_array: ArrayRef, + needle_data_type: &DataType, +) -> Result { let in_array = flatten_dictionary_haystack(in_array)?; - if let Some(filter) = instantiate_primitive_filter(&in_array)? { - return Ok(filter); + let filter = if let Some(filter) = instantiate_fixed_size_binary_filter(&in_array)? { + filter + } else if let Some(filter) = instantiate_primitive_filter(&in_array)? { + filter + } else { + Arc::new(ArrayStaticFilter::try_new(in_array)?) + }; + + // Plain inputs can call the concrete filter directly. Dictionary inputs + // share one adapter across all concrete filter types. + if matches!(needle_data_type, DataType::Dictionary(_, _)) { + Ok(Arc::new(DictionaryFilter::new(filter))) + } else { + Ok(filter) + } +} + +fn flatten_dictionary_haystack(mut in_array: ArrayRef) -> Result { + // Flatten every dictionary layer so the final value type can use a + // specialized filter. + while let Some(dictionary) = in_array.as_any_dictionary_opt() { + in_array = take(dictionary.values().as_ref(), dictionary.keys(), None)?; } - Ok(Arc::new(ArrayStaticFilter::try_new(in_array)?)) + Ok(in_array) } -fn flatten_dictionary_haystack(in_array: ArrayRef) -> Result { - // Flatten dictionary-encoded haystacks to their value type so that - // 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), +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::array::{ + BooleanArray, DictionaryArray, Int8Array, Int16Array, Int32Array, + }; + + use super::*; + + fn nested_dictionary(keys: Int16Array) -> Result { + let values: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])); + // This dictionary represents [1, 3, NULL]. + let inner: ArrayRef = Arc::new(DictionaryArray::try_new( + Int8Array::from(vec![0, 2, 1]), + values, + )?); + Ok(Arc::new(DictionaryArray::try_new(keys, inner)?)) + } + + #[test] + fn nested_dictionary_haystacks_only_include_referenced_values() -> Result<()> { + let needles = Int32Array::from(vec![1, 2, 3]); + + // The null in the inner dictionary is not referenced by the outer one. + let filter = instantiate_static_filter( + nested_dictionary(Int16Array::from(vec![0, 1]))?, + &DataType::Int32, + )?; + assert_eq!(filter.null_count(), 0); + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![true, false, true]) + ); + + // Referencing that same value gives the list normal SQL null semantics. + let filter = instantiate_static_filter( + nested_dictionary(Int16Array::from(vec![0, 2]))?, + &DataType::Int32, + )?; + assert_eq!(filter.null_count(), 1); + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![Some(true), None, None]) + ); + + Ok(()) } }