From ab4b5e0d964dd2b76a68ba515764d186bb476298 Mon Sep 17 00:00:00 2001 From: Geoffrey Claude Date: Wed, 5 Aug 2026 13:58:36 +0200 Subject: [PATCH 1/4] Add FixedSizeBinary IN LIST benchmarks --- .../physical-expr/benches/in_list_strategy.rs | 191 ++++++++++++++---- 1 file changed, 152 insertions(+), 39 deletions(-) diff --git a/datafusion/physical-expr/benches/in_list_strategy.rs b/datafusion/physical-expr/benches/in_list_strategy.rs index 762c2c97a1115..00a7c476f5877 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,19 @@ //! | 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 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; @@ -870,7 +875,7 @@ fn bench_dictionary(c: &mut Criterion) { // NULL HANDLING BENCHMARKS // ============================================================================= // -// Tests representative null-containing inputs across primitive and string cases. +// Null-containing primitive and string cases. fn bench_nulls(c: &mut Criterion) { // ========================================================================= @@ -1014,72 +1019,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); } } } From 1217f5ec322ae31940a7d53235780fa24734eb3e Mon Sep 17 00:00:00 2001 From: Geoffrey Claude Date: Sun, 9 Aug 2026 09:22:46 +0200 Subject: [PATCH 2/4] Optimize IN LIST for fixed-size binary arrays --- .../physical-expr/src/expressions/in_list.rs | 33 ++ .../in_list/fixed_size_binary_filter.rs | 371 ++++++++++++++++++ .../src/expressions/in_list/strategy.rs | 5 + 3 files changed, 409 insertions(+) create mode 100644 datafusion/physical-expr/src/expressions/in_list/fixed_size_binary_filter.rs diff --git a/datafusion/physical-expr/src/expressions/in_list.rs b/datafusion/physical-expr/src/expressions/in_list.rs index 0fb978cd0bafe..154decfd8bb89 100644 --- a/datafusion/physical-expr/src/expressions/in_list.rs +++ b/datafusion/physical-expr/src/expressions/in_list.rs @@ -38,6 +38,7 @@ use datafusion_expr::{ColumnarValue, expr_vec_fmt}; mod array_static_filter; mod branchless_filter; +mod fixed_size_binary_filter; mod primitive_filter; mod result; mod static_filter; @@ -3548,6 +3549,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/fixed_size_binary_filter.rs b/datafusion/physical-expr/src/expressions/in_list/fixed_size_binary_filter.rs new file mode 100644 index 0000000000000..421c02ea1f2d5 --- /dev/null +++ b/datafusion/physical-expr/src/expressions/in_list/fixed_size_binary_filter.rs @@ -0,0 +1,371 @@ +// 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, handle_dictionary}; + +/// Reinterpret fixed-size binary values as same-width primitive values. +/// +/// Arrow buffers are normally sufficiently aligned, making this zero-copy. A +/// valid Arrow array can still be constructed from a sliced, unaligned buffer; +/// in that case, copy each value into aligned primitive storage. +fn reinterpret_as_primitive(array: &FixedSizeBinaryArray) -> Result> +where + T: ArrowPrimitiveType, +{ + let width = size_of::(); + if 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 { + handle_dictionary!(self, v, negated); + + if v.data_type() != &self.data_type { + return Err(exec_datafusion_err!( + "FixedSizeBinary filter: expected {} array, got {}", + self.data_type, + v.data_type() + )); + } + let array = v.as_fixed_size_binary_opt().ok_or_else(|| { + exec_datafusion_err!( + "FixedSizeBinary filter: expected concrete {} array", + self.data_type + ) + })?; + let primitive = reinterpret_as_primitive::(array)?; + self.inner.contains(&primitive, negated) + } +} + +fn instantiate_for_primitive(array: &FixedSizeBinaryArray) -> Result +where + T: ArrowPrimitiveType + Send + Sync + 'static, +{ + let primitive: ArrayRef = Arc::new(reinterpret_as_primitive::(array)?); + let inner = instantiate_primitive_filter(&primitive)?.ok_or_else(|| { + internal_datafusion_err!( + "FixedSizeBinary filter: no primitive filter for {}", + primitive.data_type() + ) + })?; + Ok(Arc::new(FixedSizeBinaryFilter:: { + data_type: array.data_type().clone(), + inner, + _marker: PhantomData, + })) +} + +/// Creates an optimized filter for supported concrete `FixedSizeBinary` arrays. +pub(super) fn instantiate_fixed_size_binary_filter( + in_array: &ArrayRef, +) -> Result> { + let DataType::FixedSizeBinary(width) = in_array.data_type() else { + return Ok(None); + }; + let Some(array) = in_array.as_fixed_size_binary_opt() else { + return Ok(None); + }; + + let filter = match width { + 1 => instantiate_for_primitive::(array)?, + 2 => instantiate_for_primitive::(array)?, + 4 => instantiate_for_primitive::(array)?, + 8 => instantiate_for_primitive::(array)?, + 16 => instantiate_for_primitive::(array)?, + _ => return Ok(None), + }; + Ok(Some(filter)) +} + +#[cfg(test)] +mod tests { + use arrow::array::{DictionaryArray, Int8Array, StringArray}; + use arrow::buffer::{Buffer, MutableBuffer, NullBuffer}; + use arrow::datatypes::Int8Type; + + use super::*; + + fn value(width: i32, index: usize, miss: bool) -> Vec { + let mut value = (index as u128).to_le_bytes()[..width as usize].to_vec(); + let last = value.last_mut().unwrap(); + if miss { + *last |= 0x80; + } else { + *last &= 0x7f; + } + value + } + + fn array(width: i32, values: &[Option>]) -> FixedSizeBinaryArray { + FixedSizeBinaryArray::try_from_sparse_iter_with_size( + values.iter().map(|value| value.as_deref()), + width, + ) + .unwrap() + } + + fn make_filter(width: i32, values: &[Option>]) -> Result { + let in_array: ArrayRef = Arc::new(array(width, values)); + Ok(instantiate_fixed_size_binary_filter(&in_array)?.unwrap()) + } + + #[test] + fn filters_supported_widths_across_strategy_thresholds() -> Result<()> { + for (width, list_len) in [ + (1, 16), + (1, 17), + (2, 8), + (2, 9), + (4, 32), + (4, 33), + (8, 16), + (8, 17), + (16, 4), + (16, 5), + ] { + let mut hit = vec![0x80; width as usize]; + hit[width as usize - 1] = 0xff; + let mut miss = hit.clone(); + miss[width as usize - 1] ^= 1; + + let mut haystack = (0..list_len - 1) + .map(|index| Some(value(width, index, false))) + .collect::>(); + haystack.push(Some(hit.clone())); + let filter = make_filter(width, &haystack)?; + let needles = array(width, &[Some(hit), Some(miss), None]); + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![Some(true), Some(false), None]), + "width={width}, list_len={list_len}" + ); + } + Ok(()) + } + + #[test] + fn handles_slices_nulls_and_not_in() -> Result<()> { + let width = 16; + let parent = array( + width, + &[ + Some(value(width, 0, false)), + Some(value(width, 1, false)), + None, + Some(value(width, 2, false)), + Some(value(width, 3, false)), + Some(value(width, 4, false)), + Some(value(width, 5, false)), + Some(value(width, 6, false)), + ], + ); + // Five non-null values select the hash-set path. + let in_array: ArrayRef = Arc::new(parent.slice(1, 6)); + let filter = instantiate_fixed_size_binary_filter(&in_array)?.unwrap(); + let needles = array( + width, + &[ + Some(value(width, 2, false)), + Some(value(width, 0, false)), + Some(value(width, 6, false)), + Some(value(width, 7, false)), + None, + ], + ); + + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![Some(true), None, None, None, None]) + ); + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![Some(false), None, None, None, None]) + ); + Ok(()) + } + + #[test] + fn handles_dictionary_needles() -> Result<()> { + let filter = make_filter(4, &[Some(value(4, 7, false))])?; + let dictionary_values: ArrayRef = Arc::new(array( + 4, + &[Some(value(4, 7, false)), Some(value(4, 8, false))], + )); + let keys = Int8Array::from(vec![Some(0), Some(1), None]); + let needles = + DictionaryArray::::try_new(keys, dictionary_values).unwrap(); + + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![Some(true), Some(false), None]) + ); + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![Some(false), Some(true), None]) + ); + Ok(()) + } + + #[test] + fn rejects_unsupported_arrays() -> Result<()> { + let filter = make_filter(4, &[Some(value(4, 1, false))])?; + let wrong_width = array(8, &[Some(value(8, 1, false))]); + let error = filter + .contains(&wrong_width, false) + .unwrap_err() + .to_string(); + assert!( + error.contains("expected FixedSizeBinary(4) array, got FixedSizeBinary(8)"), + "{error}" + ); + + let wrong_type = StringArray::from(vec!["one"]); + let error = filter.contains(&wrong_type, false).unwrap_err().to_string(); + assert!( + error.contains("expected FixedSizeBinary(4) array, got Utf8"), + "{error}" + ); + + for width in [0, 3, 5, 15, 17] { + let unsupported: ArrayRef = + Arc::new(FixedSizeBinaryArray::new_null(width, 1)); + assert!( + instantiate_fixed_size_binary_filter(&unsupported)?.is_none(), + "width={width}" + ); + } + + Ok(()) + } + + fn unaligned_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/strategy.rs b/datafusion/physical-expr/src/expressions/in_list/strategy.rs index d008217ee19fa..98bd698507844 100644 --- a/datafusion/physical-expr/src/expressions/in_list/strategy.rs +++ b/datafusion/physical-expr/src/expressions/in_list/strategy.rs @@ -23,12 +23,17 @@ use arrow::datatypes::DataType; use datafusion_common::Result; use super::array_static_filter::ArrayStaticFilter; +use super::fixed_size_binary_filter::instantiate_fixed_size_binary_filter; use super::primitive_filter::instantiate_primitive_filter; use super::static_filter::StaticFilterRef; pub(super) fn instantiate_static_filter(in_array: ArrayRef) -> Result { let in_array = flatten_dictionary_haystack(in_array)?; + if let Some(filter) = instantiate_fixed_size_binary_filter(&in_array)? { + return Ok(filter); + } + if let Some(filter) = instantiate_primitive_filter(&in_array)? { return Ok(filter); } From c396e6ae465a96725363e7715702d140ca97a2ca Mon Sep 17 00:00:00 2001 From: Geoffrey Claude Date: Tue, 25 Aug 2026 13:24:27 +0200 Subject: [PATCH 3/4] Add small-batch IN LIST dictionary benchmarks Cover plain and dictionary Int32 inputs from one row through the standard 8,192-row batch. The cases use both short and longer IN lists so fixed per-batch costs are visible without losing the usual vectorized comparison. --- .../physical-expr/benches/in_list_strategy.rs | 57 ++++++++++++++++++- 1 file changed, 56 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-expr/benches/in_list_strategy.rs b/datafusion/physical-expr/benches/in_list_strategy.rs index 00a7c476f5877..24ca0ae17658e 100644 --- a/datafusion/physical-expr/benches/in_list_strategy.rs +++ b/datafusion/physical-expr/benches/in_list_strategy.rs @@ -45,6 +45,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 | +//! | 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 | @@ -871,6 +872,60 @@ 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 // ============================================================================= @@ -1204,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); From 33a3f38740def9907937357513a77f015bad168f Mon Sep 17 00:00:00 2001 From: Geoffrey Claude Date: Tue, 25 Aug 2026 14:53:29 +0200 Subject: [PATCH 4/4] Handle dictionary IN LIST inputs in one place Wrap the selected filter when the input is dictionary encoded, instead of teaching every filter implementation how to unpack dictionary keys. Plain inputs continue to use their filter directly. Flatten nested dictionary values before choosing a filter, and cover dictionary key types, nested dictionaries, nulls, timestamps, decimals, and scalar inputs. --- .../physical-expr/src/expressions/in_list.rs | 30 +++- .../in_list/array_static_filter.rs | 23 +-- .../expressions/in_list/branchless_filter.rs | 4 +- .../expressions/in_list/dictionary_filter.rs | 161 ++++++++++++++++++ .../in_list/fixed_size_binary_filter.rs | 7 +- .../expressions/in_list/primitive_filter.rs | 24 ++- .../src/expressions/in_list/static_filter.rs | 32 +--- .../src/expressions/in_list/strategy.rs | 90 ++++++++-- 8 files changed, 291 insertions(+), 80 deletions(-) create mode 100644 datafusion/physical-expr/src/expressions/in_list/dictionary_filter.rs diff --git a/datafusion/physical-expr/src/expressions/in_list.rs b/datafusion/physical-expr/src/expressions/in_list.rs index 154decfd8bb89..3c260190b8208 100644 --- a/datafusion/physical-expr/src/expressions/in_list.rs +++ b/datafusion/physical-expr/src/expressions/in_list.rs @@ -38,6 +38,7 @@ use datafusion_expr::{ColumnarValue, expr_vec_fmt}; mod array_static_filter; mod branchless_filter; +mod dictionary_filter; mod fixed_size_binary_filter; mod primitive_filter; mod result; @@ -216,7 +217,7 @@ impl InListExpr { expr, list, negated, - Some(instantiate_static_filter(array)?), + Some(instantiate_static_filter(array, &expr_data_type)?), )) } @@ -243,7 +244,7 @@ impl InListExpr { // Try to create a static filter if all list expressions are constants let static_filter = match try_evaluate_constant_list(&list, schema)? { - Some(in_array) => Some(instantiate_static_filter(in_array)?), + Some(in_array) => Some(instantiate_static_filter(in_array, &expr_data_type)?), None => None, // Non-constant expressions, fall back to dynamic evaluation }; @@ -1265,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 { 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 index 421c02ea1f2d5..ddad5a70994ac 100644 --- 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 @@ -53,7 +53,7 @@ use arrow::datatypes::{ use datafusion_common::{Result, exec_datafusion_err, internal_datafusion_err}; use super::primitive_filter::instantiate_primitive_filter; -use super::static_filter::{StaticFilter, StaticFilterRef, handle_dictionary}; +use super::static_filter::{StaticFilter, StaticFilterRef}; /// Reinterpret fixed-size binary values as same-width primitive values. /// @@ -99,8 +99,6 @@ where } 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 {}", @@ -165,6 +163,7 @@ mod tests { 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 { @@ -268,7 +267,7 @@ mod tests { #[test] fn handles_dictionary_needles() -> Result<()> { - let filter = make_filter(4, &[Some(value(4, 7, false))])?; + 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))], 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 98bd698507844..093752cce04f2 100644 --- a/datafusion/physical-expr/src/expressions/in_list/strategy.rs +++ b/datafusion/physical-expr/src/expressions/in_list/strategy.rs @@ -17,36 +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_fixed_size_binary_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) } +} - if let Some(filter) = instantiate_primitive_filter(&in_array)? { - return 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(()) } }