From 8f22528b52c9dab86bec152671cfa5e71552993d Mon Sep 17 00:00:00 2001 From: Ryu <177669517+ryux1@users.noreply.github.com> Date: Mon, 17 Aug 2026 19:29:25 +0000 Subject: [PATCH 01/19] perf: Reduce record batch memory accounting overhead (#24319) ## Which issue does this PR close? - Closes #24310. ## Rationale for this change Record batch memory accounting runs on hot execution paths. The current implementation materializes ArrayData for every array and allocates a hash set even for small batches, adding measurable overhead to queries that frequently update memory reservations. ## What changes are included in this PR? - Traverse Arrow arrays directly and recursively count their backing buffers without materializing ArrayData. - Track the first 16 buffer identities inline, then promote to a hash set for wider batches or counters spanning many batches. - Preserve shared-buffer deduplication and full buffer-capacity accounting semantics. - Add parity coverage against the previous ArrayData traversal for primitive, binary/view, list/view, fixed-size, struct, union, dictionary, map, and all legal run-end index layouts. - Add a focused Criterion benchmark across column counts, row counts, and primitive/list/struct layouts. Criterion point estimates from `cargo bench -p datafusion-common --bench record_batch_memory`, measured sequentially on upstream `main` and this PR on the same machine: ### Column count (8,192 rows, Int64) | Columns | main | this PR | Speedup | |---:|---:|---:|---:| | 1 | 81.999 ns | 10.645 ns | 7.70x | | 4 | 308.11 ns | 38.669 ns | 7.97x | | 16 | 1.2543 us | 173.55 ns | 7.23x | | 64 | 4.7457 us | 1.3914 us | 3.41x | ### Row count (4 Int64 columns) | Rows | main | this PR | Speedup | |---:|---:|---:|---:| | 1 | 299.20 ns | 39.939 ns | 7.49x | | 128 | 302.99 ns | 38.914 ns | 7.79x | | 8,192 | 307.64 ns | 39.740 ns | 7.74x | | 65,536 | 310.61 ns | 38.994 ns | 7.97x | ### Array layout (4 columns, 8,192 rows) | Layout | main | this PR | Speedup | |---|---:|---:|---:| | Primitive Int64 | 311.01 ns | 40.140 ns | 7.75x | | List of Int64 | 778.29 ns | 82.865 ns | 9.39x | | Struct of two Int64 fields | 1.0072 us | 113.70 ns | 8.86x | Each list row contains two Int64 values. Each struct column contains two Int64 child fields. ## Are these changes tested? Yes. - `cargo fmt --all -- --check` - `cargo clippy --all-targets --all-features -- -D warnings` - `cargo test -p datafusion-common utils::memory --lib` (10 passed) - `RUST_BACKTRACE=1 cargo test --profile ci --exclude datafusion-examples --exclude datafusion-benchmarks --exclude datafusion-cli --workspace --lib --tests --bins --features avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption` - Focused Criterion comparison shown above ## Are there any user-facing changes? No API or behavior changes. This reduces CPU and allocation overhead in record batch memory accounting. AI assistance: OpenAI Codex assisted with implementation and test execution. I reviewed the change and its behavior end to end. Backport provenance: - Upstream PR: https://github.com/apache/datafusion/pull/24319 - Upstream commit: c7be84a18c77b974fafc26f2cf7265526af4d18d - The upstream diff is unchanged, including its own tests and benchmark. - No fork adaptation is included in this commit. Local Arrow 59 validity-buffer regressions are deliberately deferred to the final additional-test section. - The body above is the original upstream description; its benchmark estimates and verification statements describe upstream, not this fork's measurements. Signed-off-by: discord9 --- datafusion/common/Cargo.toml | 4 + .../common/benches/record_batch_memory.rs | 190 ++++++++ datafusion/common/src/utils/memory.rs | 415 ++++++++++++++++-- 3 files changed, 577 insertions(+), 32 deletions(-) create mode 100644 datafusion/common/benches/record_batch_memory.rs diff --git a/datafusion/common/Cargo.toml b/datafusion/common/Cargo.toml index 1eb23089a4021..9ee199fe82f28 100644 --- a/datafusion/common/Cargo.toml +++ b/datafusion/common/Cargo.toml @@ -64,6 +64,10 @@ name = "scalar_to_array" harness = false name = "stats_merge" +[[bench]] +harness = false +name = "record_batch_memory" + [dependencies] arrow = { workspace = true } arrow-ipc = { workspace = true } diff --git a/datafusion/common/benches/record_batch_memory.rs b/datafusion/common/benches/record_batch_memory.rs new file mode 100644 index 0000000000000..2479d6ac987cb --- /dev/null +++ b/datafusion/common/benches/record_batch_memory.rs @@ -0,0 +1,190 @@ +// 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. + +//! Measures the CPU overhead of accounting for the backing buffers retained by +//! [`RecordBatch`]es. Batch construction is intentionally outside the timed +//! region so the benchmarks isolate buffer traversal and identity deduplication. + +use std::hint::black_box; +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int64Array, ListArray, StructArray}; +use arrow::datatypes::{DataType, Field, Int64Type, Schema}; +use arrow::record_batch::RecordBatch; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_common::utils::memory::{ + RecordBatchMemoryCounter, get_record_batch_memory_size, +}; + +fn make_batch(columns: Vec) -> RecordBatch { + let fields = columns + .iter() + .enumerate() + .map(|(index, column)| { + Field::new(format!("col_{index}"), column.data_type().clone(), false) + }) + .collect::>(); + + RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap() +} + +fn make_primitive_batch(num_rows: usize, num_columns: usize) -> RecordBatch { + let columns = (0..num_columns) + .map(|index| { + Arc::new(Int64Array::from_iter_values( + (0..num_rows).map(|value| value as i64 + index as i64), + )) as ArrayRef + }) + .collect::>(); + + make_batch(columns) +} + +fn make_list_batch(num_rows: usize, num_columns: usize) -> RecordBatch { + let columns = (0..num_columns) + .map(|column| { + Arc::new(ListArray::from_iter_primitive::( + (0..num_rows).map(|row| { + let value = row as i64 + column as i64; + Some(vec![Some(value), Some(value + 1)]) + }), + )) as ArrayRef + }) + .collect::>(); + + make_batch(columns) +} + +fn make_struct_batch(num_rows: usize, num_columns: usize) -> RecordBatch { + let columns = (0..num_columns) + .map(|column| { + let left = Arc::new(Int64Array::from_iter_values( + (0..num_rows).map(|row| row as i64 + column as i64), + )) as ArrayRef; + let right = Arc::new(Int64Array::from_iter_values( + (0..num_rows).map(|row| row as i64 - column as i64), + )) as ArrayRef; + + Arc::new(StructArray::from(vec![ + (Arc::new(Field::new("left", DataType::Int64, false)), left), + (Arc::new(Field::new("right", DataType::Int64, false)), right), + ])) as ArrayRef + }) + .collect::>(); + + make_batch(columns) +} + +fn benchmark_column_count(c: &mut Criterion) { + let mut group = c.benchmark_group("record_batch_memory_size/column_count"); + + // Each primitive column contributes a distinct backing buffer, exercising + // both the inline buffer-ID path and hash-set promotion. + for num_columns in [1, 4, 16, 64] { + let batch = make_primitive_batch(8192, num_columns); + group.bench_with_input( + BenchmarkId::from_parameter(num_columns), + &batch, + |bencher, batch| { + bencher.iter(|| get_record_batch_memory_size(black_box(batch))); + }, + ); + } + + group.finish(); +} + +fn benchmark_row_count(c: &mut Criterion) { + let mut group = c.benchmark_group("record_batch_memory_size/row_count"); + + // Buffer traversal should depend on the number of buffers, not the number + // of values stored in each buffer. + for num_rows in [1, 128, 8192, 65_536] { + let batch = make_primitive_batch(num_rows, 4); + group.bench_with_input( + BenchmarkId::from_parameter(num_rows), + &batch, + |bencher, batch| { + bencher.iter(|| get_record_batch_memory_size(black_box(batch))); + }, + ); + } + + group.finish(); +} + +fn benchmark_array_layout(c: &mut Criterion) { + let mut group = c.benchmark_group("record_batch_memory_size/array_layout"); + + // Compare direct primitive-buffer accounting with recursive traversal of + // representative nested layouts. + for (name, batch) in [ + ("primitive", make_primitive_batch(8192, 4)), + ("list", make_list_batch(8192, 4)), + ("struct", make_struct_batch(8192, 4)), + ] { + group.bench_with_input( + BenchmarkId::from_parameter(name), + &batch, + |bencher, batch| { + bencher.iter(|| get_record_batch_memory_size(black_box(batch))); + }, + ); + } + + group.finish(); +} + +fn benchmark_shared_slices(c: &mut Criterion) { + let mut group = c.benchmark_group("record_batch_memory_size/shared_slices"); + + // Model the hash-join build-side workload: one counter is reused across a + // sequence of zero-copy batch slices that retain the same backing buffers. + // Slicing happens outside the timed region; the benchmark measures repeated + // identity lookups and the one-time accounting of each shared buffer. + for num_columns in [4, 16, 64] { + let batch = make_primitive_batch(8192, num_columns); + let slices = (0..32) + .map(|index| batch.slice(index * 256, 256)) + .collect::>(); + + group.bench_with_input( + BenchmarkId::from_parameter(num_columns), + &slices, + |bencher, slices| { + bencher.iter(|| { + let mut counter = RecordBatchMemoryCounter::new(); + for batch in black_box(slices) { + black_box(counter.count_batch(black_box(batch))); + } + black_box(counter.memory_usage()) + }); + }, + ); + } + + group.finish(); +} + +criterion_group!( + benches, + benchmark_column_count, + benchmark_row_count, + benchmark_array_layout, + benchmark_shared_slices +); +criterion_main!(benches); diff --git a/datafusion/common/src/utils/memory.rs b/datafusion/common/src/utils/memory.rs index 21c084119e120..fd405e06a262e 100644 --- a/datafusion/common/src/utils/memory.rs +++ b/datafusion/common/src/utils/memory.rs @@ -19,11 +19,24 @@ use crate::error::_exec_datafusion_err; use crate::{HashSet, Result}; -use arrow::array::ArrayData; +use arrow::array::types::{ByteArrayType, ByteViewType, RunEndIndexType}; +use arrow::array::{ + Array, AsArray, GenericByteArray, GenericByteViewArray, GenericListArray, + GenericListViewArray, RunArray, +}; +use arrow::buffer::Buffer; +use arrow::datatypes::DataType; +use arrow::downcast_primitive_array; use arrow::record_batch::RecordBatch; use std::mem::size_of; use std::num::NonZero; +/// Maximum number of distinct buffer IDs retained inline before promotion to +/// a [`HashSet`]. Sixteen keeps small buffer sets allocation-free while +/// limiting linear lookup and inline storage to 16 pointer-sized entries. +/// This is a performance heuristic, not a semantic limit. +const INLINE_BUFFER_IDS: usize = 16; + /// Estimates the memory size required for a hash table prior to allocation. /// /// # Parameters @@ -151,7 +164,7 @@ pub fn get_record_batch_memory_size(batch: &RecordBatch) -> usize { pub struct RecordBatchMemoryCounter { /// Start addresses of `Buffer`s that have already been counted (instead of /// actual used data region's pointer represented by current `Array`) - counted_buffers: HashSet>, + counted_buffers: BufferIdSet, /// Total memory of all unique buffers counted so far memory_usage: usize, } @@ -164,49 +177,229 @@ impl RecordBatchMemoryCounter { /// Count `batch`, returning the memory used by its buffers that have not /// been counted before. pub fn count_batch(&mut self, batch: &RecordBatch) -> usize { - let mut total_size = 0; + let previous_memory_usage = self.memory_usage; for array in batch.columns() { - let array_data = array.to_data(); - count_array_data_memory_size( - &array_data, - &mut self.counted_buffers, - &mut total_size, - ); + self.count_array_memory_size(array.as_ref()); } - self.memory_usage += total_size; - total_size + self.memory_usage - previous_memory_usage } /// Total memory of the unique buffers of all batches counted so far. pub fn memory_usage(&self) -> usize { self.memory_usage } -} -/// Count the memory usage of `array_data` and its children recursively. -fn count_array_data_memory_size( - array_data: &ArrayData, - counted_buffers: &mut HashSet>, - total_size: &mut usize, -) { - // Count memory usage for `array_data` - for buffer in array_data.buffers() { - if counted_buffers.insert(buffer.data_ptr().addr()) { - *total_size += buffer.capacity(); - } // Otherwise the buffer's memory is already counted + fn count_buffer_memory_size(&mut self, buffer: &Buffer) { + if self.counted_buffers.insert(buffer.data_ptr().addr()) { + self.memory_usage += buffer.capacity(); + } } - if let Some(null_buffer) = array_data.nulls() - && counted_buffers.insert(null_buffer.inner().inner().data_ptr().addr()) - { - *total_size += null_buffer.inner().inner().capacity(); + /// Count the memory usage of `array` and its children recursively. + fn count_array_memory_size(&mut self, array: &dyn Array) { + if let Some(nulls) = array.nulls() { + self.count_buffer_memory_size(nulls.buffer()); + } + + downcast_primitive_array! { + array => self.count_buffer_memory_size(array.values().inner()), + DataType::Null => {} + DataType::Boolean => { + self.count_buffer_memory_size(array.as_boolean().values().inner()); + } + DataType::Binary => { + self.count_byte_array_memory_size(array.as_binary::()); + } + DataType::LargeBinary => { + self.count_byte_array_memory_size(array.as_binary::()); + } + DataType::Utf8 => { + self.count_byte_array_memory_size(array.as_string::()); + } + DataType::LargeUtf8 => { + self.count_byte_array_memory_size(array.as_string::()); + } + DataType::BinaryView => { + self.count_byte_view_array_memory_size(array.as_binary_view()); + } + DataType::Utf8View => { + self.count_byte_view_array_memory_size(array.as_string_view()); + } + DataType::FixedSizeBinary(_) => { + self.count_buffer_memory_size(array.as_fixed_size_binary().values()); + } + DataType::List(_) => { + self.count_list_array_memory_size(array.as_list::()); + } + DataType::LargeList(_) => { + self.count_list_array_memory_size(array.as_list::()); + } + DataType::ListView(_) => { + self.count_list_view_array_memory_size(array.as_list_view::()); + } + DataType::LargeListView(_) => { + self.count_list_view_array_memory_size(array.as_list_view::()); + } + DataType::FixedSizeList(_, _) => { + self.count_array_memory_size( + array.as_fixed_size_list().values().as_ref(), + ); + } + DataType::Struct(_) => { + for child in array.as_struct().columns() { + self.count_array_memory_size(child.as_ref()); + } + } + DataType::Union(_, _) => { + let array = array.as_union(); + self.count_buffer_memory_size(array.type_ids().inner()); + if let Some(offsets) = array.offsets() { + self.count_buffer_memory_size(offsets.inner()); + } + for (type_id, _) in array.fields().iter() { + self.count_array_memory_size(array.child(type_id).as_ref()); + } + } + DataType::Dictionary(_, _) => { + let array = array.as_any_dictionary(); + self.count_array_memory_size(array.keys()); + self.count_array_memory_size(array.values().as_ref()); + } + DataType::Map(_, _) => { + let array = array.as_map(); + self.count_buffer_memory_size(array.offsets().inner().inner()); + self.count_array_memory_size(array.entries()); + } + DataType::RunEndEncoded(run_ends, _) => match run_ends.data_type() { + DataType::Int16 => { + self.count_run_array_memory_size::( + array, + ); + } + DataType::Int32 => { + self.count_run_array_memory_size::( + array, + ); + } + DataType::Int64 => { + self.count_run_array_memory_size::( + array, + ); + } + // Arrow only permits Int16, Int32, and Int64 run-end indexes. A + // custom Array implementation may still expose malformed data; + // retain correct accounting for it without panicking. + _ => self.count_array_data_memory_size(&array.to_data()), + }, + // All currently supported non-primitive layouts are handled above. + // The Arrow macro requires a final arm for primitive variants that + // its nested dispatch has already consumed. Keep a safe generic + // fallback for custom or future Array implementations. + _ => self.count_array_data_memory_size(&array.to_data()), + } } - // Count all children `ArrayData` recursively - for child in array_data.child_data() { - count_array_data_memory_size(child, counted_buffers, total_size); + fn count_byte_array_memory_size( + &mut self, + array: &GenericByteArray, + ) { + self.count_buffer_memory_size(array.offsets().inner().inner()); + self.count_buffer_memory_size(array.values()); + } + + fn count_byte_view_array_memory_size( + &mut self, + array: &GenericByteViewArray, + ) { + self.count_buffer_memory_size(array.views().inner()); + for buffer in array.data_buffers() { + self.count_buffer_memory_size(buffer); + } + } + + fn count_list_array_memory_size( + &mut self, + array: &GenericListArray, + ) { + self.count_buffer_memory_size(array.offsets().inner().inner()); + self.count_array_memory_size(array.values().as_ref()); + } + + fn count_list_view_array_memory_size( + &mut self, + array: &GenericListViewArray, + ) { + self.count_buffer_memory_size(array.offsets().inner()); + self.count_buffer_memory_size(array.sizes().inner()); + self.count_array_memory_size(array.values().as_ref()); + } + + fn count_run_array_memory_size(&mut self, array: &dyn Array) { + if let Some(array) = array.as_any().downcast_ref::>() { + self.count_buffer_memory_size(array.run_ends().inner().inner()); + self.count_array_memory_size(array.values().as_ref()); + } else { + // The DataType and concrete array implementation disagree. Use the + // generic representation rather than panic while accounting memory. + self.count_array_data_memory_size(&array.to_data()); + } + } + + fn count_array_data_memory_size(&mut self, array_data: &arrow::array::ArrayData) { + for buffer in array_data.buffers() { + self.count_buffer_memory_size(buffer); + } + if let Some(nulls) = array_data.nulls() { + self.count_buffer_memory_size(nulls.buffer()); + } + for child in array_data.child_data() { + self.count_array_data_memory_size(child); + } + } +} + +/// Tracks a small number of buffers inline, avoiding a heap allocation for +/// typical batches, and promotes to a hash set when more buffers are seen. +#[derive(Debug)] +struct BufferIdSet { + inline: [Option>; INLINE_BUFFER_IDS], + len: usize, + overflow: Option>>, +} + +impl Default for BufferIdSet { + fn default() -> Self { + Self { + inline: [None; INLINE_BUFFER_IDS], + len: 0, + overflow: None, + } + } +} + +impl BufferIdSet { + fn insert(&mut self, buffer_id: NonZero) -> bool { + if let Some(overflow) = &mut self.overflow { + return overflow.insert(buffer_id); + } + + if self.inline[..self.len].contains(&Some(buffer_id)) { + return false; + } + + if self.len < INLINE_BUFFER_IDS { + self.inline[self.len] = Some(buffer_id); + self.len += 1; + return true; + } + + let mut overflow = HashSet::with_capacity(INLINE_BUFFER_IDS + 1); + overflow.extend(self.inline.iter().flatten().copied()); + let inserted = overflow.insert(buffer_id); + self.overflow = Some(overflow); + inserted } } @@ -247,10 +440,49 @@ mod tests { #[cfg(test)] mod record_batch_tests { use super::*; - use arrow::array::{Float64Array, Int32Array, ListArray}; - use arrow::datatypes::{DataType, Field, Int32Type, Schema}; + use arrow::array::{ + ArrayData, ArrayRef, BinaryViewArray, Float64Array, Int16Array, Int32Array, + Int64Array, LargeListViewArray, ListArray, ListViewArray, RunArray, StringArray, + StringViewArray, new_null_array, + }; + use arrow::datatypes::{ + DataType, Field, Int16Type, Int32Type, Int64Type, Schema, UnionFields, UnionMode, + }; use std::sync::Arc; + fn array_data_memory_size(array: &dyn Array) -> usize { + fn count( + array_data: &ArrayData, + counted_buffers: &mut HashSet>, + total_size: &mut usize, + ) { + for buffer in array_data.buffers() { + if counted_buffers.insert(buffer.data_ptr().addr()) { + *total_size += buffer.capacity(); + } + } + if let Some(nulls) = array_data.nulls() { + let buffer = nulls.inner().inner(); + if counted_buffers.insert(buffer.data_ptr().addr()) { + *total_size += buffer.capacity(); + } + } + for child in array_data.child_data() { + count(child, counted_buffers, total_size); + } + } + + let mut total_size = 0; + count(&array.to_data(), &mut HashSet::default(), &mut total_size); + total_size + } + + fn assert_array_memory_size_matches(array: &dyn Array) { + let mut counter = RecordBatchMemoryCounter::new(); + counter.count_array_memory_size(array); + assert_eq!(counter.memory_usage(), array_data_memory_size(array)); + } + #[test] fn test_get_record_batch_memory_size() { let schema = Arc::new(Schema::new(vec![ @@ -359,6 +591,125 @@ mod record_batch_tests { assert_eq!(counter.memory_usage(), get_record_batch_memory_size(&batch)); } + #[test] + fn test_record_batch_memory_counter_promotes_buffer_set() { + let fields = (0..=INLINE_BUFFER_IDS) + .map(|index| Field::new(format!("col_{index}"), DataType::Int32, false)) + .collect::>(); + let columns = (0..=INLINE_BUFFER_IDS) + .map(|value| Arc::new(Int32Array::from(vec![value as i32])) as _) + .collect::>(); + let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap(); + + let mut counter = RecordBatchMemoryCounter::new(); + assert_eq!( + counter.count_batch(&batch), + (INLINE_BUFFER_IDS + 1) * size_of::() + ); + assert!(counter.counted_buffers.overflow.is_some()); + assert_eq!(counter.count_batch(&batch), 0); + } + + #[test] + fn test_array_memory_size_matches_array_data_layouts() { + let list_field = Arc::new(Field::new_list_field(DataType::Int32, true)); + let struct_fields = vec![Field::new("value", DataType::Int32, true)].into(); + let union_fields = UnionFields::try_new( + vec![0], + vec![Field::new("value", DataType::Int32, true)], + ) + .unwrap(); + let map_entries = Arc::new(Field::new( + "entries", + DataType::Struct( + vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", DataType::Int32, true), + ] + .into(), + ), + false, + )); + let data_types = vec![ + DataType::Boolean, + DataType::Int32, + DataType::Binary, + DataType::LargeBinary, + DataType::FixedSizeBinary(4), + DataType::BinaryView, + DataType::Utf8, + DataType::LargeUtf8, + DataType::Utf8View, + DataType::List(Arc::clone(&list_field)), + DataType::LargeList(Arc::clone(&list_field)), + DataType::ListView(Arc::clone(&list_field)), + DataType::LargeListView(Arc::clone(&list_field)), + DataType::FixedSizeList(Arc::clone(&list_field), 2), + DataType::Struct(struct_fields), + DataType::Union(union_fields, UnionMode::Dense), + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + DataType::Map(map_entries, false), + ]; + + for data_type in data_types { + let array = new_null_array(&data_type, 3); + assert_array_memory_size_matches(array.as_ref()); + } + + // Exercise the view-specific buffers with concrete, non-empty values. + let view_arrays = [ + Arc::new(BinaryViewArray::from_iter_values([ + b"short".as_slice(), + b"a payload longer than twelve bytes".as_slice(), + ])) as ArrayRef, + Arc::new(StringViewArray::from_iter_values([ + "short", + "a payload longer than twelve bytes", + ])) as ArrayRef, + Arc::new(ListViewArray::from_iter_primitive::([ + Some(vec![Some(1), Some(2)]), + None, + Some(vec![Some(3)]), + ])) as ArrayRef, + Arc::new(LargeListViewArray::from_iter_primitive::( + [Some(vec![Some(1), Some(2)]), None, Some(vec![Some(3)])], + )) as ArrayRef, + ]; + + for array in view_arrays { + assert_array_memory_size_matches(array.as_ref()); + } + + let run_values = StringArray::from(vec!["alpha", "beta"]); + let run_arrays = [ + Arc::new( + RunArray::::try_new( + &Int16Array::from(vec![2_i16, 5]), + &run_values, + ) + .unwrap(), + ) as ArrayRef, + Arc::new( + RunArray::::try_new( + &Int32Array::from(vec![2_i32, 5]), + &run_values, + ) + .unwrap(), + ) as ArrayRef, + Arc::new( + RunArray::::try_new( + &Int64Array::from(vec![2_i64, 5]), + &run_values, + ) + .unwrap(), + ) as ArrayRef, + ]; + + for array in run_arrays { + assert_array_memory_size_matches(array.as_ref()); + } + } + #[test] fn test_get_record_batch_memory_size_nested_array() { let schema = Arc::new(Schema::new(vec![ From 12d59612064c988ee2ecbf0fa42f6c2a13b0f091 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sun, 20 Sep 2026 10:12:19 +0000 Subject: [PATCH 02/19] fix: account for hash-join row-index storage (#25508) ## Which issue does this PR close? Addresses the shared accounting bug identified while reviewing [#25491](https://github.com/apache/datafusion/pull/25491#discussion_r4052388920). Split out so it can be reviewed and backported independently of prepared-build reuse. ## Rationale for this change Hash joins reserve memory for their hash buckets, but omit the row-index vector used to link build rows with the same key. As a result, a join can allocate beyond its admitted budget and underreport its build memory. For example, 100 million build rows need an additional 400 MB for this vector that the memory pool currently does not see. ## What changes are included in this PR? The shared hash-map constructor now reserves the bucket estimate and row-index storage together before allocating either, and includes both in the build-memory metric. This covers ordinary CollectLeft and partitioned joins, using four bytes per row for the u32 representation and eight for u64. Perfect-hash array maps keep their existing accounting. ## What is the testing strategy for this PR? A regression gives the constructor one byte less than the complete budget and checks that it refuses the allocation without charging the pool or metric. It also verifies success at the exact limit and release on drop. The test fails against the original constructor and passes with this fix. All 533 hash-join tests, full workspace Clippy with all targets/features and warnings denied, formatting, and whitespace checks pass on this branch. The combined prepared-build branch also passed the extended workspace run with this constructor fix (11,995 Rust tests and 521 SQL test files; eight tests ignored). Local checks used Rust 1.98.1 and temporary source overrides for the requested public releases of Arrow 60.0.0, object_store 0.14.2, sqlparser 0.63.0, and compression-codecs 0.4.42 because the registry mirror lacked them. No dependency or lockfile changes are included. Ordinary-path benchmark comparisons for the combined branch, including the accounting fix, are reported in #25491; they are not an isolated performance claim for this fix. ## Are there any user-facing changes? Hash joins using these maps account for their row-index storage in memory limits and build-memory metrics. A budget that previously succeeded only because this allocation was omitted can now be rejected before allocation. Backport provenance: - Upstream PR: https://github.com/apache/datafusion/pull/25508 - Upstream commit: 43f567a976cac1c16c341b0ae90e98b2b715835d - All upstream additions and deletions are retained, including its regression test. Only patch anchors and indentation are relocated: the fork still has inline map construction instead of upstream's new_join_hashmap helper. - No helper extraction or fork interface adaptation is included here. Those changes follow in the immediately next commit. - This original-only intermediate state cannot build test targets: the upstream test references that absent helper. Checks apply to the following adapted state, not retroactively to this commit. - The description above records upstream's verification, not new fork results. Signed-off-by: discord9 --- .../physical-plan/src/joins/hash_join/exec.rs | 50 ++++++++++++++++++- 1 file changed, 48 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 08d209003ad91..66d81d63cf178 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -2326,13 +2326,27 @@ async fn collect_left_input( // Arc is used instead of Box to allow sharing with SharedBuildAccumulator for hash map pushdown let mut hashmap: Box = if num_rows > u32::MAX as usize { let estimated_hashtable_size = - estimate_memory_size::<(u64, u64)>(num_rows, fixed_size_u64)?; + estimate_memory_size::<(u64, u64)>(num_rows, fixed_size_u64)? + // Each build row also owns an index in the duplicate-key chain. + .checked_add(num_rows * size_of::()) + .ok_or_else(|| { + datafusion_common::exec_datafusion_err!( + "Hash join table size overflow" + ) + })?; reservation.try_grow(estimated_hashtable_size)?; metrics.build_mem_used.add(estimated_hashtable_size); Box::new(JoinHashMapU64::with_capacity(num_rows)) } else { let estimated_hashtable_size = - estimate_memory_size::<(u32, u64)>(num_rows, fixed_size_u32)?; + estimate_memory_size::<(u32, u64)>(num_rows, fixed_size_u32)? + // Each build row also owns an index in the duplicate-key chain. + .checked_add(num_rows * size_of::()) + .ok_or_else(|| { + datafusion_common::exec_datafusion_err!( + "Hash join table size overflow" + ) + })?; reservation.try_grow(estimated_hashtable_size)?; metrics.build_mem_used.add(estimated_hashtable_size); Box::new(JoinHashMapU32::with_capacity(num_rows)) @@ -2434,6 +2448,38 @@ async fn collect_left_input( mod tests { use super::*; + #[test] + fn hash_map_admits_row_indices_before_allocation() -> Result<()> { + use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryPool}; + + let rows = 1024; + let buckets = + estimate_memory_size::<(u32, u64)>(rows, size_of::())?; + let bytes = buckets + rows * size_of::(); + for (limit, succeeds) in [(bytes - 1, false), (bytes, true)] { + let pool: Arc = Arc::new(GreedyMemoryPool::new(limit)); + let mut reservation = MemoryConsumer::new("row indices").register(&pool); + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + let result = new_join_hashmap(rows, &mut reservation, &metrics); + if succeeds { + let map = result?; + assert_eq!(reservation.size(), bytes); + assert_eq!(metrics.build_mem_used.value(), bytes); + drop(map); + } else { + assert!(matches!( + result, + Err(datafusion_common::DataFusionError::ResourcesExhausted(_)) + )); + assert_eq!(reservation.size(), 0); + assert_eq!(metrics.build_mem_used.value(), 0); + } + drop(reservation); + assert_eq!(pool.reserved(), 0); + } + Ok(()) + } + fn assert_phj_used(metrics: &MetricsSet, use_phj: bool) { if use_phj { assert!( From 3e4ce27dcec05957c3f6df329240777aa0624e40 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 6 Nov 2026 21:14:15 +0800 Subject: [PATCH 03/19] fix: adapt #25508 hash-map construction to the fork The fork constructs the join map inline, while upstream #25508 updates and regression-tests a new_join_hashmap helper. Extract that existing constructor and wire its caller, retaining the preceding upstream allocation/accounting logic unchanged. This supplies the missing interface without mixing it into the upstream backport. Upstream PR: https://github.com/apache/datafusion/pull/25508 Upstream commit: 43f567a976cac1c16c341b0ae90e98b2b715835d This is a local fork adaptation, not an upstream-authored change. Verification of this adapted state: cargo fmt --all -- --check; full all-target/all-feature Clippy with -D warnings; datafusion-physical-plan nextest (1707 passed). The preceding original-only intermediate test target was not buildable and is not retrospectively claimed to have passed. Signed-off-by: discord9 --- .../physical-plan/src/joins/hash_join/exec.rs | 68 +++++++++++-------- 1 file changed, 38 insertions(+), 30 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 66d81d63cf178..cc09975a39559 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -2198,6 +2198,43 @@ fn should_collect_min_max_for_perfect_hash( Ok(ArrayMap::is_supported_type(&data_type)) } +/// Reserves the hash buckets and duplicate-key row indices before allocation. +fn new_join_hashmap( + num_rows: usize, + reservation: &mut MemoryReservation, + metrics: &BuildProbeJoinMetrics, +) -> Result> { + let fixed_size_u32 = size_of::(); + let fixed_size_u64 = size_of::(); + if num_rows > u32::MAX as usize { + let estimated_hashtable_size = + estimate_memory_size::<(u64, u64)>(num_rows, fixed_size_u64)? + // Each build row also owns an index in the duplicate-key chain. + .checked_add(num_rows * size_of::()) + .ok_or_else(|| { + datafusion_common::exec_datafusion_err!( + "Hash join table size overflow" + ) + })?; + reservation.try_grow(estimated_hashtable_size)?; + metrics.build_mem_used.add(estimated_hashtable_size); + Ok(Box::new(JoinHashMapU64::with_capacity(num_rows))) + } else { + let estimated_hashtable_size = + estimate_memory_size::<(u32, u64)>(num_rows, fixed_size_u32)? + // Each build row also owns an index in the duplicate-key chain. + .checked_add(num_rows * size_of::()) + .ok_or_else(|| { + datafusion_common::exec_datafusion_err!( + "Hash join table size overflow" + ) + })?; + reservation.try_grow(estimated_hashtable_size)?; + metrics.build_mem_used.add(estimated_hashtable_size); + Ok(Box::new(JoinHashMapU32::with_capacity(num_rows))) + } +} + /// Collects all batches from the left (build) side stream and creates a hash map for joining. /// /// This function is responsible for: @@ -2318,39 +2355,10 @@ async fn collect_left_input( } else { // Estimation of memory size, required for hashtable, prior to allocation. // Final result can be verified using `RawTable.allocation_info()` - let fixed_size_u32 = size_of::(); - let fixed_size_u64 = size_of::(); - // Use `u32` indices for the JoinHashMap when num_rows ≤ u32::MAX, otherwise use the // `u64` indice variant // Arc is used instead of Box to allow sharing with SharedBuildAccumulator for hash map pushdown - let mut hashmap: Box = if num_rows > u32::MAX as usize { - let estimated_hashtable_size = - estimate_memory_size::<(u64, u64)>(num_rows, fixed_size_u64)? - // Each build row also owns an index in the duplicate-key chain. - .checked_add(num_rows * size_of::()) - .ok_or_else(|| { - datafusion_common::exec_datafusion_err!( - "Hash join table size overflow" - ) - })?; - reservation.try_grow(estimated_hashtable_size)?; - metrics.build_mem_used.add(estimated_hashtable_size); - Box::new(JoinHashMapU64::with_capacity(num_rows)) - } else { - let estimated_hashtable_size = - estimate_memory_size::<(u32, u64)>(num_rows, fixed_size_u32)? - // Each build row also owns an index in the duplicate-key chain. - .checked_add(num_rows * size_of::()) - .ok_or_else(|| { - datafusion_common::exec_datafusion_err!( - "Hash join table size overflow" - ) - })?; - reservation.try_grow(estimated_hashtable_size)?; - metrics.build_mem_used.add(estimated_hashtable_size); - Box::new(JoinHashMapU32::with_capacity(num_rows)) - }; + let mut hashmap = new_join_hashmap(num_rows, &mut reservation, &metrics)?; let mut hashes_buffer = Vec::new(); let mut offset = 0; From d658e52ce8d9239a8cf3468677f6eb033f3159f9 Mon Sep 17 00:00:00 2001 From: Jay Zhan Date: Thu, 24 Sep 2026 12:41:14 +0000 Subject: [PATCH 04/19] fix: reserve the concatenated build batch and computed join keys in HashJoinExec (#25496) ## Which issue does this PR close? - Closes #25269. ## Rationale for this change A hash join reserves every build-side batch as it arrives, then copies all of them into a single batch without reserving the copy, while the input batches stay alive until the build finishes. Near the memory limit the process therefore holds about twice the build side while the memory pool sees it once, so instead of a `ResourcesExhausted` error the process can be killed by the OS. It also means nothing that reacts to pool pressure can rely on what a hash join reports. Measured on a 1.25 GB build side (TPC-H SF1 `lineitem`, 16 columns, as the build side): the pool was charged ~1.25 GB while peak RSS was ~2.46 GB. ## What changes are included in this PR? - `concat_build_batches`: reserves the copy before `concat_batches`, drops the input batches, then trims the reservation to what the concatenated batch retains (`get_record_batch_memory_size`), growing instead if the estimate was low. `build_mem_used` follows every grow and shrink. - Nothing extra is reserved when there is nothing to copy: concatenating a single batch is zero-copy, and for `Utf8View`/`BinaryView` columns only the views are copied while the data buffers stay shared with the inputs. - `try_create_array_map` becomes `array_map_key_range`: it only decides whether the perfect hash join applies and reserves the `ArrayMap`. Both the hash map and the `ArrayMap` paths now hand the batches to the helper by value, so the inputs are gone before the visited bitmap and the null-aware scope maps are allocated. - Join key arrays that do not share the buffers of the build batch (any on-expression that is not a plain column) are kept for the whole join and are now reserved too. This uses a new public `RecordBatchMemoryCounter::count_array`. This does not reduce memory usage: the build still peaks at about twice the build side, it is now counted. A memory-limited query that only fit because the copy was not counted can now fail with `ResourcesExhausted`. ## What is the testing strategy for this PR? - `join_build_concat_is_reserved` (CollectLeft and Partitioned, hash map and `ArrayMap`): fails with `ResourcesExhausted` at map + 1.5x inputs, passes at map + 3x inputs. It fails on `main` in all four variants. - `join_build_key_arrays_are_reserved` (hash map and `ArrayMap`): with a single build batch and an `a1 + 1` join key, the join fails at inputs + map + half the key array and passes with room for it; a plain column key passes at the lower limit. It fails when the reservation is removed. - Unit tests of the helper: output equals `concat_batches` in both orders, the copy is reserved and the reservation then matches the retained batch, a single batch is not reserved twice, view data buffers are not reserved twice. - Existing `single_partition_join_overallocation` / `partitioned_join_overallocation` are unchanged. ## Are there any user-facing changes? One public addition: `RecordBatchMemoryCounter::count_array`. Hash joins running close to their memory limit may now report `ResourcesExhausted` where they previously exceeded the limit unnoticed. Backport provenance: - Upstream PR: https://github.com/apache/datafusion/pull/25496 - Upstream commit: 21a3215b66a8620cb932899de419e6b43d1848e6 - Preserve upstream count_array, concat/accounting helpers and collector, plus all original tests. Resolve patch anchors and import wrapping for the fork. - The upstream rustdoc replacement's old line is absent from the fork; retain its added line at the corresponding existing helper without fabricating a deletion. Existing fork &bounds context is unchanged. - No fork candidate guard or missing Array import is mixed into this commit; those adaptations follow immediately. This original-only state has not been checked for buildability. Validate the adapted state, not this intermediate. - The original description's tests and benchmarks refer to upstream results. Signed-off-by: discord9 --- datafusion/common/src/utils/memory.rs | 20 + .../physical-plan/src/joins/hash_join/exec.rs | 554 ++++++++++++++++-- 2 files changed, 522 insertions(+), 52 deletions(-) diff --git a/datafusion/common/src/utils/memory.rs b/datafusion/common/src/utils/memory.rs index fd405e06a262e..e235a578587cb 100644 --- a/datafusion/common/src/utils/memory.rs +++ b/datafusion/common/src/utils/memory.rs @@ -186,6 +186,14 @@ impl RecordBatchMemoryCounter { self.memory_usage - previous_memory_usage } + /// Count `array`, returning the memory used by its buffers that have not + /// been counted before. + pub fn count_array(&mut self, array: &dyn Array) -> usize { + let previous_memory_usage = self.memory_usage; + self.count_array_memory_size(array); + self.memory_usage - previous_memory_usage + } + /// Total memory of the unique buffers of all batches counted so far. pub fn memory_usage(&self) -> usize { self.memory_usage @@ -483,6 +491,18 @@ mod record_batch_tests { assert_eq!(counter.memory_usage(), array_data_memory_size(array)); } + #[test] + fn test_count_array_counts_shared_buffers_once() { + let array = Int32Array::from(vec![1, 2, 3, 4, 5]); + let size = array_data_memory_size(&array); + + let mut counter = RecordBatchMemoryCounter::new(); + assert_eq!(counter.count_array(&array), size); + // A slice shares the buffer that is already counted + assert_eq!(counter.count_array(&array.slice(1, 2)), 0); + assert_eq!(counter.memory_usage(), size); + } + #[test] fn test_get_record_batch_memory_size() { let schema = Arc::new(Schema::new(vec![ diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index cc09975a39559..e1183eae45cf9 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -77,7 +77,9 @@ use arrow::util::bit_util; use arrow_schema::{DataType, Schema}; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::TreeNodeRecursion; -use datafusion_common::utils::memory::{RecordBatchMemoryCounter, estimate_memory_size}; +use datafusion_common::utils::memory::{ + RecordBatchMemoryCounter, estimate_memory_size, get_record_batch_memory_size, +}; use datafusion_common::{ JoinSide, JoinType, NullEquality, Result, assert_or_internal_err, internal_err, plan_err, project_schema, @@ -107,8 +109,14 @@ pub(crate) const HASH_JOIN_SEED: SeededRandomState = const ARRAY_MAP_CREATED_COUNT_METRIC_NAME: &str = "array_map_created_count"; +/// Decides whether the build side should be joined with an [`ArrayMap`] +/// (perfect hash join), returning the `(min, max)` key range to build it over. +/// +/// On `Some`, the memory of the [`ArrayMap`] has already been added to +/// `reservation`; the caller concatenates the build batches and creates the map +/// with [`ArrayMap::try_new`]. #[expect(clippy::too_many_arguments)] -fn try_create_array_map( +fn array_map_key_range( bounds: &Option, schema: &SchemaRef, batches: &[RecordBatch], @@ -117,7 +125,7 @@ fn try_create_array_map( perfect_hash_join_small_build_threshold: usize, perfect_hash_join_min_key_density: f64, null_equality: NullEquality, -) -> Result)>> { +) -> Result> { if on_left.len() != 1 { return Ok(None); } @@ -182,12 +190,7 @@ fn try_create_array_map( let mem_size = ArrayMap::estimate_memory_size(min_val, max_val, num_row); reservation.try_grow(mem_size)?; - let batch = concat_batches(schema, batches)?; - let left_values = evaluate_expressions_to_arrays(on_left, &batch)?; - - let array_map = ArrayMap::try_new(&left_values[0], min_val, max_val)?; - - Ok(Some((array_map, batch, left_values))) + Ok(Some((min_val, max_val))) } /// HashTable and input data for the left (build side) of a join @@ -572,7 +575,7 @@ impl From<&HashJoinExec> for HashJoinExecBuilder { /// 4. build_side.num_rows() < u32::MAX /// 5. NullEqualsNothing || (NullEqualsNull && build side doesn't contain null) /// -/// See [`try_create_array_map`] for more details. +/// See [`array_map_key_range`] for more details. /// /// Note that when using [`PartitionMode::Partitioned`], the build side is split into multiple /// partitions. This can cause a dense build side to become sparse within each partition, @@ -2185,6 +2188,7 @@ impl BuildSideState { } } +/// range and density, see [`array_map_key_range`]. fn should_collect_min_max_for_perfect_hash( on_left: &[PhysicalExprRef], schema: &SchemaRef, @@ -2235,6 +2239,75 @@ fn new_join_hashmap( } } +/// Estimates the memory newly allocated for `array`'s share of a concatenated +/// array. +/// +/// Concatenating view arrays copies only the views and keeps sharing the data +/// buffers of the inputs, which are already accounted for. +fn estimate_concat_allocation(array: &dyn Array) -> Result { + match array.data_type() { + DataType::Utf8View | DataType::BinaryView => { + let nulls = array + .nulls() + .map(|_| bit_util::ceil(array.len(), 8)) + .unwrap_or_default(); + Ok(array.len() * size_of::() + nulls) + } + _ => Ok(array.to_data().get_slice_memory_size()?), + } +} + +/// Concatenates the build side `batches` into a single batch, in reverse order +/// if `reverse` is set, keeping `reservation` ahead of the actual memory usage. +/// +/// `inputs_reserved` is the memory `reservation` already holds for `batches`. +/// The copy is reserved before it is made, and once `batches` are dropped the +/// reservation is trimmed to the memory retained by the returned batch. +fn concat_build_batches( + schema: &SchemaRef, + batches: Vec, + reverse: bool, + inputs_reserved: usize, + reservation: &mut MemoryReservation, + metrics: &BuildProbeJoinMetrics, +) -> Result { + // Concatenating a single batch is zero-copy + let copy_size = if batches.len() > 1 { + let mut copy_size = 0; + for batch in &batches { + for array in batch.columns() { + copy_size += estimate_concat_allocation(array.as_ref())?; + } + } + copy_size + } else { + 0 + }; + reservation.try_grow(copy_size)?; + metrics.build_mem_used.add(copy_size); + + let batch = if reverse { + concat_batches(schema, batches.iter().rev())? + } else { + concat_batches(schema, batches.iter())? + }; + drop(batches); + + // The inputs are gone: only hold on to what the concatenated batch retains, + // which includes any buffers it still shares with the inputs. + let held = inputs_reserved + copy_size; + let retained = get_record_batch_memory_size(&batch); + if retained > held { + reservation.try_grow(retained - held)?; + metrics.build_mem_used.add(retained - held); + } else { + reservation.shrink(held - retained); + metrics.build_mem_used.sub(held - retained); + } + + Ok(batch) +} + /// Collects all batches from the left (build) side stream and creates a hash map for joining. /// /// This function is responsible for: @@ -2253,6 +2326,13 @@ fn new_join_hashmap( /// * `probe_threads_count` - Number of threads that will probe this hash table /// * `should_compute_dynamic_filters` - Whether to compute min/max bounds for dynamic filtering /// +/// # Memory Accounting +/// Build batches are added to `reservation` as they arrive. They are then copied +/// into a single batch, see [`concat_build_batches`]: the copy is reserved +/// before it is made, and the reservation is trimmed to what the single batch +/// retains once the input batches are dropped. Join key arrays that do not +/// share the buffers of that batch are reserved as well. +/// /// # Dynamic Filter Coordination /// When `should_compute_dynamic_filters` is true, this function computes the min/max bounds /// for each join key column but does NOT update the dynamic filter. Instead, the @@ -2322,8 +2402,9 @@ async fn collect_left_input( metrics, mut reservation, bounds_accumulators, - memory_counter: _, + memory_counter, } = state; + let inputs_reserved = memory_counter.memory_usage(); // Compute bounds let mut bounds = match bounds_accumulators { @@ -2337,8 +2418,8 @@ async fn collect_left_input( _ => None, }; - let (join_hash_map, batch, left_values) = - if let Some((array_map, batch, left_value)) = try_create_array_map( + let (join_hash_map, batch, left_values) = if let Some((min_val, max_val)) = + array_map_key_range( &bounds, &schema, &batches, @@ -2348,48 +2429,75 @@ async fn collect_left_input( config.execution.perfect_hash_join_min_key_density, null_equality, )? { - array_map_created_count.add(1); - metrics.build_mem_used.add(array_map.size()); + let batch = concat_build_batches( + &schema, + batches, + false, + inputs_reserved, + &mut reservation, + &metrics, + )?; + let left_values = evaluate_expressions_to_arrays(&on_left, &batch)?; + let array_map = ArrayMap::try_new(&left_values[0], min_val, max_val)?; - (Map::ArrayMap(array_map), batch, left_value) - } else { - // Estimation of memory size, required for hashtable, prior to allocation. - // Final result can be verified using `RawTable.allocation_info()` - // Use `u32` indices for the JoinHashMap when num_rows ≤ u32::MAX, otherwise use the - // `u64` indice variant - // Arc is used instead of Box to allow sharing with SharedBuildAccumulator for hash map pushdown - let mut hashmap = new_join_hashmap(num_rows, &mut reservation, &metrics)?; - - let mut hashes_buffer = Vec::new(); - let mut offset = 0; - - let batches_iter = batches.iter().rev(); - - // Updating hashmap starting from the last batch - for batch in batches_iter.clone() { - hashes_buffer.clear(); - hashes_buffer.resize(batch.num_rows(), 0); - update_hash( - &on_left, - batch, - &mut *hashmap, - offset, - &random_state, - &mut hashes_buffer, - 0, - true, - null_equality, - )?; - offset += batch.num_rows(); - } + array_map_created_count.add(1); + metrics.build_mem_used.add(array_map.size()); - // Merge all batches into a single batch, so we can directly index into the arrays - let batch = concat_batches(&schema, batches_iter.clone())?; + (Map::ArrayMap(array_map), batch, left_values) + } else { + // Estimation of memory size, required for hashtable, prior to allocation. + // Final result can be verified using `RawTable.allocation_info()` + // Use `u32` indices for the JoinHashMap when num_rows ≤ u32::MAX, otherwise use the + // `u64` indice variant + // Arc is used instead of Box to allow sharing with SharedBuildAccumulator for hash map pushdown + let mut hashmap = new_join_hashmap(num_rows, &mut reservation, &metrics)?; + + let mut hashes_buffer = Vec::new(); + let mut offset = 0; + + // Updating hashmap starting from the last batch + for batch in batches.iter().rev() { + hashes_buffer.clear(); + hashes_buffer.resize(batch.num_rows(), 0); + update_hash( + &on_left, + batch, + &mut *hashmap, + offset, + &random_state, + &mut hashes_buffer, + 0, + true, + null_equality, + )?; + offset += batch.num_rows(); + } - let left_values = evaluate_expressions_to_arrays(&on_left, &batch)?; + // Merge all batches into a single batch, so we can directly index into the arrays + let batch = concat_build_batches( + &schema, + batches, + true, + inputs_reserved, + &mut reservation, + &metrics, + )?; - (Map::HashMap(hashmap), batch, left_values) - }; + let left_values = evaluate_expressions_to_arrays(&on_left, &batch)?; + + (Map::HashMap(hashmap), batch, left_values) + }; + + // Join keys that are plain columns share the buffers of `batch`, any other + // expression evaluates to new arrays that are kept for the whole join. + let mut key_counter = RecordBatchMemoryCounter::new(); + key_counter.count_batch(&batch); + let keys_size = left_values + .iter() + .map(|values| key_counter.count_array(values.as_ref())) + .sum::(); + reservation.try_grow(keys_size)?; + metrics.build_mem_used.add(keys_size); // Reserve additional memory for visited indices bitmap and create shared builder let visited_indices_bitmap = if with_visited_indices_bitmap { @@ -2536,7 +2644,8 @@ mod tests { }; use arrow::array::{ - Date32Array, Int32Array, Int64Array, StructArray, UInt32Array, UInt64Array, + BinaryViewArray, Date32Array, Int32Array, Int64Array, StringArray, + StringViewArray, StructArray, UInt32Array, UInt64Array, }; use arrow::buffer::NullBuffer; use arrow::datatypes::{DataType, Field}; @@ -2547,6 +2656,9 @@ mod tests { exec_err, internal_err, }; use datafusion_execution::config::SessionConfig; + use datafusion_execution::memory_pool::{ + GreedyMemoryPool, MemoryPool, UnboundedMemoryPool, + }; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_expr::Operator; use datafusion_physical_expr::expressions::{BinaryExpr, Literal}; @@ -6153,6 +6265,344 @@ mod tests { Ok(()) } + /// Build side batches of `num_batches` x `num_rows` rows with distinct + /// buffers: an Int32 key, a Utf8, a Utf8View and a BinaryView payload. + fn concat_test_batches(num_batches: usize, num_rows: usize) -> Vec { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, false), + Field::new("s", DataType::Utf8, true), + Field::new("v", DataType::Utf8View, true), + Field::new("bv", DataType::BinaryView, true), + ])); + (0..num_batches) + .map(|b| { + let start = (b * num_rows) as i32; + let keys = Int32Array::from_iter_values(start..start + num_rows as i32); + let strings = (0..num_rows) + .map(|i| (i % 7 != 0).then(|| format!("string-{b}-{i}"))) + .collect::(); + let views = (0..num_rows) + .map(|i| { + (i % 5 != 0).then(|| format!("a long string view value {b}-{i}")) + }) + .collect::(); + let binary_views = (0..num_rows) + .map(|i| { + (i % 3 != 0).then(|| format!("a long binary view value {b}-{i}")) + }) + .collect::(); + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(keys), + Arc::new(strings), + Arc::new(views), + Arc::new(binary_views), + ], + ) + .unwrap() + }) + .collect() + } + + /// Reserves `batches` the way `collect_left_input` does. + fn reserve_inputs( + batches: &[RecordBatch], + pool: &Arc, + ) -> Result<(MemoryReservation, usize)> { + let reservation = MemoryConsumer::new("HashJoinInput").register(pool); + let mut counter = RecordBatchMemoryCounter::new(); + for batch in batches { + reservation.try_grow(counter.count_batch(batch))?; + } + Ok((reservation, counter.memory_usage())) + } + + #[test] + fn concat_build_batches_matches_concat_batches() -> Result<()> { + let batches = concat_test_batches(4, 100); + let schema = batches[0].schema(); + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + + for reverse in [false, true] { + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let (mut reservation, inputs_reserved) = reserve_inputs(&batches, &pool)?; + let expected = if reverse { + concat_batches(&schema, batches.iter().rev())? + } else { + concat_batches(&schema, batches.iter())? + }; + let batch = concat_build_batches( + &schema, + batches.clone(), + reverse, + inputs_reserved, + &mut reservation, + &metrics, + )?; + assert_eq!(batch, expected); + } + Ok(()) + } + + #[test] + fn concat_build_batches_reserves_copy() -> Result<()> { + let batches = concat_test_batches(4, 1000); + let schema = batches[0].schema(); + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + let inputs: usize = batches.iter().map(get_record_batch_memory_size).sum(); + + // The inputs fit, but not together with their concatenated copy + let pool: Arc = Arc::new(GreedyMemoryPool::new(inputs * 5 / 4)); + let (mut reservation, inputs_reserved) = reserve_inputs(&batches, &pool)?; + assert_eq!(inputs_reserved, inputs); + let err = concat_build_batches( + &schema, + batches.clone(), + false, + inputs_reserved, + &mut reservation, + &metrics, + ) + .unwrap_err(); + assert_contains!(err.to_string(), "Resources exhausted"); + drop(reservation); + + // With room for the copy, the reservation ends up at what is retained + let pool: Arc = Arc::new(GreedyMemoryPool::new(inputs * 2)); + let (mut reservation, inputs_reserved) = reserve_inputs(&batches, &pool)?; + let batch = concat_build_batches( + &schema, + batches, + false, + inputs_reserved, + &mut reservation, + &metrics, + )?; + assert_eq!(reservation.size(), get_record_batch_memory_size(&batch)); + assert_eq!(pool.reserved(), reservation.size()); + Ok(()) + } + + /// Concatenating a single batch is zero-copy, so nothing more is reserved. + #[test] + fn concat_build_batches_single_batch_not_reserved_twice() -> Result<()> { + let batches = concat_test_batches(1, 1000); + let schema = batches[0].schema(); + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + let inputs = get_record_batch_memory_size(&batches[0]); + + let pool: Arc = Arc::new(GreedyMemoryPool::new(inputs)); + let (mut reservation, inputs_reserved) = reserve_inputs(&batches, &pool)?; + let batch = concat_build_batches( + &schema, + batches, + true, + inputs_reserved, + &mut reservation, + &metrics, + )?; + assert_eq!(batch.num_rows(), 1000); + assert_eq!(reservation.size(), inputs); + Ok(()) + } + + /// Only the views of a view array are copied, its data buffers stay shared. + #[rstest] + fn concat_build_batches_view_data_not_reserved_twice( + #[values("v", "bv")] column: &str, + ) -> Result<()> { + let batches = concat_test_batches(4, 1000); + let column = batches[0].schema().index_of(column)?; + let batches = batches + .into_iter() + .map(|batch| batch.project(&[column])) + .collect::, _>>()?; + let schema = batches[0].schema(); + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + let inputs: usize = batches.iter().map(get_record_batch_memory_size).sum(); + let views = 4 * 1000 * size_of::(); + assert!(inputs > 2 * views); + + let pool: Arc = + Arc::new(GreedyMemoryPool::new(inputs + views * 5 / 4)); + let (mut reservation, inputs_reserved) = reserve_inputs(&batches, &pool)?; + let batch = concat_build_batches( + &schema, + batches, + false, + inputs_reserved, + &mut reservation, + &metrics, + )?; + assert_eq!(reservation.size(), get_record_batch_memory_size(&batch)); + Ok(()) + } + + /// The build side is concatenated into a single batch, and that copy must + /// be visible to the memory pool while the input batches are still alive. + #[rstest] + #[tokio::test] + async fn join_build_concat_is_reserved( + #[values(PartitionMode::CollectLeft, PartitionMode::Partitioned)] + mode: PartitionMode, + #[values(false, true)] use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let num_rows = 4000; + let values = (0..num_rows).collect::>(); + let batches = (0..4) + .map(|_| build_table_i32(("a1", &values), ("b1", &values), ("c1", &values))) + .collect::>(); + let inputs: usize = batches.iter().map(get_record_batch_memory_size).sum(); + let left_schema = batches[0].schema(); + let right = build_table( + ("a2", &vec![10, 11]), + ("b2", &vec![12, 13]), + ("c2", &vec![14, 15]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("a1", &left_schema)?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + + // Reserved on top of the inputs and their copy + let build_rows = batches.len() * num_rows as usize; + let map = if use_perfect_hash_join_as_possible { + ArrayMap::estimate_memory_size(0, num_rows as u64 - 1, build_rows) + } else { + estimate_memory_size::<(u32, u64)>(build_rows, size_of::())? + + build_rows * size_of::() + }; + + for (limit, fits) in [(map + inputs * 3 / 2, false), (map + inputs * 3, true)] { + let left = TestMemoryExec::try_new_exec( + std::slice::from_ref(&batches), + Arc::clone(&left_schema), + None, + )?; + let join = HashJoinExec::try_new( + left, + Arc::clone(&right), + on.clone(), + None, + &JoinType::Inner, + None, + mode, + NullEquality::NullEqualsNothing, + false, + )?; + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(limit, 1.0) + .build_arc()?; + let task_ctx = prepare_task_ctx(8192, use_perfect_hash_join_as_possible); + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(task_ctx.session_config().clone()) + .with_runtime(runtime), + ); + + let result = common::collect(join.execute(0, task_ctx)?).await; + if fits { + result?; + } else { + assert_contains!( + result.unwrap_err().to_string(), + "Resources exhausted: Additional allocation failed for HashJoinInput" + ); + } + } + Ok(()) + } + + /// Join keys that are not plain columns evaluate to new arrays, which are + /// kept for the whole join and must be reserved. + #[rstest] + #[tokio::test] + async fn join_build_key_arrays_are_reserved( + #[values(false, true)] use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let num_rows = 16000; + let values = (0..num_rows).collect::>(); + // A single build batch, so that there is no concatenated copy + let batch = build_table_i32(("a1", &values), ("b1", &values), ("c1", &values)); + let inputs = get_record_batch_memory_size(&batch); + let keys = inputs / 3; + let left_schema = batch.schema(); + let right = build_table( + ("a2", &vec![10, 11]), + ("b2", &vec![12, 13]), + ("c2", &vec![14, 15]), + ); + + let map = if use_perfect_hash_join_as_possible { + ArrayMap::estimate_memory_size(1, num_rows as u64, num_rows as usize) + } else { + estimate_memory_size::<(u32, u64)>( + num_rows as usize, + size_of::(), + )? + num_rows as usize * size_of::() + }; + + for (computed_key, limit, fits) in [ + (false, inputs + map + keys / 2, true), + (true, inputs + map + keys / 2, false), + (true, inputs + map + keys * 2, true), + ] { + let column = Arc::new(Column::new_with_schema("a1", &left_schema)?) as _; + let left_key: PhysicalExprRef = if computed_key { + Arc::new(BinaryExpr::new( + column, + Operator::Plus, + Arc::new(Literal::new(ScalarValue::Int32(Some(1)))), + )) + } else { + column + }; + let on = vec![( + left_key, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + let left = TestMemoryExec::try_new_exec( + &[vec![batch.clone()]], + Arc::clone(&left_schema), + None, + )?; + let join = HashJoinExec::try_new( + left, + Arc::clone(&right), + on, + None, + &JoinType::Inner, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + false, + )?; + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(limit, 1.0) + .build_arc()?; + let task_ctx = prepare_task_ctx(8192, use_perfect_hash_join_as_possible); + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(task_ctx.session_config().clone()) + .with_runtime(runtime), + ); + + let result = common::collect(join.execute(0, task_ctx)?).await; + if fits { + result?; + } else { + assert_contains!( + result.unwrap_err().to_string(), + "Resources exhausted: Additional allocation failed for HashJoinInput" + ); + } + } + Ok(()) + } + fn build_table_struct( struct_name: &str, field_name_and_values: (&str, &Vec>), From 6e033ea407c805c1181877d85108fc51e68a081d Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 6 Nov 2026 21:21:38 +0800 Subject: [PATCH 05/19] fix: adapt #25496 imports and perfect-hash eligibility to the fork Import Array for the upstream concat estimator: this trait was already imported in the upstream parent but not in the older fork. Use the fork's existing supported-type candidate helper when checking the array-map key range, preserving upstream's eligibility requirement. Upstream PR: https://github.com/apache/datafusion/pull/25496 Upstream commit: 21a3215b66a8620cb932899de419e6b43d1848e6 This commit contains only fork adaptation, not upstream-authored changes. The adapted state compiled in full all-target/all-feature Clippy. Previously started common/physical-plan nextest also completed: 2289 passed. Signed-off-by: discord9 --- datafusion/physical-plan/src/joins/hash_join/exec.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index e1183eae45cf9..4f313f457c62a 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -69,7 +69,7 @@ use crate::{ metrics::{ExecutionPlanMetricsSet, MetricsSet}, }; -use arrow::array::{ArrayRef, BooleanBufferBuilder}; +use arrow::array::{Array, ArrayRef, BooleanBufferBuilder}; use arrow::compute::concat_batches; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; @@ -126,7 +126,7 @@ fn array_map_key_range( perfect_hash_join_min_key_density: f64, null_equality: NullEquality, ) -> Result> { - if on_left.len() != 1 { + if !should_collect_min_max_for_perfect_hash(on_left, schema)? { return Ok(None); } From 276f9302b3b00eb4413d1c8890a8d9428977f700 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Sat, 26 Sep 2026 15:15:00 +0000 Subject: [PATCH 06/19] feat: support reusable prepared hash-join builds (#25491) ## Which issue does this PR close? Part of [Comet #6013](https://github.com/apache/datafusion-comet/issues/6013). [Comet #6037](https://github.com/apache/datafusion-comet/pull/6037) uses this API. The shared row-index accounting fix in [#25508](https://github.com/apache/datafusion/pull/25508) is included in the base branch. ## Rationale for this change Independent joins against the same build snapshot repeat preparation and retain separate hash tables. `CollectLeft` shares a build among partitions of one join, but cannot share it across independently created plans. For example, eight concurrent requests joining events to the same users snapshot can share one users lookup while probing their own event batches. Preparing that lookup once avoids seven redundant builds. Applications can replace the snapshot while existing consumers finish with the old one. This supports both Comet broadcast tasks and other DataFusion applications. Comet creates independent native joins for its executor tasks, so their existing builds can overlap across task slots. The intended benefit is sharing one retained build batch and hash table and avoiding repeated materialization and hashing work. Removing repeated builds does not imply a proportional latency improvement: cold consumers wait for preparation, then probe independently. Latency depends on scheduling, contention, build size, coordination, and whether a prepared build is already retained. This PR does not establish end-to-end Comet latency, process CPU-time, or RSS savings. ## What changes are included in this PR? `HashJoinExec::prepare_build` returns an immutable `PreparedHashJoinBuild` that compatible joins attach explicitly: ```rust let prepared = first_join .prepare_build(users_snapshot, durable_pool, config) .await?; let first = first_join.builder() .with_prepared_build(Arc::clone(&prepared)) .build()?; let second = second_join.builder() .with_prepared_build(prepared) .build()?; ``` Consumers share build rows, lookup storage, and the retained memory reservation. Probe progress, residual predicates, and dynamic filters remain independent. Preparation charges retained data and temporary working memory; errors and cancellation release unfinished work. Concatenation alignment is charged once per output buffer, so small input batches do not multiply that reservation. Supported builds use non-spilling `CollectLeft` inner joins, direct-column keys with matching types, and fixed-width or UTF-8 build columns. Attachment and execution check schema, keys, and null equality, replaces the unused build subtree with a schema placeholder, and preserves the handle through supported resets and projections. Prepared state is process-local and cannot be serialized. The caller owns snapshot identity, concurrent preparation, and eviction. Input buffers must survive producer cleanup, and the supplied pool must cover every consumer. Attach after physical optimization, supply fresh probe plans and dynamic-filter expressions, and retain a build handle while any associated filter outlives its plan. Compatibility checks cannot verify snapshot contents. ## What is the testing strategy for this PR? Fifteen focused tests cover equivalence with ordinary joins, both null-equality modes, empty and duplicate keys, batch boundaries, concurrent consumers with independent filters and predicates, plan rewrites, public key changes after attachment, admission failures, small-batch concatenation under a fixed memory limit, cancellation, and buffer/reservation lifetimes. The latest regressions exercise under-admission without shrinking the reservation, checked copy/scratch arithmetic, and IN-list buffer sharing across one/two keys and one/two input batches. Existing serialization coverage rejects prepared state; the API example is a doctest. The benchmark checks outputs across 24 cases. Prepared copy and scratch sizing retain checked arithmetic. A concat output that exceeds its admitted bound returns an error in every build profile before shrinking the reservation. Numeric IN-list membership shares the retained batch's key buffers; charging its full array size again would duplicate that payload charge. General array/schema metadata and consumer-local filter expressions are outside the buffer reservation. Validation of [8a964a88f4](https://github.com/apache/datafusion/commit/8a964a88f4a656a6c745b8bfd8c8e543a670ad51) passed: all 15 focused prepared-build tests in the normal `ci` profile; `cargo fmt --all`; full workspace Clippy with all targets/features and `-D warnings`; the complete repository lint suite, including Rust and HTML documentation builds; and the extended workspace suite (12,159 Rust tests and all 524 SQL logic files, with 8 Rust tests ignored). All 15 prepared-build tests also passed in the optimized `release-nonlto` profile. The public-key-mutation regression was verified to fail before restoring execution-time validation: the ordinary join returned its two expected rows while the stale prepared lookup returned none. With the guard restored, the mutated descriptor is rejected with a planning error. Local validation used Rust 1.98.1 and temporary source overrides for the official Arrow 60.0.0 (`ef1fa157`), object_store 0.14.2 (`279572ea`), and sqlparser 0.63.0 (`85b1a6f2`) releases. Dependency overrides and generated lockfile changes are not included. The rebased head has not been benchmarked. ```bash cargo bench -p datafusion-physical-plan --features test_utils --profile release-nonlto --bench prepared_hash_join ``` ## Are there any user-facing changes? Opt-in Rust APIs and an `EXPLAIN` marker for prepared builds. Ordinary SQL planning and join selection are unchanged. Backport provenance: - Upstream PR: https://github.com/apache/datafusion/pull/25491 - Upstream commit: acf5c89e6454469c763c0fb5539e5a294d3439ed - Preserve the upstream prepared-build API, benchmark and original tests. The three new Rust files and serialization test are upstream-exact. - Resolve conflicts against the older fork's existing owner/codec structure. Upstream-only neighboring null-mark/ProbeCompletion/RightAnti documentation and null-aware parameter context have no counterpart; do not fabricate them. The serialization guard uses self.prepared_build before encoding children, retaining upstream's rejection and message. - The immediately following commit contains the older-fork probe-state and collector-call adaptations. Additional local tests are deferred to the end. - The upstream description's tests and benchmark measurements are upstream results, not a claim that this original-only intermediate state was tested. Signed-off-by: discord9 --- datafusion/physical-plan/Cargo.toml | 5 + .../benches/prepared_hash_join.rs | 272 ++++++ .../physical-plan/src/joins/hash_join/exec.rs | 296 ++++-- .../src/joins/hash_join/exec/prepared.rs | 367 +++++++ .../joins/hash_join/exec/prepared/tests.rs | 905 ++++++++++++++++++ .../physical-plan/src/joins/hash_join/mod.rs | 2 +- .../joins/hash_join/partitioned_hash_eval.rs | 17 +- .../src/joins/hash_join/stream.rs | 4 +- datafusion/physical-plan/src/joins/mod.rs | 3 +- datafusion/proto/tests/cases/plans/joins.rs | 46 + 10 files changed, 1809 insertions(+), 108 deletions(-) create mode 100644 datafusion/physical-plan/benches/prepared_hash_join.rs create mode 100644 datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs create mode 100644 datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index edde8d73cde75..50d02162cfee9 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -140,6 +140,11 @@ harness = false name = "hash_join_semi_anti" required-features = ["test_utils"] +[[bench]] +harness = false +name = "prepared_hash_join" +required-features = ["test_utils"] + [[bench]] harness = false name = "multi_group_by" diff --git a/datafusion/physical-plan/benches/prepared_hash_join.rs b/datafusion/physical-plan/benches/prepared_hash_join.rs new file mode 100644 index 0000000000000..113a2cf75b929 --- /dev/null +++ b/datafusion/physical-plan/benches/prepared_hash_join.rs @@ -0,0 +1,272 @@ +// 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. + +//! Independent INNER joins over the same in-memory build snapshot. +//! +//! SQL planning does not expose prepared builds, so these benchmarks exercise +//! the public execution API. Cases vary build size, integer/string keys, and +//! concurrent consumers. Every consumer probes 4096 rows with 50% matches. +//! `cold` includes preparation; `warm` retains a build across iterations. +//! Input generation is excluded, but plan construction and draining outputs are +//! timed. This does not measure Comet decoding, its cache, or its single-flight +//! wait. Sparse integer keys avoid perfect-hash selection. + +use std::hint::black_box; +use std::sync::Arc; +use std::time::Duration; + +use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray}; +use arrow::datatypes::{DataType, Field, Schema}; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_common::{JoinType, Result}; +use datafusion_common_runtime::SpawnedTask; +use datafusion_execution::TaskContext; +use datafusion_physical_expr::expressions::col; +use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::joins::{ + HashJoinExec, HashJoinExecBuilder, PartitionMode, PreparedHashJoinBuild, +}; +use datafusion_physical_plan::memory::MemoryStream; +use datafusion_physical_plan::test::TestMemoryExec; +use futures::TryStreamExt; +use tokio::runtime::Builder; +use tokio::sync::Barrier; + +const PROBE_ROWS: usize = 4096; + +struct Workload { + build: Vec, + probe: RecordBatch, + expected_sum: i64, +} + +impl Workload { + fn new(build_rows: usize, strings: bool) -> Self { + let batch = |keys: Vec| { + let array: ArrayRef = if strings { + Arc::new(StringArray::from_iter_values( + keys.iter().map(|key| format!("key-{key:08}")), + )) + } else { + Arc::new(Int64Array::from_iter_values( + keys.iter().map(|key| (key * 17) as i64), + )) + }; + RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("key", array.data_type().clone(), false), + Field::new("payload", DataType::Int64, false), + ])), + vec![ + array, + Arc::new(Int64Array::from_iter_values( + keys.iter().map(|key| *key as i64), + )), + ], + ) + .unwrap() + }; + let build = batch((0..build_rows).collect()); + let probe = batch( + (0..PROBE_ROWS) + .map(|i| (i / 2) % build_rows + (i % 2) * build_rows) + .collect(), + ); + Self { + build: (0..build_rows) + .step_by(8192) + .map(|offset| build.slice(offset, (build_rows - offset).min(8192))) + .collect(), + probe, + expected_sum: (0..PROBE_ROWS / 2).map(|i| (i % build_rows) as i64).sum(), + } + } + + fn join(&self) -> Result { + let schema = self.probe.schema(); + HashJoinExecBuilder::new( + TestMemoryExec::try_new_exec( + std::slice::from_ref(&self.build), + Arc::clone(&schema), + None, + )?, + TestMemoryExec::try_new_exec( + &[vec![self.probe.clone()]], + Arc::clone(&schema), + None, + )?, + vec![(col("key", &schema)?, col("key", &schema)?)], + JoinType::Inner, + ) + .with_partition_mode(PartitionMode::CollectLeft) + .build() + } + + async fn prepare( + &self, + join: &HashJoinExec, + context: &TaskContext, + ) -> Result> { + join.prepare_build( + Box::pin(MemoryStream::try_new( + self.build.clone(), + self.probe.schema(), + None, + )?), + Arc::clone(context.memory_pool()), + Arc::clone(context.session_config().options()), + ) + .await + } +} + +#[derive(Clone, Copy)] +enum Mode { + Rebuild, + Cold, + Warm, +} + +impl Mode { + fn name(self) -> &'static str { + match self { + Self::Rebuild => "rebuild", + Self::Cold => "cold", + Self::Warm => "warm", + } + } +} + +async fn run( + workload: &Workload, + consumers: usize, + mode: Mode, + warm: Option>, + context: Arc, +) -> Result<(usize, i64)> { + let joins = (0..consumers) + .map(|_| workload.join()) + .collect::>>()?; + let prepared = match mode { + Mode::Rebuild => None, + Mode::Cold => Some(workload.prepare(&joins[0], &context).await?), + Mode::Warm => warm, + }; + let start = Arc::new(Barrier::new(consumers)); + let tasks = joins + .into_iter() + .map(|join| { + let join = match &prepared { + Some(build) => join + .builder() + .with_prepared_build(Arc::clone(build)) + .build()?, + None => join, + }; + let context = Arc::clone(&context); + let start = Arc::clone(&start); + Ok(SpawnedTask::spawn(async move { + start.wait().await; + let mut stream = join.execute(0, context)?; + let mut result = (0, 0); + while let Some(batch) = stream.try_next().await? { + result.0 += batch.num_rows(); + let payload = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + result.1 += payload.values().iter().sum::(); + } + Ok::<_, datafusion_common::DataFusionError>(result) + })) + }) + .collect::>>()?; + let mut result = (0, 0); + for task in tasks { + let (rows, sum) = task.await.unwrap()?; + result.0 += rows; + result.1 += sum; + } + Ok(result) +} + +fn benchmark(c: &mut Criterion) { + let runtime = Builder::new_multi_thread() + .worker_threads(4) + .enable_all() + .build() + .unwrap(); + let mut group = c.benchmark_group("prepared_hash_join"); + group.sample_size(20); + group.warm_up_time(Duration::from_millis(500)); + group.measurement_time(Duration::from_secs(2)); + + for (key_type, strings) in [("int64", false), ("utf8", true)] { + for rows in [64, 65536] { + let workload = Workload::new(rows, strings); + for consumers in [1, 4] { + for mode in [Mode::Rebuild, Mode::Cold, Mode::Warm] { + let context = Arc::new(TaskContext::default()); + let warm = matches!(mode, Mode::Warm).then(|| { + runtime + .block_on( + workload.prepare(&workload.join().unwrap(), &context), + ) + .unwrap() + }); + let output = runtime + .block_on(run( + &workload, + consumers, + mode, + warm.clone(), + Arc::clone(&context), + )) + .unwrap(); + assert_eq!( + output, + ( + consumers * PROBE_ROWS / 2, + consumers as i64 * workload.expected_sum + ) + ); + let case = format!("{key_type}/{rows}/{consumers}_consumers"); + group.bench_function(BenchmarkId::new(mode.name(), case), |b| { + b.iter(|| { + black_box( + runtime + .block_on(run( + &workload, + consumers, + mode, + warm.clone(), + Arc::clone(&context), + )) + .unwrap(), + ) + }); + }); + } + } + } + } + group.finish(); +} + +criterion_group!(benches, benchmark); +criterion_main!(benches); diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 4f313f457c62a..c23e1d81b2016 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -76,6 +76,7 @@ use arrow::record_batch::RecordBatch; use arrow::util::bit_util; use arrow_schema::{DataType, Schema}; use datafusion_common::config::ConfigOptions; +use datafusion_common::stats::Precision; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::utils::memory::{ RecordBatchMemoryCounter, estimate_memory_size, get_record_batch_memory_size, @@ -103,6 +104,9 @@ use parking_lot::Mutex; use super::partitioned_hash_eval::SeededRandomState; +mod prepared; +pub use prepared::PreparedHashJoinBuild; + /// Hard-coded seed to ensure hash values from the hash join differ from `RepartitionExec`, avoiding collisions. pub(crate) const HASH_JOIN_SEED: SeededRandomState = SeededRandomState::with_seed(12210250226015887276); @@ -193,48 +197,47 @@ fn array_map_key_range( Ok(Some((min_val, max_val))) } -/// HashTable and input data for the left (build side) of a join -pub(super) struct JoinLeftData { - /// The hash table with indices into `batch` - /// Arc is used to allow sharing with SharedBuildAccumulator for hash map pushdown - pub(super) map: Arc, - /// The input rows for the build side +/// Immutable build buffers and their durable reservation. +struct JoinBuildData { + /// Hash table with row indices into `batch`, also shared with dynamic filters. + map: Arc, + /// The input rows for the build side. batch: RecordBatch, - /// The build side on expressions values + /// Evaluated build-side key expressions. values: Vec, - /// Shared bitmap builder for visited left indices + /// Bounds computed from the build side; absent for an empty partition. + bounds: Option, + /// IN-list values or a hash-table reference used for filter pushdown. + membership: PushdownStrategy, + // Keep the reservation after the allocations it accounts for. + reservation: MemoryReservation, +} + +/// A build lease and mutable bookkeeping for one join execution. +pub(super) struct JoinLeftData { + /// Shared bitmap for visited left indices within this consuming join. visited_indices_bitmap: SharedBitmapBuilder, /// Counter of running probe-threads, potentially /// able to update `visited_indices_bitmap` probe_threads_counter: AtomicUsize, - /// We need to keep this field to maintain accurate memory accounting, even though we don't directly use it. - /// Without holding onto this reservation, the recorded memory usage would become inconsistent with actual usage. - /// This could hide potential out-of-memory issues, especially when upstream operators increase their memory consumption. - /// The MemoryReservation ensures proper tracking of memory resources throughout the join operation's lifecycle. - _reservation: MemoryReservation, - /// Bounds computed from the build side for dynamic filter pushdown. - /// If the partition is empty (no rows) this will be None. - /// If the partition has some rows this will be Some with the bounds for each join key column. - pub(super) bounds: Option, - /// Membership testing strategy for filter pushdown - /// Contains either InList values for small build sides or hash table reference for large build sides - pub(super) membership: PushdownStrategy, /// Shared atomic flag indicating if any probe partition saw data (for null-aware anti joins) /// This is shared across all probe partitions to provide global knowledge pub(super) probe_side_non_empty: AtomicBool, /// Shared atomic flag indicating if any probe partition saw NULL in join keys (for null-aware anti joins) pub(super) probe_side_has_null: AtomicBool, + // Drop the reservation after the mutable state it also accounts for. + build: Arc, } impl JoinLeftData { /// return a reference to the map pub(super) fn map(&self) -> &Map { - &self.map + &self.build.map } /// returns a reference to the build side batch pub(super) fn batch(&self) -> &RecordBatch { - &self.batch + &self.build.batch } /// Returns `true` if the build side physically contains rows. @@ -256,7 +259,7 @@ impl JoinLeftData { /// returns a reference to the build side expressions values pub(super) fn values(&self) -> &[ArrayRef] { - &self.values + &self.build.values } /// returns a reference to the visited indices bitmap @@ -266,7 +269,12 @@ impl JoinLeftData { /// returns a reference to the InList values for filter pushdown pub(super) fn membership(&self) -> &PushdownStrategy { - &self.membership + &self.build.membership + } + + /// Borrow bounds without sharing a consuming join's dynamic-filter state. + pub(super) fn bounds(&self) -> Option<&PartitionBounds> { + self.build.bounds.as_ref() } /// Decrements the counter of running threads, and returns `true` @@ -309,6 +317,7 @@ impl HashJoinExecBuilder { filter: None, join_type, left_fut: Default::default(), + prepared_build: None, random_state: HASH_JOIN_SEED, mode: PartitionMode::Auto, fetch: None, @@ -400,13 +409,20 @@ impl HashJoinExecBuilder { children.len() == 2, "wrong number of children passed into `HashJoinExecBuilder`" ); + if self.exec.prepared_build.is_some() + && !Arc::ptr_eq(&self.exec.left, &children[0]) + { + return plan_err!( + "Cannot replace the build child after attaching a prepared hash-join build" + ); + } self.preserve_properties &= has_same_children_properties(&self.exec, &children)?; self.exec.right = children.swap_remove(1); self.exec.left = children.swap_remove(0); Ok(self) } - /// Reset runtime state. + /// Reset task-local runtime state while retaining the immutable prepared build. pub fn reset_state(mut self) -> Self { self.exec.left_fut = Default::default(); self.exec.dynamic_filter = None; @@ -422,10 +438,18 @@ impl HashJoinExecBuilder { /// Build resulting execution plan. pub fn build(self) -> Result { let Self { - exec, + mut exec, preserve_properties, } = self; + if let Some(prepared) = &exec.prepared_build { + prepared.validate(&exec)?; + exec.left_fut = Default::default(); + exec.metrics = ExecutionPlanMetricsSet::new(); + if let Some(filter) = &mut exec.dynamic_filter { + filter.build_accumulator = OnceLock::new(); + } + } // Validate null_aware flag if exec.null_aware { let join_type = exec.join_type(); @@ -454,6 +478,7 @@ impl HashJoinExecBuilder { filter, join_type, left_fut, + prepared_build, random_state, mode, metrics, @@ -501,6 +526,7 @@ impl HashJoinExecBuilder { join_type, join_schema, left_fut, + prepared_build, random_state, mode, metrics, @@ -531,6 +557,7 @@ impl From<&HashJoinExec> for HashJoinExecBuilder { join_type: exec.join_type, join_schema: Arc::clone(&exec.join_schema), left_fut: Arc::clone(&exec.left_fut), + prepared_build: exec.prepared_build.clone(), random_state: exec.random_state.clone(), mode: exec.mode, metrics: exec.metrics.clone(), @@ -760,6 +787,8 @@ pub struct HashJoinExec { /// Each output stream waits on the `OnceAsync` to signal the completion of /// the hash table creation. left_fut: Arc>, + /// Immutable build attached by an embedding executor. + prepared_build: Option>, /// Shared the `SeededRandomState` for the hashing algorithm (seeds preserved for serialization) random_state: SeededRandomState, /// Partitioning mode to use @@ -803,6 +832,7 @@ impl fmt::Debug for HashJoinExec { .field("join_type", &self.join_type) .field("join_schema", &self.join_schema) .field("left_fut", &self.left_fut) + .field("prepared_build", &self.prepared_build) .field("random_state", &self.random_state) .field("mode", &self.mode) .field("metrics", &self.metrics) @@ -1229,6 +1259,12 @@ impl DisplayAs for HashJoinExec { .map_or_else(String::new, |f| format!(", fetch={f}")); let display_null_aware = if self.null_aware { ", null_aware" } else { "" }; + let display_prepared = self + .prepared_build + .as_ref() + .map_or_else(String::new, |prepared| { + format!(", prepared_build={} rows", prepared.num_rows()) + }); let on = self .on .iter() @@ -1237,7 +1273,7 @@ impl DisplayAs for HashJoinExec { .join(", "); write!( f, - "HashJoinExec: mode={:?}, join_type={:?}, on=[{}]{}{}{}{}{}", + "HashJoinExec: mode={:?}, join_type={:?}, on=[{}]{}{}{}{}{}{}", self.mode, self.join_type, on, @@ -1246,6 +1282,7 @@ impl DisplayAs for HashJoinExec { display_null_equality, display_fetch, display_null_aware, + display_prepared, ) } DisplayFormatType::TreeRender => { @@ -1264,6 +1301,10 @@ impl DisplayAs for HashJoinExec { writeln!(f, "on={on}")?; + if let Some(prepared) = &self.prepared_build { + writeln!(f, "prepared_build={} rows", prepared.num_rows())?; + } + if self.null_equality() == NullEquality::NullEqualsNull { writeln!(f, "NullsEqual: true")?; } @@ -1489,8 +1530,16 @@ impl ExecutionPlan for HashJoinExec { .flatten() .flatten(); - let left_fut = match self.mode { - PartitionMode::CollectLeft => self.left_fut.try_once(|| { + let left_fut = match (&self.prepared_build, self.mode) { + (Some(prepared), _) => { + // Public join fields can change after builder validation. + prepared.validate(self)?; + let prepared = Arc::clone(prepared); + self.left_fut.try_once(|| { + Ok(async move { Ok(prepared.probe_data(right_partitions)) }) + })? + } + (None, PartitionMode::CollectLeft) => self.left_fut.try_once(|| { let left_stream = self.left.execute(0, Arc::clone(&context))?; let reservation = @@ -1503,14 +1552,15 @@ impl ExecutionPlan for HashJoinExec { join_metrics.clone(), reservation, need_produce_result_in_final(self.join_type), - self.right().output_partitioning().partition_count(), + right_partitions, enable_dynamic_filter_pushdown, Arc::clone(context.session_config().options()), self.null_equality, array_map_created_count, + BuildMode::Ordinary, )) })?, - PartitionMode::Partitioned => { + (None, PartitionMode::Partitioned) => { let left_stream = self.left.execute(partition, Arc::clone(&context))?; let reservation = @@ -1528,9 +1578,10 @@ impl ExecutionPlan for HashJoinExec { Arc::clone(context.session_config().options()), self.null_equality, array_map_created_count, + BuildMode::Ordinary, )) } - PartitionMode::Auto => { + (None, PartitionMode::Auto) => { return plan_err!( "Invalid HashJoinExec, unsupported PartitionMode {:?} in execute()", PartitionMode::Auto @@ -1613,11 +1664,17 @@ impl ExecutionPlan for HashJoinExec { input_stats: &[Arc], _args: &StatisticsArgs, ) -> Result> { - let left_stats = Arc::clone(&input_stats[0]); - let right_stats = Arc::clone(&input_stats[1]); + let left_stats = if let Some(prepared) = &self.prepared_build { + // The declared left child supplies the schema but is never executed. + Statistics::new_unknown(self.left.schema().as_ref()) + .with_num_rows(Precision::Exact(prepared.num_rows())) + } else { + input_stats[0].as_ref().clone() + }; + let right_stats = input_stats[1].as_ref().clone(); let stats = estimate_join_statistics( - Arc::unwrap_or_clone(left_stats), - Arc::unwrap_or_clone(right_stats), + left_stats, + right_stats, &self.on, self.null_equality, &self.join_type, @@ -1640,6 +1697,9 @@ impl ExecutionPlan for HashJoinExec { if self.contains_projection() { return Ok(None); } + if self.prepared_build.is_some() { + return try_embed_projection(projection, self); + } let schema = self.schema(); if let Some(JoinData { @@ -1847,6 +1907,11 @@ impl ExecutionPlan for HashJoinExec { ) -> Result> { use datafusion_proto_models::protobuf; + if self.prepared_build.is_some() { + return datafusion_common::not_impl_err!( + "HashJoinExec with a prepared build cannot be serialized" + ); + } let left = ctx.encode_child(self.left())?; let right = ctx.encode_child(self.right())?; @@ -2148,6 +2213,12 @@ impl CollectLeftAccumulator { } } +#[derive(PartialEq, Eq)] +enum BuildMode { + Ordinary, + Prepared, +} + /// State for collecting the build-side data during hash join struct BuildSideState { batches: Vec, @@ -2346,7 +2417,7 @@ fn concat_build_batches( #[expect(clippy::too_many_arguments)] async fn collect_left_input( random_state: RandomState, - left_stream: SendableRecordBatchStream, + mut left_stream: SendableRecordBatchStream, on_left: Vec, metrics: BuildProbeJoinMetrics, reservation: MemoryReservation, @@ -2356,13 +2427,15 @@ async fn collect_left_input( config: Arc, null_equality: NullEquality, array_map_created_count: Count, + mode: BuildMode, ) -> Result { let schema = left_stream.schema(); + let prepared = mode == BuildMode::Prepared; let should_collect_min_max_for_phj = should_collect_min_max_for_perfect_hash(&on_left, &schema)?; - let initial = BuildSideState::try_new( + let mut state = BuildSideState::try_new( metrics, reservation, on_left.clone(), @@ -2370,41 +2443,52 @@ async fn collect_left_input( should_compute_dynamic_filters || should_collect_min_max_for_phj, )?; - let state = left_stream - .try_fold(initial, |mut state, batch| async move { - // Update accumulators if computing bounds - if let Some(ref mut accumulators) = state.bounds_accumulators { - for accumulator in accumulators { - accumulator.update_batch(&batch)?; - } + let mut max_batch_rows = 0; + while let Some(batch) = left_stream.try_next().await? { + max_batch_rows = max_batch_rows.max(batch.num_rows()); + if let Some(accumulators) = &mut state.bounds_accumulators { + for accumulator in accumulators { + accumulator.update_batch(&batch)?; } - - // Decide if we spill or not - let batch_size = state.memory_counter.count_batch(&batch); - // Reserve memory for incoming batch - state.reservation.try_grow(batch_size)?; - // Update metrics - state.metrics.build_mem_used.add(batch_size); - state.metrics.build_input_batches.add(1); - state.metrics.build_input_rows.add(batch.num_rows()); - // Update row count - state.num_rows += batch.num_rows(); - // Push batch to output - state.batches.push(batch); - Ok(state) - }) - .await?; - - // Extract fields from state + } + let batch_size = state.memory_counter.count_batch(&batch); + state.reservation.try_grow(batch_size)?; + state.metrics.build_mem_used.add(batch_size); + state.metrics.build_input_batches.add(1); + state.metrics.build_input_rows.add(batch.num_rows()); + state.num_rows += batch.num_rows(); + state.batches.push(batch); + } + drop(left_stream); + if prepared && state.batches.is_empty() { + // Even empty UTF-8 output retains an offset buffer. + let empty = RecordBatch::new_empty(Arc::clone(&schema)); + state + .reservation + .try_grow(state.memory_counter.count_batch(&empty))?; + state.batches.push(empty); + } + let input_bytes = state.memory_counter.memory_usage(); + + // Bind the reservation first so error paths release allocations before their charge. let BuildSideState { - batches, + mut reservation, + mut batches, num_rows, metrics, - mut reservation, bounds_accumulators, - memory_counter, + memory_counter: _, } = state; - let inputs_reserved = memory_counter.memory_usage(); + + // Admit concatenation copies while the original batches are retained. + // Arrow keeps a single batch as an inexpensive slice. + let copy_bytes = if prepared && batches.len() > 1 { + let copy_bytes = prepared::prepared_copy_bytes(&batches)?; + reservation.try_grow(copy_bytes)?; + copy_bytes + } else { + 0 + }; // Compute bounds let mut bounds = match bounds_accumulators { @@ -2429,14 +2513,18 @@ async fn collect_left_input( config.execution.perfect_hash_join_min_key_density, null_equality, )? { - let batch = concat_build_batches( - &schema, - batches, - false, - inputs_reserved, - &mut reservation, - &metrics, - )?; + let batch = if prepared { + concat_batches(&schema, batches.iter())? + } else { + concat_build_batches( + &schema, + std::mem::take(&mut batches), + false, + input_bytes, + &mut reservation, + &metrics, + )? + }; let left_values = evaluate_expressions_to_arrays(&on_left, &batch)?; let array_map = ArrayMap::try_new(&left_values[0], min_val, max_val)?; @@ -2452,7 +2540,16 @@ async fn collect_left_input( // Arc is used instead of Box to allow sharing with SharedBuildAccumulator for hash map pushdown let mut hashmap = new_join_hashmap(num_rows, &mut reservation, &metrics)?; - let mut hashes_buffer = Vec::new(); + let scratch_reservation = reservation.new_empty(); + if prepared { + scratch_reservation.try_grow(prepared::prepared_scratch_bytes( + max_batch_rows, + on_left.len(), + null_equality, + )?)?; + } + // The maximum is known: avoid geometric growth and its excess capacity. + let mut hashes_buffer = Vec::with_capacity(max_batch_rows); let mut offset = 0; // Updating hashmap starting from the last batch @@ -2474,14 +2571,18 @@ async fn collect_left_input( } // Merge all batches into a single batch, so we can directly index into the arrays - let batch = concat_build_batches( - &schema, - batches, - true, - inputs_reserved, - &mut reservation, - &metrics, - )?; + let batch = if prepared { + concat_batches(&schema, batches.iter().rev())? + } else { + concat_build_batches( + &schema, + std::mem::take(&mut batches), + true, + input_bytes, + &mut reservation, + &metrics, + )? + }; let left_values = evaluate_expressions_to_arrays(&on_left, &batch)?; @@ -2544,15 +2645,30 @@ async fn collect_left_input( bounds = None; } + if prepared { + drop(batches); + // Prepared keys are direct columns. IN-list arrays share these batch + // buffers, including through multi-key StructArray children. + let retained = RecordBatchMemoryCounter::new().count_batch(&batch); + prepared::reconcile_prepared_copy_reservation( + &reservation, + input_bytes, + copy_bytes, + retained, + )?; + } + let data = JoinLeftData { - map, - batch, - values: left_values, + build: Arc::new(JoinBuildData { + map, + batch, + values: left_values, + bounds, + membership, + reservation, + }), visited_indices_bitmap: Mutex::new(visited_indices_bitmap), probe_threads_counter: AtomicUsize::new(probe_threads_count), - _reservation: reservation, - bounds, - membership, probe_side_non_empty: AtomicBool::new(false), probe_side_has_null: AtomicBool::new(false), }; diff --git a/datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs b/datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs new file mode 100644 index 0000000000000..fb46cedf96cf1 --- /dev/null +++ b/datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs @@ -0,0 +1,367 @@ +// 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. + +//! Explicit immutable build reuse for embedding executors. + +use super::*; +use crate::spill::spill_manager::GetSlicedSize; +use datafusion_common::exec_datafusion_err; +use datafusion_execution::memory_pool::MemoryPool; + +/// An immutable, fully prepared broadcast build, independent of any probe task. +/// +/// Created by [`HashJoinExec::prepare_build`]. The embedding executor owns cache +/// identity, admission, single-flight coordination, cancellation and eviction. +/// This object retains its input buffers and memory reservation until its last +/// lease is dropped; it never retains an input stream or task context. Prepared +/// builds support fixed-width and UTF-8 build columns, with direct-column keys +/// and non-spilling INNER joins. Residual conditions belong to each consuming +/// join; null-aware joins remain unsupported. +/// +/// Hash-join gathers copy supported build columns into output buffers, +/// including contiguous selections. Output batches can therefore outlive this +/// object without retaining unaccounted cached payload. View, dictionary and +/// nested build columns remain unsupported. UTF-8 and fixed-size binary keys +/// use hash-table membership filters instead of copying range or IN-list values. +pub struct PreparedHashJoinBuild { + build: Arc, + keys: Vec, + null_equality: NullEquality, +} + +impl fmt::Debug for PreparedHashJoinBuild { + /// Describe immutable metadata without dumping table contents. + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("PreparedHashJoinBuild") + .field("schema", &self.build.batch.schema()) + .field("keys", &self.keys) + .field("rows", &self.num_rows()) + .field("reserved_bytes", &self.reserved_bytes()) + .finish() + } +} + +impl PreparedHashJoinBuild { + /// Return the retained build reservation, excluding all per-probe state. + pub fn reserved_bytes(&self) -> usize { + self.build.reservation.size() + } + + /// Return the complete build row count, including duplicate and null keys. + pub fn num_rows(&self) -> usize { + self.build.batch.num_rows() + } + + /// Create independent mutable state for one consuming join. + pub(super) fn probe_data(&self, probe_threads: usize) -> JoinLeftData { + JoinLeftData { + build: Arc::clone(&self.build), + null_aware_scope_map: None, + null_value_build_rows: None, + visited_indices_bitmap: Mutex::new(BooleanBufferBuilder::new(0)), + null_indices_bitmap: Mutex::new(BooleanBufferBuilder::new(0)), + probe_completion: ProbeCompletion::new(probe_threads), + build_side_has_null: false, + } + } + + /// Validate the build descriptor and current execution restrictions without + /// consuming input or modifying either plan. Cache identity is caller-owned. + pub(super) fn validate(&self, join: &HashJoinExec) -> Result<()> { + let keys = prepared_key_indices(join)?; + if join.left.schema() != self.build.batch.schema() + || keys != self.keys + || join.null_equality != self.null_equality + { + return plan_err!( + "Prepared hash-join build does not match schema, keys or null equality" + ); + } + if let Some(filter) = &join.dynamic_filter { + let probe_keys = join.on.iter().map(|(_, right)| right); + if !filter.filter.children().into_iter().eq(probe_keys) { + return plan_err!( + "Prepared hash-join dynamic filter keys do not match probe keys" + ); + } + } + Ok(()) + } +} + +impl HashJoinExec { + /// Prepare the supplied build snapshot using an embedding executor's durable pool. + /// + /// The caller owns snapshot identity: this consumes `input`, never `self.left`. + /// Schema and key compatibility cannot establish that two inputs contain the + /// same data. Attaching the wrong snapshot can silently change query results. + /// + /// # Caller contract + /// + /// * Input buffers must remain valid independently of producer task cleanup. + /// Keep producer allocations charged until the stream releases them. + /// * Supply a pool whose lifetime covers every consumer. It accounts for + /// retained input, hash storage, and preparation's copy/scratch buffers. + /// * Coordinate concurrent preparation, publication, and invalidation in the + /// embedding executor. This method does not provide a cache. + /// + /// Errors and cancellation release unfinished work. `config` controls + /// perfect-map and dynamic-filter choices. Byte keys use hash membership + /// to avoid copying payloads into range bounds and IN-list literals. + /// Consumer-local dynamic-filter allocations are outside this reservation. + /// + /// # Example + /// + /// Given two independently planned, compatible joins over different probe + /// inputs, a service can prepare one dimension snapshot for both consumers. + /// + /// ```no_run + /// # use std::sync::Arc; + /// # use datafusion_common::{Result, config::ConfigOptions}; + /// # use datafusion_execution::memory_pool::MemoryPool; + /// # use datafusion_physical_plan::SendableRecordBatchStream; + /// # use datafusion_physical_plan::joins::HashJoinExec; + /// # async fn example(first_join: HashJoinExec, second_join: HashJoinExec, + /// # dimension_snapshot: SendableRecordBatchStream, + /// # durable_pool: Arc, config: Arc) -> Result<()> { + /// let prepared = first_join + /// .prepare_build(dimension_snapshot, durable_pool, config) + /// .await?; + /// let first = first_join.builder() + /// .with_prepared_build(Arc::clone(&prepared)) + /// .build()?; + /// let second = second_join.builder() + /// .with_prepared_build(prepared) + /// .build()?; + /// # let _ = (first, second); + /// # Ok(()) + /// # } + /// ``` + pub async fn prepare_build( + &self, + input: SendableRecordBatchStream, + pool: Arc, + config: Arc, + ) -> Result> { + let keys = prepared_key_indices(self)?; + let schema = self.left.schema(); + if input.schema() != schema { + return plan_err!( + "Prepared hash-join input schema does not match build schema" + ); + } + let byte_keys = keys.iter().any(|&key| { + matches!( + schema.field(key).data_type(), + DataType::Utf8 | DataType::FixedSizeBinary(_) + ) + }); + // Range accumulation and IN-list publication materialize ScalarValue + // copies of byte keys. Hash membership borrows the admitted table instead. + let config = if byte_keys { + let mut config = config.as_ref().clone(); + config.optimizer.hash_join_inlist_pushdown_max_size = 0; + Arc::new(config) + } else { + config + }; + let metrics_set = ExecutionPlanMetricsSet::new(); + let metrics = BuildProbeJoinMetrics::new(0, &metrics_set); + let count = MetricBuilder::new(&metrics_set) + .counter(ARRAY_MAP_CREATED_COUNT_METRIC_NAME, 0); + let reservation = MemoryConsumer::new("PreparedHashJoinBuild").register(&pool); + let data = collect_left_input( + self.random_state.random_state().clone(), + input, + self.on.iter().map(|(left, _)| Arc::clone(left)).collect(), + metrics, + reservation, + false, + 0, + !byte_keys, + config, + self.null_equality, + None, + count, + BuildMode::Prepared, + ) + .await?; + Ok(Arc::new(PreparedHashJoinBuild { + build: data.build, + keys, + null_equality: self.null_equality, + })) + } +} + +/// Bound copy allocations, including validity, offsets and alignment. Aliased +/// columns count separately because concatenation materializes each column. +pub(super) fn prepared_copy_bytes(batches: &[RecordBatch]) -> Result { + let overflow = || exec_datafusion_err!("Prepared hash-join copy size overflow"); + let mut bytes = 0usize; + for batch in batches { + // Concat can materialize validity for an all-valid input when another input + // contains nulls. Arrow's slice measurement only counts existing bitmaps. + let missing_validity = batch + .columns() + .iter() + .filter(|array| array.nulls().is_none()) + .count(); + let validity_bytes = batch + .num_rows() + .div_ceil(8) + .checked_mul(missing_validity) + .ok_or_else(overflow)?; + bytes = bytes + .checked_add(batch.get_sliced_size()?) + .and_then(|bytes| bytes.checked_add(validity_bytes)) + .ok_or_else(overflow)?; + } + // These are flat arrays: allow one alignment unit per output buffer, + // including potential validity. Concat allocates each buffer only once. + let buffers = batches[0] + .columns() + .iter() + .try_fold(0usize, |buffers, array| { + buffers + .checked_add(array.to_data().buffers().len()) + .and_then(|buffers| buffers.checked_add(1)) + .ok_or_else(overflow) + })?; + buffers + .checked_mul(64) + .and_then(|padding| bytes.checked_add(padding)) + .ok_or_else(overflow) +} + +/// Bound hashes and flat-key null masks for the largest input batch. +pub(super) fn prepared_scratch_bytes( + max_batch_rows: usize, + num_keys: usize, + null_equality: NullEquality, +) -> Result { + let overflow = || exec_datafusion_err!("Prepared hash-join scratch size overflow"); + let hash_bytes = max_batch_rows + .checked_mul(size_of::()) + .ok_or_else(overflow)?; + if null_equality == NullEquality::NullEqualsNull { + return Ok(hash_bytes); + } + // Allow one logical null mask per key plus the combined mask. + let masks = num_keys.checked_add(1).ok_or_else(overflow)?; + let mask_bytes = max_batch_rows + .div_ceil(8) + .checked_add(64) + .and_then(|bytes| bytes.checked_mul(masks)) + .ok_or_else(overflow)?; + hash_bytes.checked_add(mask_bytes).ok_or_else(overflow) +} + +/// Release the unused copy allowance only after validating the retained size. +pub(super) fn reconcile_prepared_copy_reservation( + reservation: &MemoryReservation, + input_bytes: usize, + copy_bytes: usize, + retained: usize, +) -> Result<()> { + let allowance = input_bytes.checked_add(copy_bytes).ok_or_else(|| { + exec_datafusion_err!("Prepared hash-join copy allowance size overflow") + })?; + if retained > allowance { + return internal_err!( + "Prepared hash-join concat exceeded its admitted copy bound" + ); + } + reservation.shrink(allowance - retained); + Ok(()) +} + +impl HashJoinExecBuilder { + /// Attach a fully prepared build to a fresh, compatible join execution. + /// + /// [`Self::build`] checks schema, keys, null equality, and supported join modes. + /// Attaching resets execution state and replaces the unused left subtree with + /// an empty schema placeholder. Residual predicates remain consumer-local. + /// + /// # Caller contract + /// + /// * Choose the correct build snapshot; compatibility checks do not verify + /// input identity. See [`HashJoinExec::prepare_build`]. + /// * Supply fresh dynamic-filter expressions and probe plans for independent + /// executions. An existing expression handle is preserved for its probe + /// consumers, while the build-report accumulator is reset. + /// * Retain a prepared-build lease if a dynamic-filter expression outlives + /// this plan. Its array/map references preserve allocation lifetime, but + /// the lease is what preserves the corresponding memory charge. + /// * Attach after physical optimizations that rewrite the build child. + /// Probe-only rewrites must preserve the attached join's `left()`; replacing + /// it or making incompatible key, mode, or join-type changes fails. + pub fn with_prepared_build(mut self, prepared: Arc) -> Self { + // This removes the ignored child's ordering/equivalences and preserves + // its identity through plan resets. + self.exec.left = Arc::new(crate::empty::EmptyExec::new(self.exec.left.schema())); + self.exec.prepared_build = Some(prepared); + self.preserve_properties = false; + self + } +} + +/// Validate eligibility and return build-key indices in join-key order. +/// INNER joins need no shared build-match bitmap. +fn prepared_key_indices(join: &HashJoinExec) -> Result> { + if join.join_type != JoinType::Inner + || join.mode != PartitionMode::CollectLeft + || join.null_aware + { + return plan_err!("Prepared hash-join builds require a CollectLeft INNER join"); + } + let left = join.left.schema(); + let right = join.right.schema(); + if left.fields().iter().any(|field| { + let ty = field.data_type(); + ty.primitive_width().is_none() + && !matches!( + ty, + DataType::Null + | DataType::Boolean + | DataType::FixedSizeBinary(_) + | DataType::Utf8 + ) + }) { + return plan_err!( + "Prepared hash-join builds require fixed-width or UTF-8 build columns" + ); + } + join.on + .iter() + .map(|(l, r)| { + let Some(l) = l.downcast_ref::() else { + return plan_err!("Prepared hash-join builds require direct column keys"); + }; + let Some(r) = r.downcast_ref::() else { + return plan_err!("Prepared hash-join builds require direct column keys"); + }; + if l.data_type(&left)? != r.data_type(&right)? { + return plan_err!("Prepared hash-join key types must match"); + } + Ok(l.index()) + }) + .collect() +} + +#[cfg(test)] +mod tests; diff --git a/datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs b/datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs new file mode 100644 index 0000000000000..e83f953a4c9f0 --- /dev/null +++ b/datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs @@ -0,0 +1,905 @@ +// 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. + +//! Regressions for immutable preparation, task isolation and buffer ownership. + +use super::*; +use crate::{ + common, + display::DisplayableExecutionPlan, + empty::EmptyExec, + filter::FilterExec, + memory::MemoryStream, + sorts::sort::SortExec, + statistics::{StatisticsArgs, StatisticsContext}, + stream::RecordBatchStreamAdapter, + test::TestMemoryExec, +}; +use arrow::array::{ + BooleanArray, FixedSizeBinaryArray, Int64Array, NullArray, StringArray, + StringViewArray, StructArray, +}; +use arrow::buffer::Buffer; +use arrow::compute::kernels::sort::SortOptions; +use arrow_schema::Field; +use datafusion_common::test_util::batches_to_sort_string; +use datafusion_common::{DataFusionError, assert_batches_eq}; +use datafusion_common_runtime::SpawnedTask; +use datafusion_execution::memory_pool::GreedyMemoryPool; +use datafusion_expr::Operator; +use datafusion_physical_expr::expressions::BinaryExpr; +use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr}; +use futures::{StreamExt, stream}; + +fn batch(keys: Vec>) -> RecordBatch { + let rows = keys.len(); + RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int64, true), + Field::new("payload", DataType::Int64, false), + ])), + vec![ + Arc::new(Int64Array::from(keys)), + Arc::new(Int64Array::from_iter_values( + (0..rows).map(|i| i as i64 + 10), + )), + ], + ) + .unwrap() +} + +fn join(build_schema: SchemaRef, probe: RecordBatch) -> Result { + let right = + TestMemoryExec::try_new_exec(&[vec![probe.clone()]], probe.schema(), None)?; + HashJoinExecBuilder::new( + Arc::new(EmptyExec::new(build_schema)), + right, + vec![( + Arc::new(Column::new("key", 0)), + Arc::new(Column::new("key", 0)), + )], + JoinType::Inner, + ) + .with_partition_mode(PartitionMode::CollectLeft) + .build() +} + +fn input(batches: Vec, schema: SchemaRef) -> SendableRecordBatchStream { + Box::pin(MemoryStream::try_new(batches, schema, None).unwrap()) +} + +async fn prepare( + join: &HashJoinExec, + batches: Vec, + pool: Arc, +) -> Result> { + join.prepare_build( + input(batches, join.left().schema()), + pool, + Arc::new(ConfigOptions::default()), + ) + .await +} + +fn buffer_addresses(arrays: &[ArrayRef]) -> Vec { + arrays + .iter() + .flat_map(|array| { + let data = array.to_data(); + data.buffers() + .iter() + .chain(data.nulls().map(|nulls| nulls.inner().inner())) + .filter(|buffer| buffer.capacity() != 0) + .map(|buffer| buffer.data_ptr().as_ptr() as usize) + .collect::>() + }) + .collect() +} + +async fn run(join: &HashJoinExec) -> Result> { + let partitions = (0..join.properties().output_partitioning().partition_count()).map( + |partition| async move { + common::collect(join.execute(partition, Arc::new(TaskContext::default()))?) + .await + }, + ); + Ok(futures::future::try_join_all(partitions) + .await? + .into_iter() + .flatten() + .collect()) +} + +fn with_probe_filter(join: HashJoinExec) -> Result { + let filter = HashJoinExec::create_dynamic_filter(&join.on); + let probe = Arc::new(FilterExec::try_new( + Arc::clone(&filter) as _, + Arc::clone(join.right()), + )?); + join.builder() + .with_new_children(vec![Arc::clone(join.left()), probe])? + .with_dynamic_filter(Some(HashJoinExecDynamicFilter { + filter, + build_accumulator: OnceLock::new(), + })) + .build() +} + +fn payload_filter(op: Operator) -> JoinFilter { + JoinFilter::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("build_payload", 0)), + op, + Arc::new(Column::new("probe_payload", 1)), + )), + vec![ + ColumnIndex { + index: 1, + side: JoinSide::Left, + }, + ColumnIndex { + index: 1, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![ + Field::new("build_payload", DataType::Int64, true), + Field::new("probe_payload", DataType::Int64, true), + ])), + ) +} + +#[tokio::test] +async fn prepared_build_placeholder_properties_and_reset() -> Result<()> { + let ascending = |name, index| PhysicalSortExpr { + expr: Arc::new(Column::new(name, index)), + options: SortOptions::default(), + }; + let source_schema = + Arc::new(Schema::new(vec![Field::new("key", DataType::Int64, false)])); + let placeholder = ProjectionExec::try_new( + vec![ + ( + Arc::new(Column::new("key", 0)) as PhysicalExprRef, + "key".into(), + ), + ( + Arc::new(Column::new("key", 0)) as PhysicalExprRef, + "payload".into(), + ), + ], + Arc::new(EmptyExec::new(source_schema)), + )?; + let left: Arc = Arc::new(SortExec::new( + LexOrdering::new(vec![ascending("payload", 1)]).unwrap(), + Arc::new(placeholder), + )); + let build = RecordBatch::try_new( + left.schema(), + vec![ + Arc::new(Int64Array::from(vec![1, 1])), + Arc::new(Int64Array::from(vec![2, 1])), + ], + )?; + let probe_schema = + Arc::new(Schema::new(vec![Field::new("key", DataType::Int64, false)])); + let probe = RecordBatch::try_new( + Arc::clone(&probe_schema), + vec![Arc::new(Int64Array::from(vec![1]))], + )?; + let right: Arc = Arc::new(SortExec::new( + LexOrdering::new(vec![ascending("key", 0)]).unwrap(), + TestMemoryExec::try_new_exec(&[vec![probe]], probe_schema, None)?, + )); + let base = HashJoinExecBuilder::new( + left, + right, + vec![( + Arc::new(Column::new("key", 0)), + Arc::new(Column::new("key", 0)), + )], + JoinType::Inner, + ) + .with_partition_mode(PartitionMode::CollectLeft) + .build()?; + let build_key: PhysicalExprRef = Arc::new(Column::new("key", 0)); + let build_payload: PhysicalExprRef = Arc::new(Column::new("payload", 1)); + let asserted_order = [ascending("key", 2), ascending("payload", 1)]; + let base_properties = base.properties().equivalence_properties(); + assert!( + base_properties + .eq_group() + .exprs_equal(&build_key, &build_payload) + ); + assert!(base_properties.ordering_satisfy(asserted_order.clone())?); + + let prepared = + prepare(&base, vec![build], Arc::new(GreedyMemoryPool::new(1 << 20))).await?; + let attached = base.builder().with_prepared_build(prepared).build()?; + assert!( + DisplayableExecutionPlan::new(&attached) + .indent(false) + .to_string() + .contains("prepared_build=2 rows") + ); + let properties = attached.properties().equivalence_properties(); + assert!( + !properties + .eq_group() + .exprs_equal(&build_key, &build_payload) + ); + assert!(properties.ordering_satisfy([ascending("key", 2)])?); + assert!(!properties.ordering_satisfy(asserted_order)?); + let stats = StatisticsContext::new().compute(&attached, &StatisticsArgs::new())?; + assert_ne!(stats.num_rows.get_value(), Some(&0)); + let output = run(&attached).await?; + assert_eq!(output.iter().map(RecordBatch::num_rows).sum::(), 2); + // Resetting the unused projection/sort must retain the prepared rows. + let reset = crate::execution_plan::reset_plan_states(Arc::new(attached))?; + assert_eq!( + common::collect(reset.execute(0, Arc::new(TaskContext::default()))?).await?, + output + ); + Ok(()) +} + +#[tokio::test] +async fn prepared_build_projection_pushdown_preserves_rows() -> Result<()> { + let build = batch(vec![Some(1)]); + let base = join(build.schema(), build.clone())?; + let prepared = + prepare(&base, vec![build], Arc::new(GreedyMemoryPool::new(1 << 20))).await?; + let task = Arc::new(base.builder().with_prepared_build(prepared).build()?); + let projection = ProjectionExec::try_new( + vec![ + ( + Arc::new(Column::new("key", 0)) as PhysicalExprRef, + "build_key".to_owned(), + ), + ( + Arc::new(Column::new("key", 2)) as PhysicalExprRef, + "probe_key".to_owned(), + ), + ], + Arc::clone(&task) as _, + )?; + let transformed = task + .try_swapping_with_projection(&projection)? + .unwrap_or_else(|| Arc::new(projection)); + let output = + common::collect(transformed.execute(0, Arc::new(TaskContext::default()))?) + .await?; + assert_batches_eq!( + [ + "+-----------+-----------+", + "| build_key | probe_key |", + "+-----------+-----------+", + "| 1 | 1 |", + "+-----------+-----------+", + ], + &output + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn prepared_build_reuses_data_with_independent_dynamic_filters() -> Result<()> { + let build = batch(vec![Some(1), Some(1), None, Some(-3)]); + let base = join(build.schema(), batch(vec![Some(1)]))?; + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 20)); + let prepared = prepare(&base, vec![build.clone()], Arc::clone(&pool)).await?; + let bytes = pool.reserved(); + assert_eq!(prepared.num_rows(), 4); + assert_eq!(bytes, prepared.reserved_bytes()); + assert!(bytes > 0); + + let mut plans = Vec::new(); + let mut expected = Vec::new(); + for (op, probe) in [ + (Operator::Gt, batch(vec![Some(1), Some(2), None])), + (Operator::Lt, batch(vec![Some(1), Some(1)])), + ] { + let right = TestMemoryExec::try_new_exec( + &[ + vec![probe.slice(0, 1)], + vec![probe.slice(1, probe.num_rows() - 1)], + ], + probe.schema(), + None, + )?; + let task = base + .builder() + .with_new_children(vec![Arc::clone(base.left()), right])? + .with_filter(Some(payload_filter(op))) + .build()?; + let ordinary = task + .builder() + .reset_state() + .with_new_children(vec![ + TestMemoryExec::try_new_exec( + &[vec![build.clone()]], + build.schema(), + None, + )?, + Arc::clone(task.right()), + ])? + .build()?; + expected.push(run(&ordinary).await?); + let task = with_probe_filter(task)? + .builder() + .with_prepared_build(Arc::clone(&prepared)) + .build()?; + plans.push(Arc::new(task)); + } + let cancelled = base + .builder() + .with_prepared_build(Arc::clone(&prepared)) + .build()?; + drop(cancelled.execute(0, Arc::new(TaskContext::default()))?); + drop(cancelled); + assert_eq!(pool.reserved(), bytes); + let start = Arc::new(tokio::sync::Barrier::new(plans.len())); + let tasks = plans.iter().map(|plan| { + let plan = Arc::clone(plan); + let start = Arc::clone(&start); + SpawnedTask::spawn(async move { + start.wait().await; + run(&plan).await + }) + }); + let output = futures::future::join_all(tasks) + .await + .into_iter() + .map(|task| task.expect("probe task panicked")) + .collect::>>()?; + for (index, plan) in plans.iter().enumerate() { + assert_eq!( + batches_to_sort_string(&output[index]), + batches_to_sort_string(&expected[index]) + ); + let metrics = plan.metrics().unwrap(); + assert_eq!( + metrics.sum_by_name("build_input_rows").unwrap().as_usize(), + 0 + ); + assert_eq!(metrics.output_rows(), Some(1)); + let dynamic = plan.dynamic_filter.as_ref().unwrap(); + assert!(futures::poll!(Box::pin(dynamic.filter.wait_complete())).is_ready()); + } + assert_eq!(pool.reserved(), bytes); + drop(plans); + drop(prepared); + assert_eq!(pool.reserved(), 0); + Ok(()) +} + +#[tokio::test] +async fn prepared_build_errors_and_cancellation_release_reservations() -> Result<()> { + let build = batch(vec![Some(1), Some(2)]); + let base = join(build.schema(), build.clone())?; + let exhausted: Arc = Arc::new(GreedyMemoryPool::new(0)); + assert!(matches!( + prepare(&base, vec![build.clone()], Arc::clone(&exhausted)).await, + Err(DataFusionError::ResourcesExhausted(_)) + )); + assert_eq!(exhausted.reserved(), 0); + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 20)); + let failing = stream::iter(vec![ + Ok(build.clone()), + datafusion_common::exec_err!("producer failed"), + ]); + let result = base + .prepare_build( + Box::pin(RecordBatchStreamAdapter::new(build.schema(), failing)), + Arc::clone(&pool), + Arc::new(ConfigOptions::default()), + ) + .await; + assert!(result.is_err()); + assert_eq!(pool.reserved(), 0); + let pending = stream::iter(vec![Ok(build.clone())]).chain(stream::pending()); + let mut preparation = Box::pin(base.prepare_build( + Box::pin(RecordBatchStreamAdapter::new(build.schema(), pending)), + Arc::clone(&pool), + Arc::new(ConfigOptions::default()), + )); + assert!(futures::poll!(&mut preparation).is_pending()); + assert!(pool.reserved() > 0); + drop(preparation); + assert_eq!(pool.reserved(), 0); + Ok(()) +} + +#[tokio::test] +async fn prepared_build_shares_single_batch_but_owns_output() -> Result<()> { + for array in [ + Arc::new(Int64Array::from(vec![ + Some(0), + Some(1), + None, + Some(2), + Some(2), + Some(3), + ])) as ArrayRef, + Arc::new(StringArray::from(vec![ + Some("prefix"), + Some("a"), + None, + Some("b"), + Some("b"), + Some("suffix"), + ])), + ] { + let build = RecordBatch::try_from_iter([("key", array.slice(1, 4))])?; + let probe = RecordBatch::try_from_iter([("key", array.slice(1, 3))])?; + let base = join(build.schema(), probe)? + .builder() + .with_null_equality(NullEquality::NullEqualsNull) + .build()?; + for sources in [ + vec![build.clone()], + vec![build.slice(0, 2), build.slice(2, 2)], + ] { + let single_batch = sources.len() == 1; + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 20)); + let prepared = prepare(&base, sources, Arc::clone(&pool)).await?; + let addresses = buffer_addresses(prepared.build.batch.columns()); + if single_batch { + assert_eq!(addresses, buffer_addresses(build.columns())); + } + let bytes = prepared.reserved_bytes(); + let task = base.builder().with_prepared_build(prepared).build()?; + let output = run(&task).await?; + assert_eq!(output.iter().map(RecordBatch::num_rows).sum::(), 4); + assert!( + output + .iter() + .flat_map(|batch| buffer_addresses(&batch.columns()[..1])) + .all(|address| !addresses.contains(&address)) + ); + assert_eq!(pool.reserved(), bytes); + drop(task); + assert_eq!(pool.reserved(), 0); + for batch in output { + assert_eq!(batch.column(0).to_data(), batch.column(1).to_data()); + } + } + } + Ok(()) +} + +#[tokio::test] +async fn prepared_build_validates_descriptor_and_builder_changes() -> Result<()> { + let build = batch(vec![Some(1), Some(2)]); + let base = with_probe_filter(join(build.schema(), build.clone())?)?; + let prepared = prepare( + &base, + vec![build.clone()], + Arc::new(GreedyMemoryPool::new(1 << 20)), + ) + .await?; + let keys: JoinOn = vec![( + Arc::new(Column::new("payload", 1)), + Arc::new(Column::new("payload", 1)), + )]; + for builder in [ + base.builder() + .with_null_equality(NullEquality::NullEqualsNull), + base.builder().with_type(JoinType::Right), + base.builder() + .with_partition_mode(PartitionMode::Partitioned), + base.builder().with_on(keys), + ] { + assert!( + builder + .with_prepared_build(Arc::clone(&prepared)) + .build() + .is_err() + ); + } + let attached = Arc::new(base.builder().with_prepared_build(prepared).build()?); + // Changing only the probe key must not reuse its existing dynamic filter. + assert!( + attached + .builder() + .with_on(vec![( + Arc::new(Column::new("key", 0)), + Arc::new(Column::new("payload", 1)), + )]) + .build() + .is_err() + ); + let replacement = Arc::new(EmptyExec::new(build.schema())) as Arc; + assert!( + Arc::clone(&attached) + .replace_children( + vec![replacement, Arc::clone(attached.right())], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + .is_err() + ); + let right = TestMemoryExec::try_new_exec( + &[vec![batch(vec![Some(1), Some(3)])]], + build.schema(), + None, + )?; + let new_probe = attached + .builder() + .with_new_children(vec![Arc::clone(attached.left()), right])? + .build()?; + assert_eq!( + run(&new_probe) + .await? + .iter() + .map(RecordBatch::num_rows) + .sum::(), + 1 + ); + Ok(()) +} + +#[tokio::test] +async fn prepared_build_revalidates_public_keys_at_execution() -> Result<()> { + let build = batch(vec![Some(1), Some(2)]); + let base = join(build.schema(), build.clone())?; + let keys: JoinOn = vec![( + Arc::new(Column::new("payload", 1)), + Arc::new(Column::new("payload", 1)), + )]; + let ordinary = base + .builder() + .with_on(keys.clone()) + .with_new_children(vec![ + TestMemoryExec::try_new_exec(&[vec![build.clone()]], build.schema(), None)?, + Arc::clone(base.right()), + ])? + .build()?; + assert_batches_eq!( + [ + "+-----+---------+-----+---------+", + "| key | payload | key | payload |", + "+-----+---------+-----+---------+", + "| 1 | 10 | 1 | 10 |", + "| 2 | 11 | 2 | 11 |", + "+-----+---------+-----+---------+", + ], + &run(&ordinary).await? + ); + + let prepared = + prepare(&base, vec![build], Arc::new(GreedyMemoryPool::new(1 << 20))).await?; + let mut attached = base.builder().with_prepared_build(prepared).build()?; + attached.on = keys; + let error = run(&attached) + .await + .expect_err("changed public keys must not probe the old prepared table"); + assert!(matches!(error, DataFusionError::Plan(_)), "{error}"); + assert!( + error + .to_string() + .contains("does not match schema, keys or null equality") + ); + Ok(()) +} + +#[tokio::test] +async fn prepared_build_rejects_view_payload_before_consumption() -> Result<()> { + let build = RecordBatch::try_from_iter([ + ("key", Arc::new(Int64Array::from(vec![1])) as ArrayRef), + ("payload", Arc::new(StringViewArray::from(vec!["x"]))), + ])?; + let base = join(build.schema(), batch(vec![Some(1)]))?; + let never_polled = + stream::poll_fn(|_| -> std::task::Poll>> { + panic!("unsupported input must not be polled") + }); + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 20)); + assert!( + base.prepare_build( + Box::pin(RecordBatchStreamAdapter::new(build.schema(), never_polled)), + Arc::clone(&pool), + Arc::new(ConfigOptions::default()), + ) + .await + .is_err() + ); + assert_eq!(pool.reserved(), 0); + Ok(()) +} + +#[tokio::test] +async fn prepared_build_matches_collect_left() -> Result<()> { + let split = |batch: &RecordBatch| { + (0..batch.num_rows()) + .step_by(2) + .map(|offset| batch.slice(offset, 2.min(batch.num_rows() - offset))) + .collect::>() + }; + for (build_keys, probe_keys) in [ + (vec![], vec![Some(1)]), + (vec![Some(1)], vec![]), + (vec![None, None], vec![None, Some(1)]), + ( + vec![None, Some(0), Some(0), Some(-4096), Some(4096)], + vec![Some(0), None, Some(4096), Some(2)], + ), + ] { + let build = batch(build_keys); + let probe = batch(probe_keys); + let sources = split(&build); + for nulls in [ + NullEquality::NullEqualsNothing, + NullEquality::NullEqualsNull, + ] { + let base = join(build.schema(), probe.clone())? + .builder() + .with_new_children(vec![ + TestMemoryExec::try_new_exec( + std::slice::from_ref(&sources), + build.schema(), + None, + )?, + TestMemoryExec::try_new_exec(&[split(&probe)], probe.schema(), None)?, + ])? + .with_null_equality(nulls) + .build()?; + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 20)); + let prepared = prepare(&base, sources.clone(), Arc::clone(&pool)).await?; + for batch_size in [1, 7] { + let context = Arc::new( + TaskContext::default().with_session_config( + datafusion_execution::config::SessionConfig::new() + .with_batch_size(batch_size), + ), + ); + let ordinary = base.builder().reset_state().build()?; + let attached = base + .builder() + .with_prepared_build(Arc::clone(&prepared)) + .build()?; + let mut expected = + common::collect(ordinary.execute(0, Arc::clone(&context))?).await?; + let mut actual = common::collect(attached.execute(0, context)?).await?; + assert!(actual.iter().all(|batch| batch.num_rows() <= batch_size)); + expected.retain(|batch| batch.num_rows() != 0); + actual.retain(|batch| batch.num_rows() != 0); + assert_eq!( + batches_to_sort_string(&actual), + batches_to_sort_string(&expected) + ); + } + drop(prepared); + assert_eq!(pool.reserved(), 0); + } + } + Ok(()) +} + +#[tokio::test] +async fn prepared_concat_admission_with_small_batches() -> Result<()> { + let rows = 65_536; + let build = batch((0..rows).map(|i| Some(i as i64 * 32)).collect()); + let base = join(build.schema(), batch(vec![Some(0)]))?; + for batch_rows in [32, 64, 8192] { + let batches = (0..rows) + .step_by(batch_rows) + .map(|offset| build.slice(offset, batch_rows)) + .collect(); + // Small input batches should not multiply concatenation's alignment charge. + let pool: Arc = + Arc::new(GreedyMemoryPool::new(9 * 1024 * 1024 / 2)); + let prepared = prepare(&base, batches, Arc::clone(&pool)).await?; + assert_eq!(prepared.num_rows(), rows); + assert_eq!(pool.reserved(), prepared.reserved_bytes()); + drop(prepared); + assert_eq!(pool.reserved(), 0); + } + Ok(()) +} + +#[test] +fn prepared_concat_admits_validity_for_non_nullable_input_arrays() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "key", + DataType::Boolean, + true, + )])); + // A single null in the second batch creates validity for every row in the first. + // Boolean values make that extra allocation as large as the values themselves. + let batches = [ + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(BooleanArray::from(vec![true; 100_000]))], + )?, + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(BooleanArray::from(vec![None]))], + )?, + ]; + assert!(batches[0].column(0).nulls().is_none()); + assert!(batches[1].column(0).nulls().is_some()); + let admitted = prepared_copy_bytes(&batches)?; + let copied = concat_batches(&schema, &batches)?; + assert_eq!(copied.column(0).null_count(), 1); + let allocated = copied.column(0).get_buffer_memory_size(); + assert!( + admitted >= allocated, + "copy admission {admitted} is smaller than allocated capacity {allocated}" + ); + Ok(()) +} + +#[tokio::test] +async fn prepared_inlist_shares_accounted_build_buffers() -> Result<()> { + let build = batch(vec![Some(1), Some(1_000_000), None, Some(1)]); + for keys in [1, 2] { + let on = (0..keys) + .map(|index| { + let name = build.schema().field(index).name().clone(); + let key: PhysicalExprRef = Arc::new(Column::new(&name, index)); + (Arc::clone(&key), key) + }) + .collect(); + let base = join(build.schema(), build.clone())? + .builder() + .with_on(on) + .build()?; + for batches in [ + vec![build.clone()], + vec![build.slice(0, 2), build.slice(2, 2)], + ] { + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 20)); + let prepared = prepare(&base, batches.clone(), Arc::clone(&pool)).await?; + let PushdownStrategy::InList(values) = &prepared.build.membership else { + panic!("expected IN-list membership"); + }; + let columns = prepared.build.batch.columns(); + if keys == 1 { + assert!(Arc::ptr_eq(values, &columns[0])); + } else { + let values = values.as_any().downcast_ref::().unwrap(); + for (value, column) in values.columns().iter().zip(columns) { + assert!(Arc::ptr_eq(value, column)); + } + } + let bytes = prepared.reserved_bytes(); + let mut counter = RecordBatchMemoryCounter::new(); + let batch_bytes = counter.count_batch(&prepared.build.batch); + assert!(bytes > batch_bytes, "hash storage must also be reserved"); + let membership_batch = + RecordBatch::try_from_iter([("membership", Arc::clone(values))])?; + assert_eq!(counter.count_batch(&membership_batch), 0); + drop(membership_batch); + + // Disabling IN-list membership changes no retained key buffers. + let mut config = ConfigOptions::default(); + config.optimizer.hash_join_inlist_pushdown_max_size = 0; + let map = base + .prepare_build( + input(batches, build.schema()), + Arc::clone(&pool), + Arc::new(config), + ) + .await?; + assert!(matches!(map.build.membership, PushdownStrategy::Map(_))); + assert_eq!(map.reserved_bytes(), bytes); + assert_eq!(pool.reserved(), 2 * bytes); + drop(map); + let lease = Arc::clone(&prepared); + drop(prepared); + assert_eq!(pool.reserved(), bytes); + drop(lease); + assert_eq!(pool.reserved(), 0); + } + } + Ok(()) +} + +#[test] +fn prepared_concat_rejects_under_admission_without_shrinking() -> Result<()> { + let pool: Arc = Arc::new(GreedyMemoryPool::new(1024)); + let reservation = MemoryConsumer::new("prepared concat test").register(&pool); + // 128 bytes of map storage plus 128 bytes of input/copy allowance. + reservation.try_grow(256)?; + let error = + reconcile_prepared_copy_reservation(&reservation, 64, 64, 129).unwrap_err(); + assert!(matches!(error, DataFusionError::Internal(_))); + assert!( + error + .to_string() + .contains("exceeded its admitted copy bound") + ); + assert_eq!(reservation.size(), 256); + assert!(reconcile_prepared_copy_reservation(&reservation, usize::MAX, 1, 0).is_err()); + assert_eq!(reservation.size(), 256); + reconcile_prepared_copy_reservation(&reservation, 64, 64, 128)?; + assert_eq!(reservation.size(), 256); + reconcile_prepared_copy_reservation(&reservation, 64, 64, 80)?; + assert_eq!(reservation.size(), 208); + drop(reservation); + assert_eq!(pool.reserved(), 0); + Ok(()) +} + +#[test] +fn prepared_sizing_rejects_overflow() -> Result<()> { + use NullEquality::{NullEqualsNothing, NullEqualsNull}; + for (rows, keys, equality) in [ + (usize::MAX / 8 + 1, 1, NullEqualsNull), + (0, usize::MAX, NullEqualsNothing), + (0, usize::MAX / 64, NullEqualsNothing), + (usize::MAX / 8, 0, NullEqualsNothing), + ] { + let error = prepared_scratch_bytes(rows, keys, equality).unwrap_err(); + assert!(error.to_string().contains("scratch size overflow")); + } + assert!(prepared_scratch_bytes(usize::MAX / 8, 1, NullEqualsNull).is_ok()); + + // Null arrays can describe huge row counts without allocating data buffers. + // Exercise both per-batch validity multiplication and cumulative addition. + let array: ArrayRef = Arc::new(NullArray::new(usize::MAX)); + for (columns, batches) in [(8, 1), (4, 2)] { + let batch = RecordBatch::try_from_iter( + (0..columns).map(|i| (format!("key{i}"), Arc::clone(&array))), + )?; + let error = prepared_copy_bytes(&vec![batch; batches]).unwrap_err(); + assert!(error.to_string().contains("copy size overflow")); + } + Ok(()) +} + +#[tokio::test] +async fn prepared_byte_keys_use_hash_membership() -> Result<()> { + for array in [ + Arc::new(StringArray::from(vec!["aaaa", "bbbb"])) as ArrayRef, + Arc::new(FixedSizeBinaryArray::try_new( + 4, + Buffer::from_slice_ref(b"aaaabbbb"), + None, + )?), + ] { + let build = RecordBatch::try_from_iter([("key", array.slice(0, 1))])?; + let probe = RecordBatch::try_from_iter([("key", array)])?; + let base = join(build.schema(), probe)?; + let mut options = ConfigOptions::default(); + options.optimizer.hash_join_inlist_pushdown_max_size = 1024; + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 20)); + let prepared = base + .prepare_build( + input(vec![build.clone()], build.schema()), + Arc::clone(&pool), + Arc::new(options), + ) + .await?; + assert!(prepared.build.bounds.is_none()); + assert!(matches!( + &prepared.build.membership, + PushdownStrategy::Map(_) + )); + let task = with_probe_filter(base)? + .builder() + .with_prepared_build(prepared) + .build()?; + let output = run(&task).await?; + assert_eq!(output.iter().map(RecordBatch::num_rows).sum::(), 1); + let filter = &task.dynamic_filter.as_ref().unwrap().filter; + assert!(futures::poll!(Box::pin(filter.wait_complete())).is_ready()); + drop(task); + assert_eq!(pool.reserved(), 0); + } + Ok(()) +} diff --git a/datafusion/physical-plan/src/joins/hash_join/mod.rs b/datafusion/physical-plan/src/joins/hash_join/mod.rs index b915802ea4015..4687d8545ec69 100644 --- a/datafusion/physical-plan/src/joins/hash_join/mod.rs +++ b/datafusion/physical-plan/src/joins/hash_join/mod.rs @@ -17,7 +17,7 @@ //! [`HashJoinExec`] Partitioned Hash Join Operator -pub use exec::{HashJoinExec, HashJoinExecBuilder}; +pub use exec::{HashJoinExec, HashJoinExecBuilder, PreparedHashJoinBuild}; pub use partitioned_hash_eval::{HashExpr, HashTableLookupExpr, SeededRandomState}; mod exec; diff --git a/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs b/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs index 60a25fc2efcff..fdd66d0e54ae2 100644 --- a/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs +++ b/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs @@ -309,26 +309,15 @@ impl Hash for HashTableLookupExpr { self.on_columns.dyn_hash(state); self.description.hash(state); self.random_state.seed().hash(state); - // Note that we compare hash_map by pointer equality. - // Actually comparing the contents of the hash maps would be expensive. - // The way these hash maps are used in actuality is that HashJoinExec creates - // one per partition per query execution, thus it is never possible for two different - // hash maps to have the same content in practice. - // Theoretically this is a public API and users could create identical hash maps, - // but that seems unlikely and not worth paying the cost of deep comparison all the time. + // Hash immutable map identity, consistently with pointer equality below. Arc::as_ptr(&self.map).hash(state); } } impl PartialEq for HashTableLookupExpr { fn eq(&self, other: &Self) -> bool { - // Note that we compare hash_map by pointer equality. - // Actually comparing the contents of the hash maps would be expensive. - // The way these hash maps are used in actuality is that HashJoinExec creates - // one per partition per query execution, thus it is never possible for two different - // hash maps to have the same content in practice. - // Theoretically this is a public API and users could create identical hash maps, - // but that seems unlikely and not worth paying the cost of deep comparison all the time. + // Sharing an immutable map implies equal contents, including across + // prepared-build consumers. Distinct maps need no expensive deep comparison. self.on_columns == other.on_columns && self.description == other.description && self.random_state.seed() == other.random_state.seed() diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs index 686939537e73e..fc2bcd82e775d 100644 --- a/datafusion/physical-plan/src/joins/hash_join/stream.rs +++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs @@ -556,8 +556,8 @@ impl HashJoinStream { let pushdown = left_data.membership().clone(); let bounds = left_data - .bounds - .clone() + .bounds() + .cloned() .unwrap_or_else(|| PartitionBounds::new(vec![])); // Arrow tracks null counts per array, so this costs no data scan. let keys_have_null = left_data diff --git a/datafusion/physical-plan/src/joins/mod.rs b/datafusion/physical-plan/src/joins/mod.rs index e4f7e2e123e0e..39b80cd3d0cc0 100644 --- a/datafusion/physical-plan/src/joins/mod.rs +++ b/datafusion/physical-plan/src/joins/mod.rs @@ -21,7 +21,8 @@ use arrow::array::BooleanBufferBuilder; pub use cross_join::CrossJoinExec; use datafusion_physical_expr::PhysicalExprRef; pub use hash_join::{ - HashExpr, HashJoinExec, HashJoinExecBuilder, HashTableLookupExpr, SeededRandomState, + HashExpr, HashJoinExec, HashJoinExecBuilder, HashTableLookupExpr, + PreparedHashJoinBuild, SeededRandomState, }; pub use nested_loop_join::{NestedLoopJoinExec, NestedLoopJoinExecBuilder}; use parking_lot::Mutex; diff --git a/datafusion/proto/tests/cases/plans/joins.rs b/datafusion/proto/tests/cases/plans/joins.rs index 941e8832952c6..f048772bc4e31 100644 --- a/datafusion/proto/tests/cases/plans/joins.rs +++ b/datafusion/proto/tests/cases/plans/joins.rs @@ -456,3 +456,49 @@ async fn roundtrip_logical_plan_sort_merge_join() -> Result<()> { let plan = ctx.sql(query).await?.create_physical_plan().await?; roundtrip_test(plan) } + +#[tokio::test] +async fn prepared_hash_join_cannot_be_serialized_without_its_build() -> Result<()> { + use datafusion::arrow::array::Int64Array; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::physical_plan::memory::MemoryStream; + use datafusion_common::config::ConfigOptions; + use datafusion_execution::memory_pool::GreedyMemoryPool; + + let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Int64, false)])); + let join = HashJoinExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&schema))), + Arc::new(EmptyExec::new(Arc::clone(&schema))), + vec![( + Arc::new(Column::new("col", 0)), + Arc::new(Column::new("col", 0)), + )], + None, + &JoinType::Inner, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + false, + )?; + roundtrip_test(Arc::new(join.builder().build()?))?; + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(vec![1_i64]))], + )?; + let prepared = join + .prepare_build( + Box::pin(MemoryStream::try_new(vec![batch], schema, None)?), + Arc::new(GreedyMemoryPool::new(1 << 20)), + Arc::new(ConfigOptions::default()), + ) + .await?; + let attached = join.builder().with_prepared_build(prepared).build()?; + let error = roundtrip_test(Arc::new(attached)) + .expect_err("serializing the empty build placeholder would lose prepared rows"); + assert!(matches!( + error, + datafusion_common::DataFusionError::NotImplemented(_) + )); + assert!(error.to_string().contains("prepared build"), "{error}"); + Ok(()) +} From 2e16b818c0915242acd6cd14de06d4124c284f24 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 6 Nov 2026 21:21:38 +0800 Subject: [PATCH 07/19] fix: adapt #25491 prepared joins to fork probe state Use the older fork's execution-local probe counters and NULL-aware atomic flags instead of upstream-parent fields unavailable here; remove the absent collector argument. Prepared eligibility remains non-null-aware CollectLeft INNER only. Preserve the ordinary join's existing semantics. Document the actual fork ownership and probe state, keep initializer/drop ordering aligned, and relocate the unchanged #25508 test after test imports. No additional local regression tests are included. Upstream PR: https://github.com/apache/datafusion/pull/25491 Upstream commit: acf5c89e6454469c763c0fb5539e5a294d3439ed This is a local fork adaptation. Compilation will be checked on the final adapted series; no per-intermediate test-pass claim is made. Signed-off-by: discord9 --- .../physical-plan/src/joins/hash_join/exec.rs | 91 +++++++++---------- .../src/joins/hash_join/exec/prepared.rs | 9 +- 2 files changed, 48 insertions(+), 52 deletions(-) diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index c23e1d81b2016..2f5d8abf91269 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -209,21 +209,19 @@ struct JoinBuildData { bounds: Option, /// IN-list values or a hash-table reference used for filter pushdown. membership: PushdownStrategy, - // Keep the reservation after the allocations it accounts for. + // Declared last so accounted-for build allocations are released first. reservation: MemoryReservation, } /// A build lease and mutable bookkeeping for one join execution. pub(super) struct JoinLeftData { - /// Shared bitmap for visited left indices within this consuming join. + /// Shared bitmap for visited build indices within this join execution. visited_indices_bitmap: SharedBitmapBuilder, - /// Counter of running probe-threads, potentially - /// able to update `visited_indices_bitmap` + /// Counter of running probe partitions that may update the bitmap. probe_threads_counter: AtomicUsize, - /// Shared atomic flag indicating if any probe partition saw data (for null-aware anti joins) - /// This is shared across all probe partitions to provide global knowledge + /// Whether any probe partition saw data (for null-aware anti joins). pub(super) probe_side_non_empty: AtomicBool, - /// Shared atomic flag indicating if any probe partition saw NULL in join keys (for null-aware anti joins) + /// Whether any probe partition saw a NULL key (for null-aware anti joins). pub(super) probe_side_has_null: AtomicBool, // Drop the reservation after the mutable state it also accounts for. build: Arc, @@ -2259,7 +2257,9 @@ impl BuildSideState { } } -/// range and density, see [`array_map_key_range`]. +/// Whether a single build key is a candidate for perfect hashing. +/// The final decision also depends on the observed range and density; see +/// [`array_map_key_range`]. fn should_collect_min_max_for_perfect_hash( on_left: &[PhysicalExprRef], schema: &SchemaRef, @@ -2647,8 +2647,7 @@ async fn collect_left_input( if prepared { drop(batches); - // Prepared keys are direct columns. IN-list arrays share these batch - // buffers, including through multi-key StructArray children. + // Prepared keys are direct columns; numeric IN-list arrays share batch buffers. let retained = RecordBatchMemoryCounter::new().count_batch(&batch); prepared::reconcile_prepared_copy_reservation( &reservation, @@ -2659,6 +2658,10 @@ async fn collect_left_input( } let data = JoinLeftData { + visited_indices_bitmap: Mutex::new(visited_indices_bitmap), + probe_threads_counter: AtomicUsize::new(probe_threads_count), + probe_side_non_empty: AtomicBool::new(false), + probe_side_has_null: AtomicBool::new(false), build: Arc::new(JoinBuildData { map, batch, @@ -2667,10 +2670,6 @@ async fn collect_left_input( membership, reservation, }), - visited_indices_bitmap: Mutex::new(visited_indices_bitmap), - probe_threads_counter: AtomicUsize::new(probe_threads_count), - probe_side_non_empty: AtomicBool::new(false), - probe_side_has_null: AtomicBool::new(false), }; Ok(data) @@ -2680,38 +2679,6 @@ async fn collect_left_input( mod tests { use super::*; - #[test] - fn hash_map_admits_row_indices_before_allocation() -> Result<()> { - use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryPool}; - - let rows = 1024; - let buckets = - estimate_memory_size::<(u32, u64)>(rows, size_of::())?; - let bytes = buckets + rows * size_of::(); - for (limit, succeeds) in [(bytes - 1, false), (bytes, true)] { - let pool: Arc = Arc::new(GreedyMemoryPool::new(limit)); - let mut reservation = MemoryConsumer::new("row indices").register(&pool); - let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); - let result = new_join_hashmap(rows, &mut reservation, &metrics); - if succeeds { - let map = result?; - assert_eq!(reservation.size(), bytes); - assert_eq!(metrics.build_mem_used.value(), bytes); - drop(map); - } else { - assert!(matches!( - result, - Err(datafusion_common::DataFusionError::ResourcesExhausted(_)) - )); - assert_eq!(reservation.size(), 0); - assert_eq!(metrics.build_mem_used.value(), 0); - } - drop(reservation); - assert_eq!(pool.reserved(), 0); - } - Ok(()) - } - fn assert_phj_used(metrics: &MetricsSet, use_phj: bool) { if use_phj { assert!( @@ -2786,6 +2753,38 @@ mod tests { use rstest::*; use rstest_reuse::*; + #[test] + fn hash_map_admits_row_indices_before_allocation() -> Result<()> { + use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryPool}; + + let rows = 1024; + let buckets = + estimate_memory_size::<(u32, u64)>(rows, size_of::())?; + let bytes = buckets + rows * size_of::(); + for (limit, succeeds) in [(bytes - 1, false), (bytes, true)] { + let pool: Arc = Arc::new(GreedyMemoryPool::new(limit)); + let mut reservation = MemoryConsumer::new("row indices").register(&pool); + let metrics = BuildProbeJoinMetrics::new(0, &ExecutionPlanMetricsSet::new()); + let result = new_join_hashmap(rows, &mut reservation, &metrics); + if succeeds { + let map = result?; + assert_eq!(reservation.size(), bytes); + assert_eq!(metrics.build_mem_used.value(), bytes); + drop(map); + } else { + assert!(matches!( + result, + Err(datafusion_common::DataFusionError::ResourcesExhausted(_)) + )); + assert_eq!(reservation.size(), 0); + assert_eq!(metrics.build_mem_used.value(), 0); + } + drop(reservation); + assert_eq!(pool.reserved(), 0); + } + Ok(()) + } + #[derive(Debug)] struct PartitionedTestExec { cache: Arc, diff --git a/datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs b/datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs index fb46cedf96cf1..a19ad5fdbcc1a 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs @@ -70,12 +70,10 @@ impl PreparedHashJoinBuild { pub(super) fn probe_data(&self, probe_threads: usize) -> JoinLeftData { JoinLeftData { build: Arc::clone(&self.build), - null_aware_scope_map: None, - null_value_build_rows: None, visited_indices_bitmap: Mutex::new(BooleanBufferBuilder::new(0)), - null_indices_bitmap: Mutex::new(BooleanBufferBuilder::new(0)), - probe_completion: ProbeCompletion::new(probe_threads), - build_side_has_null: false, + probe_threads_counter: AtomicUsize::new(probe_threads), + probe_side_non_empty: AtomicBool::new(false), + probe_side_has_null: AtomicBool::new(false), } } @@ -195,7 +193,6 @@ impl HashJoinExec { !byte_keys, config, self.null_equality, - None, count, BuildMode::Prepared, ) From 105af0566254099bca9b6890f507f09d9b754db6 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Mon, 24 Aug 2026 17:22:40 +0000 Subject: [PATCH 08/19] fix: honor Parquet byte-array statistics ordering (#24525) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Why are the changes needed? ### Which issue does this PR close? Part of #10586. This fixes the unsigned string/binary byte-array case; it does not claim to resolve every Parquet ordering issue. It is also the correctness prerequisite for #24526. ### Rationale for this change A Parquet scan can skip a file, row group, or page when its statistics prove that no row can satisfy the filter. If those statistics use a different comparison order from the query, that proof is invalid: DataFusion can silently discard a row that should be returned. For example, suppose a Parquet row group contains `'aé'`, `'az'`, and `'b'`: ```sql SELECT s FROM t WHERE s = 'az'; -- Expected: one row containing 'az' ``` Parquet's deprecated byte-array `min`/`max` fields use signed comparison, whereas Arrow compares strings using unsigned UTF-8 bytes. The same values therefore have two different orders: | Comparison | Values in ascending order | Minimum / maximum | | --- | --- | --- | | Legacy signed-byte order | `aé < az < b` | `aé` / `b` | | Arrow's unsigned-byte order | `az < aé < b` | `az` / `b` | The first UTF-8 byte of `é` is `0xC3`: it sorts before `z`'s `0x7A` as a signed byte, but after it as an unsigned byte. If DataFusion interprets the legacy interval `['aé', 'b']` using Arrow's ordering, it sees `'az' < 'aé'` and can incorrectly skip the row group. Merely checking that `min <= max` does not help: the two reported endpoints are still in ascending order. The newer `min_value`/`max_value` fields and page-index bounds also need an ordering declaration that the reader understands. A missing or unknown `column_orders` entry is not enough to justify assuming Arrow's ordering. The [Parquet statistics definition](https://github.com/apache/parquet-format/blob/master/src/main/thrift/parquet.thrift) and [logical-type ordering rules](https://github.com/apache/parquet-format/blob/master/LogicalTypes.md) describe these distinctions. This bug is independently observable with an ordinary equality filter; it does not require a large `IN` list. ## What changes were proposed in this PR? ### What changes are included in this PR? The change makes a recognized comparison order a prerequisite for using unsigned byte-array bounds. At the point where Parquet metadata becomes Arrow statistics, DataFusion checks whether the column's physical/logical type and footer establish the expected unsigned order. It also rejects row-group bounds taken from the deprecated signed-order fields. When that evidence is missing, min/max is reported as unknown, so pruning keeps potentially matching data instead of guessing. The rule is applied at the granularity where the statistics are used. An unsafe row group prevents DataFusion from claiming a trustworthy bound for the whole file, but it does not make other row groups' valid bounds unusable. Page-index bounds are checked against the footer independently. The same safeguards cover static and runtime row-group pruning, as well as the inverse predicates used to decide whether every row already satisfies a filter. Null counts and unrelated columns' statistics remain available. The row-group pruning API gains a metadata-aware entry point so callers can supply the footer needed for this decision. The existing entry point remains source-compatible and behaves conservatively when that information is unavailable. Signed logical types such as decimals retain their existing behavior. ### Are there any user-facing changes? Queries no longer discard matching data because of these untrustworthy byte-array bounds. Older files, or files with an unrecognized ordering, may require more scanning. Modern files with trustworthy bounds retain min/max pruning. There is no file-format change or breaking public API change. ## How was this PR tested? ### Are these changes tested? The regression uses actual serialized Parquet files containing the example above, with modern, deprecated, missing-order, and unknown-order metadata. It checks that the matching `az` row survives file, row-group, runtime, and page pruning, while valid statistics can still eliminate unrelated data. On unchanged Apache `f1f0449a`, the adapted equality regression fails because the deprecated-order case loses `az`; the modern-statistics control passes. The dedicated Parquet-crate run passed 231 unit tests and four doctests. Its seven focused ordering tests also cover mixed safe/unsafe row groups, null counts, fixed-length binary and UUID, signed decimal, and logical types with undefined ordering. Formatting, all-targets/all-features Clippy with warnings denied, and `./dev/rust_lint.sh` passed. The extended workspace run passed 10,666 Rust tests, with eight ignored, and all 503 SQL-logic files. The existing metadata benchmark was run on Apache `f1f0449a` and this patch using the same valid modern-footer fixture. Across nine full-statistics cases there was no material regression; the largest case, with 256 columns and 128 row groups, measured 1.330 ms before and 1.332 ms after. Both runs used Rust 1.97.0, `release-nonlto`, 20 samples, and separate build directories on an Apple M5 Max.
Validation commands ```sh cargo test --locked --profile ci -p datafusion-datasource-parquet cargo bench --locked --profile release-nonlto -p datafusion-datasource-parquet \ --bench parquet_metadata_statistics -- \ metadata_full --sample-size 20 --warm-up-time 0.5 --measurement-time 1 --noplot RUST_BACKTRACE=1 cargo test --locked --profile ci \ --exclude datafusion-examples --exclude datafusion-benchmarks --exclude datafusion-cli \ --workspace --lib --tests --bins \ --features avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption ```
Backport provenance: - Upstream PR: https://github.com/apache/datafusion/pull/24525 - Upstream commit: 26b40dd781d2eeab1acc38370a7069a8b9ebe450 - Upstream self-tests and all additions/removals retained; only unavoidable fork patch context resolved. - Any fork-specific adaptation is kept in a separate following commit. Signed-off-by: discord9 --- .../benches/parquet_metadata_statistics.rs | 9 +- datafusion/datasource-parquet/src/metadata.rs | 74 +- datafusion/datasource-parquet/src/mod.rs | 2 + .../datasource-parquet/src/opener/mod.rs | 5 +- .../datasource-parquet/src/page_filter.rs | 15 + .../datasource-parquet/src/push_decoder.rs | 5 + .../src/row_group_filter.rs | 121 ++- .../src/statistics_order_tests.rs | 749 ++++++++++++++++++ 8 files changed, 956 insertions(+), 24 deletions(-) create mode 100644 datafusion/datasource-parquet/src/statistics_order_tests.rs diff --git a/datafusion/datasource-parquet/benches/parquet_metadata_statistics.rs b/datafusion/datasource-parquet/benches/parquet_metadata_statistics.rs index 46ebd100fde88..82147abfd9a01 100644 --- a/datafusion/datasource-parquet/benches/parquet_metadata_statistics.rs +++ b/datafusion/datasource-parquet/benches/parquet_metadata_statistics.rs @@ -28,6 +28,7 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use criterion::{BatchSize, BenchmarkId, Criterion, criterion_group, criterion_main}; use datafusion_datasource_parquet::metadata::DFParquetMetadata; use parquet::arrow::ArrowSchemaConverter; +use parquet::basic::ColumnOrder; use parquet::data_type::ByteArray; use parquet::file::metadata::{ ColumnChunkMetaData, FileMetaData, ParquetMetaData, RowGroupMetaData, @@ -169,13 +170,19 @@ fn make_synthetic_metadata( }) .collect::>(); + // Model a modern writer so byte-array bounds have a trustworthy order. + let column_orders = schema_descr + .columns() + .iter() + .map(|column| ColumnOrder::TYPE_DEFINED_ORDER(column.sort_order())) + .collect(); let file_metadata = FileMetaData::new( 1, (spec.row_groups * ROWS_PER_GROUP) as i64, Some("datafusion parquet metadata benchmark".to_string()), None, schema_descr, - None, + Some(column_orders), ); ParquetMetaData::new(file_metadata, row_groups) diff --git a/datafusion/datasource-parquet/src/metadata.rs b/datafusion/datasource-parquet/src/metadata.rs index 38db1e22ee6c0..41ceba128519d 100644 --- a/datafusion/datasource-parquet/src/metadata.rs +++ b/datafusion/datasource-parquet/src/metadata.rs @@ -43,12 +43,13 @@ use object_store::{ObjectMeta, ObjectStore}; use parquet::DecodeResult; use parquet::arrow::arrow_reader::statistics::StatisticsConverter; use parquet::arrow::{parquet_column, parquet_to_arrow_schema}; +use parquet::basic::{ColumnOrder, SortOrder, Type as PhysicalType}; use parquet::file::metadata::{ PageIndexPolicy, ParquetMetaData, ParquetMetaDataPushDecoder, ParquetMetaDataReader, RowGroupMetaData, SortingColumn, }; use parquet::file::statistics::Statistics as ParquetStatistics; -use parquet::schema::types::SchemaDescriptor; +use parquet::schema::types::{ColumnDescriptor, SchemaDescriptor}; use std::any::Any; use std::sync::Arc; @@ -57,6 +58,60 @@ use std::sync::Arc; /// would be too unreliable otherwise. const PARTIAL_NDV_THRESHOLD: f64 = 0.75; +fn requires_unsigned_byte_array_order(column: &ColumnDescriptor) -> bool { + matches!( + column.physical_type(), + PhysicalType::BYTE_ARRAY | PhysicalType::FIXED_LEN_BYTE_ARRAY + ) && column.sort_order() != SortOrder::SIGNED +} + +/// Whether a column's min/max bounds lack a comparison order matching Arrow's. +/// +/// The deprecated Parquet `min`/`max` fields use signed comparison, unlike +/// Arrow's string and binary comparisons. Even the modern bounds cannot be +/// interpreted without the corresponding footer `column_orders` entry. +/// Signed logical types, such as decimals, retain their existing behavior. +/// Columns with undefined sort orders, such as `INT96`, never have usable +/// min/max bounds regardless of their physical type. The `INT96` check is +/// defensive because parquet-rs does not currently expose those bounds. +pub(crate) fn has_untrusted_min_max_order( + parquet_schema: &SchemaDescriptor, + column_orders: Option<&[ColumnOrder]>, + parquet_column_index: usize, +) -> bool { + let column = parquet_schema.column(parquet_column_index); + if column.sort_order() == SortOrder::UNDEFINED { + return true; + } + requires_unsigned_byte_array_order(&column) + && (column.sort_order() != SortOrder::UNSIGNED + || column_orders + .and_then(|orders| orders.get(parquet_column_index)) + .copied() + != Some(ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNSIGNED))) +} + +/// Whether any row group provides byte-array bounds in the deprecated signed +/// order, or the column's logical order is undefined. +pub(crate) fn has_untrusted_byte_array_stats<'a>( + parquet_schema: &SchemaDescriptor, + parquet_column_index: Option, + row_groups: impl IntoIterator, +) -> bool { + parquet_column_index.is_some_and(|index| { + let column = parquet_schema.column(index); + requires_unsigned_byte_array_order(&column) + && (column.sort_order() != SortOrder::UNSIGNED + || row_groups.into_iter().any(|group| { + group.column(index).statistics().is_some_and(|stats| { + stats.is_min_max_deprecated() + && (stats.min_bytes_opt().is_some() + || stats.max_bytes_opt().is_some()) + }) + })) + }) +} + /// Handles fetching Parquet file schema, metadata and statistics /// from object store. /// @@ -496,6 +551,23 @@ impl<'a> DFParquetMetadata<'a> { file_metadata.schema_descr(), ) { Ok(stats_converter) => { + let parquet_index = stats_converter.parquet_column_index(); + if parquet_index.is_some_and(|index| { + has_untrusted_min_max_order( + file_metadata.schema_descr(), + file_metadata.column_orders().map(Vec::as_slice), + index, + ) + }) || has_untrusted_byte_array_stats( + file_metadata.schema_descr(), + parquet_index, + row_groups_metadata, + ) { + // The remaining row groups cannot establish bounds + // for the whole file. Keep unrelated statistics. + min_accs[idx] = None; + max_accs[idx] = None; + } let mut accumulators = StatisticsAccumulators { min_accs: &mut min_accs, max_accs: &mut max_accs, diff --git a/datafusion/datasource-parquet/src/mod.rs b/datafusion/datasource-parquet/src/mod.rs index 35f831230b305..23348f94aae8f 100644 --- a/datafusion/datasource-parquet/src/mod.rs +++ b/datafusion/datasource-parquet/src/mod.rs @@ -42,6 +42,8 @@ mod schema_coercion; mod sink; mod sort; pub mod source; +#[cfg(test)] +mod statistics_order_tests; mod supported_predicates; #[cfg(test)] mod test_util; diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index 693e9bd2cbf31..1c59f345e30bf 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -1105,10 +1105,9 @@ impl FiltersPreparedParquetOpen { // If there is a predicate that can be evaluated against the metadata if let Some(predicate) = self.pruning_predicate.as_ref().map(|p| p.as_ref()) { if prepared.enable_row_group_stats_pruning { - row_groups.prune_by_statistics( + row_groups.prune_by_statistics_with_metadata( &prepared.physical_file_schema, - loaded.reader_metadata.parquet_schema(), - rg_metadata, + &file_metadata, predicate, &prepared.file_metrics, ); diff --git a/datafusion/datasource-parquet/src/page_filter.rs b/datafusion/datasource-parquet/src/page_filter.rs index 6bc1aca667981..5372ae0a0007d 100644 --- a/datafusion/datasource-parquet/src/page_filter.rs +++ b/datafusion/datasource-parquet/src/page_filter.rs @@ -22,6 +22,7 @@ use std::sync::Arc; use super::metrics::ParquetFileMetrics; use crate::ParquetAccessPlan; +use crate::metadata::has_untrusted_min_max_order; use arrow::array::BooleanArray; use arrow::{ @@ -490,6 +491,7 @@ struct PagesPruningStatistics<'a> { column_index: &'a ParquetColumnIndex, offset_index: &'a ParquetOffsetIndex, page_offsets: &'a Vec, + trusted_min_max: bool, } impl<'a> PagesPruningStatistics<'a> { @@ -529,6 +531,12 @@ impl<'a> PagesPruningStatistics<'a> { return None; }; let page_offsets = offset_index_metadata.page_locations(); + let file_metadata = parquet_metadata.file_metadata(); + let trusted_min_max = !has_untrusted_min_max_order( + file_metadata.schema_descr(), + file_metadata.column_orders().map(Vec::as_slice), + parquet_column_index, + ); Some(Self { row_group_index, @@ -537,6 +545,7 @@ impl<'a> PagesPruningStatistics<'a> { column_index, offset_index, page_offsets, + trusted_min_max, }) } @@ -563,6 +572,9 @@ impl<'a> PagesPruningStatistics<'a> { } impl PruningStatistics for PagesPruningStatistics<'_> { fn min_values(&self, _column: &datafusion_common::Column) -> Option { + if !self.trusted_min_max { + return None; + } match self.converter.data_page_mins( self.column_index, self.offset_index, @@ -577,6 +589,9 @@ impl PruningStatistics for PagesPruningStatistics<'_> { } fn max_values(&self, _column: &datafusion_common::Column) -> Option { + if !self.trusted_min_max { + return None; + } match self.converter.data_page_maxes( self.column_index, self.offset_index, diff --git a/datafusion/datasource-parquet/src/push_decoder.rs b/datafusion/datasource-parquet/src/push_decoder.rs index 74d8997198872..d91d86a657902 100644 --- a/datafusion/datasource-parquet/src/push_decoder.rs +++ b/datafusion/datasource-parquet/src/push_decoder.rs @@ -211,6 +211,11 @@ impl RowGroupPruner { .collect::>(); let stats = RowGroupPruningStatistics { parquet_schema: self.parquet_metadata.file_metadata().schema_descr(), + column_orders: self + .parquet_metadata + .file_metadata() + .column_orders() + .map(Vec::as_slice), row_group_metadatas, arrow_schema: self.arrow_schema.as_ref(), // Match the existing static row-group pruning behavior: when a diff --git a/datafusion/datasource-parquet/src/row_group_filter.rs b/datafusion/datasource-parquet/src/row_group_filter.rs index ddf71bb7e6d95..80db13edcda81 100644 --- a/datafusion/datasource-parquet/src/row_group_filter.rs +++ b/datafusion/datasource-parquet/src/row_group_filter.rs @@ -20,7 +20,9 @@ use std::sync::Arc; use super::{ParquetAccessPlan, ParquetFileMetrics, RowGroupAccess}; use crate::bloom_filter::BloomFilterStatistics; +use crate::metadata::{has_untrusted_byte_array_stats, has_untrusted_min_max_order}; use arrow::array::{ArrayRef, BooleanArray, UInt64Array}; +use arrow::compute::nullif; use arrow::datatypes::Schema; use datafusion_common::pruning::PruningStatistics; use datafusion_common::{Column, Result, ScalarValue}; @@ -31,7 +33,8 @@ use datafusion_physical_expr::utils::collect_columns; use datafusion_physical_expr::{PhysicalExpr, PhysicalExprSimplifier}; use datafusion_pruning::{PruningPredicate, PruningPredicateBuilder}; use parquet::arrow::arrow_reader::statistics::StatisticsConverter; -use parquet::file::metadata::RowGroupMetaData; +use parquet::basic::ColumnOrder; +use parquet::file::metadata::{ParquetMetaData, RowGroupMetaData}; use parquet::schema::types::SchemaDescriptor; /// Reduces the [`ParquetAccessPlan`] based on row group level metadata. @@ -252,8 +255,10 @@ impl RowGroupAccessPlanFilter { /// /// Updates this set to mark row groups that should not be scanned /// - /// Note: This method currently ignores ColumnOrder - /// + /// This method has no file footer, so it cannot establish the sort order + /// of string or binary min/max statistics. Such bounds are ignored. Use + /// [`Self::prune_by_statistics_with_metadata`] when the full metadata is + /// available. Null counts and other columns' statistics remain usable. /// /// # Panics /// if `groups.len() != self.len()` @@ -264,6 +269,51 @@ impl RowGroupAccessPlanFilter { groups: &[RowGroupMetaData], predicate: &PruningPredicate, metrics: &ParquetFileMetrics, + ) { + self.prune_by_statistics_inner( + arrow_schema, + parquet_schema, + groups, + None, + predicate, + metrics, + ); + } + + /// Prune row groups using statistics and the comparison orders recorded + /// in the Parquet file footer. + /// + /// Unlike [`Self::prune_by_statistics`], this method can use string and + /// binary bounds when their ordering is known to match Arrow's ordering. + /// + /// # Panics + /// if `metadata.num_row_groups() != self.len()` + pub fn prune_by_statistics_with_metadata( + &mut self, + arrow_schema: &Schema, + metadata: &ParquetMetaData, + predicate: &PruningPredicate, + metrics: &ParquetFileMetrics, + ) { + let file_metadata = metadata.file_metadata(); + self.prune_by_statistics_inner( + arrow_schema, + file_metadata.schema_descr(), + metadata.row_groups(), + file_metadata.column_orders().map(Vec::as_slice), + predicate, + metrics, + ); + } + + fn prune_by_statistics_inner( + &mut self, + arrow_schema: &Schema, + parquet_schema: &SchemaDescriptor, + groups: &[RowGroupMetaData], + column_orders: Option<&[ColumnOrder]>, + predicate: &PruningPredicate, + metrics: &ParquetFileMetrics, ) { // scoped timer updates on drop let _timer_guard = metrics.statistics_eval_time.timer(); @@ -278,6 +328,7 @@ impl RowGroupAccessPlanFilter { let pruning_stats = RowGroupPruningStatistics { parquet_schema, + column_orders, row_group_metadatas, arrow_schema, // Preserve the existing row-group pruning behavior. This path only @@ -304,8 +355,7 @@ impl RowGroupAccessPlanFilter { // Check if any of the matched row groups are fully contained by the predicate self.identify_fully_matched_row_groups( &fully_contained_candidates_original_idx, - arrow_schema, - parquet_schema, + &pruning_stats, groups, predicate, metrics, @@ -330,8 +380,7 @@ impl RowGroupAccessPlanFilter { fn identify_fully_matched_row_groups( &mut self, candidate_row_group_indices: &[usize], - arrow_schema: &Schema, - parquet_schema: &SchemaDescriptor, + pruning_stats: &RowGroupPruningStatistics<'_>, groups: &[RowGroupMetaData], predicate: &PruningPredicate, metrics: &ParquetFileMetrics, @@ -339,6 +388,7 @@ impl RowGroupAccessPlanFilter { if candidate_row_group_indices.is_empty() { return; } + let arrow_schema = pruning_stats.arrow_schema; let mut inverted_expr: Arc = Arc::new(NotExpr::new(Arc::clone(predicate.orig_expr()))); @@ -381,7 +431,8 @@ impl RowGroupAccessPlanFilter { }; let inverted_pruning_stats = RowGroupPruningStatistics { - parquet_schema, + parquet_schema: pruning_stats.parquet_schema, + column_orders: pruning_stats.column_orders, row_group_metadatas: candidate_row_group_indices .iter() .map(|&i| &groups[i]) @@ -463,6 +514,7 @@ impl RowGroupAccessPlanFilter { /// duplicating the statistics-to-`PruningStatistics` plumbing. pub(crate) struct RowGroupPruningStatistics<'a> { pub(crate) parquet_schema: &'a SchemaDescriptor, + pub(crate) column_orders: Option<&'a [ColumnOrder]>, pub(crate) row_group_metadatas: Vec<&'a RowGroupMetaData>, pub(crate) arrow_schema: &'a Schema, pub(crate) missing_null_counts_as_zero: bool, @@ -470,14 +522,11 @@ pub(crate) struct RowGroupPruningStatistics<'a> { impl<'a> RowGroupPruningStatistics<'a> { /// Return an iterator over the row group metadata - fn metadata_iter(&'a self) -> impl Iterator + 'a { + fn metadata_iter(&self) -> impl Iterator + '_ { self.row_group_metadatas.iter().copied() } - fn statistics_converter<'b>( - &'a self, - column: &'b Column, - ) -> Result> { + fn statistics_converter(&self, column: &Column) -> Result> { Ok(StatisticsConverter::try_new( &column.name, self.arrow_schema, @@ -485,19 +534,53 @@ impl<'a> RowGroupPruningStatistics<'a> { )? .with_missing_null_counts_as_zero(self.missing_null_counts_as_zero)) } + + fn min_max_statistics_converter( + &self, + column: &Column, + ) -> Option> { + let converter = self.statistics_converter(column).ok()?; + let parquet_index = converter.parquet_column_index(); + if parquet_index.is_some_and(|index| { + has_untrusted_min_max_order(self.parquet_schema, self.column_orders, index) + }) { + return None; + } + Some(converter) + } + + fn mask_untrusted_byte_array_stats( + &self, + parquet_index: Option, + values: ArrayRef, + ) -> Option { + if !has_untrusted_byte_array_stats( + self.parquet_schema, + parquet_index, + self.metadata_iter(), + ) { + return Some(values); + } + // A file may mix legacy and modern row-group statistics. Keep the + // modern bounds usable rather than discarding this entire column. + let mask = BooleanArray::from_iter(self.metadata_iter().map(|group| { + has_untrusted_byte_array_stats(self.parquet_schema, parquet_index, [group]) + })); + nullif(values.as_ref(), &mask).ok() + } } impl PruningStatistics for RowGroupPruningStatistics<'_> { fn min_values(&self, column: &Column) -> Option { - self.statistics_converter(column) - .and_then(|c| Ok(c.row_group_mins(self.metadata_iter())?)) - .ok() + let converter = self.min_max_statistics_converter(column)?; + let values = converter.row_group_mins(self.metadata_iter()).ok()?; + self.mask_untrusted_byte_array_stats(converter.parquet_column_index(), values) } fn max_values(&self, column: &Column) -> Option { - self.statistics_converter(column) - .and_then(|c| Ok(c.row_group_maxes(self.metadata_iter())?)) - .ok() + let converter = self.min_max_statistics_converter(column)?; + let values = converter.row_group_maxes(self.metadata_iter()).ok()?; + self.mask_untrusted_byte_array_stats(converter.parquet_column_index(), values) } fn num_containers(&self) -> usize { diff --git a/datafusion/datasource-parquet/src/statistics_order_tests.rs b/datafusion/datasource-parquet/src/statistics_order_tests.rs new file mode 100644 index 0000000000000..bf9df532f1b58 --- /dev/null +++ b/datafusion/datasource-parquet/src/statistics_order_tests.rs @@ -0,0 +1,749 @@ +// 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. + +//! Regression tests for interpreting Parquet byte-array statistics orders. + +use std::io::Write; +use std::sync::Arc; + +use arrow::array::{BooleanArray, record_batch}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use bytes::Bytes; +use datafusion_common::pruning::{PrunableStatistics, PruningStatistics}; +use datafusion_common::stats::Precision; +use datafusion_common::{Column, ScalarValue, Statistics}; +use datafusion_expr::{Expr, col, lit}; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::planner::logical2physical; +use datafusion_physical_plan::metrics::{Count, ExecutionPlanMetricsSet}; +use datafusion_pruning::{MAX_IN_LIST_SIZE, PruningPredicate, PruningPredicateBuilder}; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use parquet::basic::{ColumnOrder, LogicalType, SortOrder, Type as PhysicalType}; +use parquet::data_type::{ByteArray, FixedLenByteArray}; +use parquet::file::metadata::{ + ColumnChunkMetaData, ColumnIndexBuilder, FileMetaData, OffsetIndexBuilder, + PageIndexPolicy, ParquetMetaData, ParquetMetaDataReader, ParquetMetaDataWriter, + RowGroupMetaData, +}; +use parquet::file::properties::{EnabledStatistics, WriterProperties}; +use parquet::file::statistics::Statistics as ParquetStatistics; +use parquet::file::writer::TrackedWrite; +use parquet::schema::types::{SchemaDescriptor, Type as ParquetType}; + +use crate::RowGroupAccessPlanFilter; +use crate::metadata::{DFParquetMetadata, has_untrusted_min_max_order}; +use crate::push_decoder::RowGroupPruner; +use crate::row_group_filter::RowGroupPruningStatistics; +use crate::{PagePruningAccessPlanFilter, ParquetAccessPlan, ParquetFileMetrics}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum StatisticsOrder { + Modern, + Deprecated, + Missing, + Unknown, +} + +struct TestFile { + bytes: Bytes, + schema: SchemaRef, + metadata: Arc, +} + +impl TestFile { + fn new(order: StatisticsOrder) -> Self { + let batch = record_batch!( + ( + "s", + Utf8, + vec![ + Some("aé"), + Some("az"), + Some("b"), + None, + None, + None, + Some("d"), + Some("e"), + Some("f"), + ] + ), + ("n", Int32, [1, 2, 3, 10, 11, 12, 20, 21, 22]) + ) + .unwrap(); + let schema = batch.schema(); + let properties = WriterProperties::builder() + .set_max_row_group_row_count(Some(3)) + .set_data_page_row_count_limit(3) + .set_write_batch_size(3) + .set_dictionary_enabled(false) + .set_statistics_enabled(EnabledStatistics::Page) + .build(); + let mut original = Vec::new(); + let mut writer = + ArrowWriter::try_new(&mut original, Arc::clone(&schema), Some(properties)) + .unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let original = Bytes::from(original); + let metadata = read_metadata(&original); + assert_eq!(metadata.num_row_groups(), 3); + + if order == StatisticsOrder::Modern { + return Self { + bytes: original, + schema, + metadata: Arc::new(metadata), + }; + } + + // Signed-byte comparison gives ["aé", "b"] for the first row + // group's ["aé", "az", "b"]. The endpoints are not inverted in + // unsigned order, but the interval wrongly excludes "az". + let mut row_groups = metadata.row_groups().to_vec(); + let mut columns = row_groups[0].columns().to_vec(); + columns[0] = columns[0] + .clone() + .into_builder() + .set_statistics(ParquetStatistics::byte_array( + Some(ByteArray::from("aé")), + Some(ByteArray::from("b")), + None, + Some(0), + order == StatisticsOrder::Deprecated, + )) + .build() + .unwrap(); + row_groups[0] = row_groups[0] + .clone() + .into_builder() + .set_column_metadata(columns) + .build() + .unwrap(); + + let mut column_index = metadata.column_index().unwrap().clone(); + if matches!(order, StatisticsOrder::Missing | StatisticsOrder::Unknown) { + let mut index = ColumnIndexBuilder::new(PhysicalType::BYTE_ARRAY); + index.append(false, "aé".as_bytes().to_vec(), b"b".to_vec(), 0); + column_index[0][0] = index.build().unwrap(); + } + let metadata = metadata + .into_builder() + .set_row_groups(row_groups) + .set_column_index(Some(column_index)) + .build(); + + // Keep the real data pages, and serialize the replacement statistics + // and page indexes at their actual file offsets. + let mut bytes = Vec::new(); + let mut tracked = TrackedWrite::new(&mut bytes); + tracked + .write_all(&original[..footer_start(&original)]) + .unwrap(); + ParquetMetaDataWriter::new_with_tracked(tracked, &metadata) + .finish() + .unwrap(); + + // parquet-rs always writes TYPEORDER and cannot write an unknown + // ColumnOrder. The final Thrift field is column_orders (field 7): a + // two-element list of unions, followed by the FileMetaData STOP. + // Alter just that field to model old and future writers, then verify + // the decoded footer below. No data or page-index offsets change. + if matches!(order, StatisticsOrder::Missing | StatisticsOrder::Unknown) { + let end = bytes.len() - 8; + let encoded_orders = [0x19, 0x2c, 0x1c, 0, 0, 0x1c, 0, 0, 0]; + let start = end - encoded_orders.len(); + assert_eq!(&bytes[start..end], &encoded_orders); + let metadata_start = footer_start(&bytes); + if order == StatisticsOrder::Missing { + bytes.drain(start..end - 1); + let new_end = bytes.len() - 8; + let metadata_len = (new_end - metadata_start) as u32; + bytes[new_end..new_end + 4].copy_from_slice(&metadata_len.to_le_bytes()); + } else { + // Change the first union member from field 1 (TYPEORDER) to + // an unrecognized field 2. The numeric column stays known. + bytes[start + 2] = 0x2c; + } + } + + let bytes = Bytes::from(bytes); + let metadata = read_metadata(&bytes); + let expected_order = match order { + StatisticsOrder::Missing => ColumnOrder::UNDEFINED, + StatisticsOrder::Unknown => ColumnOrder::UNKNOWN, + _ => ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNSIGNED), + }; + assert_eq!(metadata.file_metadata().column_order(0), expected_order); + assert_eq!( + metadata + .row_group(0) + .column(0) + .statistics() + .unwrap() + .is_min_max_deprecated(), + order == StatisticsOrder::Deprecated, + ); + Self { + bytes, + schema, + metadata: Arc::new(metadata), + } + } + + fn predicate(&self, expr: &Expr) -> (Arc, PruningPredicate) { + let physical = logical2physical(expr, &self.schema); + let pruning = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&self.schema)) + .try_build(Arc::clone(&physical)) + .unwrap(); + (physical, pruning) + } + + fn statistics(&self) -> Statistics { + DFParquetMetadata::statistics_from_parquet_metadata(&self.metadata, &self.schema) + .unwrap() + } + + fn file_matches(&self, predicate: &PruningPredicate) -> bool { + let stats = PrunableStatistics::new( + vec![Arc::new(self.statistics())], + Arc::clone(&self.schema), + ); + predicate.prune(&stats).unwrap()[0] + } + + fn row_group_plan(&self, predicate: &PruningPredicate) -> ParquetAccessPlan { + let mut filter = RowGroupAccessPlanFilter::new(ParquetAccessPlan::new_all( + self.metadata.num_row_groups(), + )); + filter.prune_by_statistics_with_metadata( + &self.schema, + &self.metadata, + predicate, + &metrics(), + ); + filter.build() + } + + fn page_plan( + &self, + physical: &Arc, + plan: ParquetAccessPlan, + ) -> ParquetAccessPlan { + PagePruningAccessPlanFilter::new(physical, Arc::clone(&self.schema)) + .prune_plan_with_page_index( + plan, + &self.schema, + self.metadata.file_metadata().schema_descr(), + &self.metadata, + &metrics(), + ) + } + + fn matching_rows( + &self, + physical: &Arc, + plan: ParquetAccessPlan, + ) -> usize { + let mut builder = ParquetRecordBatchReaderBuilder::try_new(self.bytes.clone()) + .unwrap() + .with_row_groups(plan.row_group_indexes()); + if let Some(selection) = plan + .into_overall_row_selection(self.metadata.row_groups()) + .unwrap() + { + builder = builder.with_row_selection(selection); + } + builder + .build() + .unwrap() + .map(|batch| { + let batch = batch.unwrap(); + let matches = physical + .evaluate(&batch) + .unwrap() + .into_array(batch.num_rows()) + .unwrap(); + matches + .as_any() + .downcast_ref::() + .unwrap() + .true_count() + }) + .sum() + } +} + +fn footer_start(bytes: &[u8]) -> usize { + let end = bytes.len() - 8; + let metadata_len = u32::from_le_bytes(bytes[end..end + 4].try_into().unwrap()); + end - metadata_len as usize +} + +fn read_metadata(bytes: &Bytes) -> ParquetMetaData { + ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Required) + .parse_and_finish(bytes) + .unwrap() +} + +fn metrics() -> ParquetFileMetrics { + ParquetFileMetrics::new( + 0, + "statistics-order.parquet", + &ExecutionPlanMetricsSet::new(), + ) +} + +#[test] +fn byte_array_order_preserves_matching_rows_at_every_pruning_level() { + for order in [ + StatisticsOrder::Modern, + StatisticsOrder::Deprecated, + StatisticsOrder::Missing, + StatisticsOrder::Unknown, + ] { + let file = TestFile::new(order); + let (physical, predicate) = file.predicate(&col("s").eq(lit("az"))); + let all = ParquetAccessPlan::new_all(file.metadata.num_row_groups()); + assert_eq!(file.matching_rows(&physical, all.clone()), 1); + assert!(file.file_matches(&predicate), "order={order:?}"); + let row_groups = file.row_group_plan(&predicate); + assert!(row_groups.should_scan(0), "order={order:?}"); + assert!(!row_groups.is_fully_matched(0), "order={order:?}"); + assert_eq!( + row_groups.row_group_indexes(), + if matches!(order, StatisticsOrder::Modern | StatisticsOrder::Deprecated) { + vec![0] + } else { + vec![0, 2] + }, + ); + assert_eq!(file.matching_rows(&physical, row_groups.clone()), 1); + assert_eq!( + file.matching_rows(&physical, file.page_plan(&physical, all)), + 1 + ); + assert_eq!( + file.matching_rows(&physical, file.page_plan(&physical, row_groups)), + 1, + "order={order:?}", + ); + + let mut runtime_pruner = RowGroupPruner::new( + physical, + Arc::clone(&file.schema), + Arc::clone(&file.metadata), + Count::new(), + Count::new(), + MAX_IN_LIST_SIZE, + ); + assert!(!runtime_pruner.should_prune(&[0]), "order={order:?}"); + assert!(runtime_pruner.should_prune(&[1]), "all-null row group"); + } +} + +#[test] +fn byte_array_order_keeps_null_counts_and_unrelated_column_bounds() { + for order in [ + StatisticsOrder::Deprecated, + StatisticsOrder::Missing, + StatisticsOrder::Unknown, + ] { + let file = TestFile::new(order); + let statistics = file.statistics(); + let string = &statistics.column_statistics[0]; + assert_eq!(string.min_value, Precision::Absent, "order={order:?}"); + assert_eq!(string.max_value, Precision::Absent, "order={order:?}"); + assert_eq!(string.null_count, Precision::Exact(3)); + assert_eq!( + statistics.column_statistics[1].min_value, + Precision::Exact(ScalarValue::Int32(Some(1))), + ); + assert_eq!( + statistics.column_statistics[1].max_value, + Precision::Exact(ScalarValue::Int32(Some(22))), + ); + + let (physical, predicate) = file.predicate(&col("s").is_null()); + assert_eq!(file.row_group_plan(&predicate).row_group_indexes(), vec![1]); + let page_plan = file.page_plan(&physical, ParquetAccessPlan::new_all(3)); + assert_eq!(page_plan.row_group_indexes(), vec![1]); + assert_eq!(file.matching_rows(&physical, page_plan), 3); + + let expr = col("s").eq(lit("az")).and(col("n").eq(lit(99))); + let (physical, predicate) = file.predicate(&expr); + assert!(!file.file_matches(&predicate)); + assert!( + file.row_group_plan(&predicate) + .row_group_indexes() + .is_empty() + ); + assert!( + file.page_plan(&physical, ParquetAccessPlan::new_all(3)) + .row_group_indexes() + .is_empty(), + ); + } +} + +#[test] +fn byte_array_order_modern_bounds_and_null_only_groups_remain_usable() { + let file = TestFile::new(StatisticsOrder::Modern); + let statistics = file.statistics(); + assert_eq!( + statistics.column_statistics[0].min_value, + Precision::Exact(ScalarValue::Utf8(Some("az".to_owned()))), + ); + let (physical, predicate) = file.predicate(&col("s").eq(lit("az"))); + assert_eq!(file.row_group_plan(&predicate).row_group_indexes(), vec![0]); + assert_eq!( + file.page_plan(&physical, ParquetAccessPlan::new_all(3)) + .row_group_indexes(), + vec![0], + ); + + // The compatibility API has no footer. It still uses null counts, but + // cannot assume that even modern-looking byte-array bounds are unsigned. + let mut filter = RowGroupAccessPlanFilter::new(ParquetAccessPlan::new_all(3)); + filter.prune_by_statistics( + &file.schema, + file.metadata.file_metadata().schema_descr(), + file.metadata.row_groups(), + &predicate, + &metrics(), + ); + assert_eq!(filter.build().row_group_indexes(), vec![0, 2]); +} + +#[test] +fn byte_array_order_guard_follows_parquet_type_not_arrow_representation() { + let file = TestFile::new(StatisticsOrder::Deprecated); + let metadata = file.metadata.file_metadata(); + let column = Column::from_name("s"); + for data_type in [ + DataType::Utf8, + DataType::LargeUtf8, + DataType::Utf8View, + DataType::Binary, + DataType::LargeBinary, + DataType::BinaryView, + DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)), + ] { + let schema = Schema::new(vec![Field::new("s", data_type, true)]); + let stats = RowGroupPruningStatistics { + parquet_schema: metadata.schema_descr(), + column_orders: metadata.column_orders().map(Vec::as_slice), + row_group_metadatas: file.metadata.row_groups().iter().collect(), + arrow_schema: &schema, + missing_null_counts_as_zero: true, + }; + for values in [stats.min_values(&column), stats.max_values(&column)] { + let values = values.unwrap(); + assert!(values.is_null(0)); + assert!(values.is_null(1)); + assert!(!values.is_null(2)); + } + assert!(stats.null_counts(&column).is_some()); + } +} + +fn single_column_metadata( + parquet_type: ParquetType, + statistics: ParquetStatistics, + order: Option, +) -> ParquetMetaData { + let physical_type = parquet_type.get_physical_type(); + let schema = Arc::new(SchemaDescriptor::new(Arc::new( + ParquetType::group_type_builder("schema") + .with_fields(vec![Arc::new(parquet_type)]) + .build() + .unwrap(), + ))); + let mut column_index = ColumnIndexBuilder::new(physical_type); + column_index.append( + false, + statistics.min_bytes_opt().unwrap().to_vec(), + statistics.max_bytes_opt().unwrap().to_vec(), + 0, + ); + let mut offset_index = OffsetIndexBuilder::new(); + offset_index.append_row_count(3); + offset_index.append_offset_and_size(0, 1); + let column = ColumnChunkMetaData::builder(schema.column(0)) + .set_num_values(3) + .set_statistics(statistics) + .build() + .unwrap(); + let group = RowGroupMetaData::builder(Arc::clone(&schema)) + .set_num_rows(3) + .set_column_metadata(vec![column]) + .build() + .unwrap(); + ParquetMetaData::new( + FileMetaData::new(1, 3, None, None, schema, order.map(|order| vec![order])), + vec![group], + ) + .into_builder() + .set_column_index(Some(vec![vec![column_index.build().unwrap()]])) + .set_offset_index(Some(vec![vec![offset_index.build()]])) + .build() +} + +#[test] +fn fixed_byte_array_and_uuid_orders_guard_bounds_but_not_null_counts() { + let schema = Arc::new(Schema::new(vec![Field::new( + "s", + DataType::FixedSizeBinary(16), + true, + )])); + let padded = |prefix: &[u8]| { + let mut value = vec![0; 16]; + value[..prefix.len()].copy_from_slice(prefix); + value + }; + let min = padded("aé".as_bytes()); + let max = padded(b"b"); + let value = ScalarValue::FixedSizeBinary(16, Some(padded(b"az"))); + let physical = logical2physical(&col("s").eq(lit(value)), &schema); + let predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(Arc::clone(&physical)) + .unwrap(); + + for logical_type in [None, Some(LogicalType::Uuid)] { + for deprecated in [false, true] { + for order in [ + None, + Some(ColumnOrder::UNKNOWN), + Some(ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::SIGNED)), + Some(ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNSIGNED)), + ] { + let parquet_type = ParquetType::primitive_type_builder( + "s", + PhysicalType::FIXED_LEN_BYTE_ARRAY, + ) + .with_length(16) + .with_logical_type(logical_type.clone()) + .build() + .unwrap(); + let statistics = ParquetStatistics::fixed_len_byte_array( + Some(FixedLenByteArray::from(min.clone())), + Some(FixedLenByteArray::from(max.clone())), + None, + Some(0), + deprecated, + ); + let metadata = single_column_metadata(parquet_type, statistics, order); + let trusted = + order == Some(ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNSIGNED)); + let statistics = DFParquetMetadata::statistics_from_parquet_metadata( + &metadata, &schema, + ) + .unwrap(); + assert_eq!( + statistics.column_statistics[0].min_value, + if trusted && !deprecated { + Precision::Exact(ScalarValue::FixedSizeBinary( + 16, + Some(min.clone()), + )) + } else { + Precision::Absent + }, + ); + assert_eq!( + statistics.column_statistics[0].null_count, + Precision::Exact(0), + ); + + let mut row_groups = + RowGroupAccessPlanFilter::new(ParquetAccessPlan::new_all(1)); + row_groups.prune_by_statistics_with_metadata( + &schema, + &metadata, + &predicate, + &metrics(), + ); + assert_eq!(row_groups.build().should_scan(0), !trusted || deprecated); + let pages = + PagePruningAccessPlanFilter::new(&physical, Arc::clone(&schema)) + .prune_plan_with_page_index( + ParquetAccessPlan::new_all(1), + &schema, + metadata.file_metadata().schema_descr(), + &metadata, + &metrics(), + ); + // Modern page indexes are independent of legacy row-group + // bounds, but still need a recognized footer order. + assert_eq!(pages.should_scan(0), !trusted); + } + } + } +} + +#[test] +fn signed_decimal_byte_array_statistics_remain_usable() { + let schema = Arc::new(Schema::new(vec![Field::new( + "s", + DataType::Decimal128(10, 0), + true, + )])); + for physical_type in [PhysicalType::BYTE_ARRAY, PhysicalType::FIXED_LEN_BYTE_ARRAY] { + let parquet_type = ParquetType::primitive_type_builder("s", physical_type) + .with_length(16) + .with_logical_type(Some(LogicalType::decimal(0, 10))) + .with_precision(10) + .with_scale(0) + .build() + .unwrap(); + let min = (-2_i128).to_be_bytes().to_vec(); + let max = 3_i128.to_be_bytes().to_vec(); + let statistics = match physical_type { + PhysicalType::BYTE_ARRAY => ParquetStatistics::byte_array( + Some(ByteArray::from(min)), + Some(ByteArray::from(max)), + None, + Some(0), + true, + ), + _ => ParquetStatistics::fixed_len_byte_array( + Some(FixedLenByteArray::from(min)), + Some(FixedLenByteArray::from(max)), + None, + Some(0), + true, + ), + }; + let metadata = single_column_metadata(parquet_type, statistics, None); + let statistics = + DFParquetMetadata::statistics_from_parquet_metadata(&metadata, &schema) + .unwrap(); + assert_eq!( + statistics.column_statistics[0].min_value, + Precision::Exact(ScalarValue::Decimal128(Some(-2), 10, 0)), + ); + let physical = logical2physical( + &col("s").eq(lit(ScalarValue::Decimal128(Some(4), 10, 0))), + &schema, + ); + let predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(physical) + .unwrap(); + let mut row_groups = RowGroupAccessPlanFilter::new(ParquetAccessPlan::new_all(1)); + row_groups.prune_by_statistics_with_metadata( + &schema, + &metadata, + &predicate, + &metrics(), + ); + assert!(!row_groups.build().should_scan(0)); + } +} + +#[test] +fn undefined_int96_order_is_never_trusted() { + let parquet_type = ParquetType::primitive_type_builder("s", PhysicalType::INT96) + .build() + .unwrap(); + let schema = SchemaDescriptor::new(Arc::new( + ParquetType::group_type_builder("schema") + .with_fields(vec![Arc::new(parquet_type)]) + .build() + .unwrap(), + )); + assert_eq!(schema.column(0).sort_order(), SortOrder::UNDEFINED); + + for order in [ + None, + Some(ColumnOrder::UNDEFINED), + Some(ColumnOrder::UNKNOWN), + Some(ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNDEFINED)), + Some(ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::SIGNED)), + Some(ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNSIGNED)), + ] { + assert!( + has_untrusted_min_max_order( + &schema, + order.as_ref().map(std::slice::from_ref), + 0 + ), + "order={order:?}", + ); + } +} + +#[test] +fn undefined_logical_byte_array_order_is_not_a_bound() { + let parquet_type = ParquetType::primitive_type_builder("s", PhysicalType::BYTE_ARRAY) + .with_logical_type(Some(LogicalType::_Unknown { field_id: 100 })) + .build() + .unwrap(); + let statistics = ParquetStatistics::byte_array( + Some(ByteArray::from("aé")), + Some(ByteArray::from("b")), + None, + Some(0), + false, + ); + let metadata = single_column_metadata( + parquet_type, + statistics, + Some(ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNDEFINED)), + ); + let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, true)])); + let file_statistics = + DFParquetMetadata::statistics_from_parquet_metadata(&metadata, &schema).unwrap(); + assert_eq!( + file_statistics.column_statistics[0].min_value, + Precision::Absent + ); + assert_eq!( + file_statistics.column_statistics[0].max_value, + Precision::Absent + ); + assert_eq!( + file_statistics.column_statistics[0].null_count, + Precision::Exact(0) + ); + let physical = logical2physical(&col("s").eq(lit("az")), &schema); + let predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(Arc::clone(&physical)) + .unwrap(); + let mut row_groups = RowGroupAccessPlanFilter::new(ParquetAccessPlan::new_all(1)); + row_groups.prune_by_statistics_with_metadata( + &schema, + &metadata, + &predicate, + &metrics(), + ); + assert!(row_groups.build().should_scan(0)); + let pages = PagePruningAccessPlanFilter::new(&physical, Arc::clone(&schema)) + .prune_plan_with_page_index( + ParquetAccessPlan::new_all(1), + &schema, + metadata.file_metadata().schema_descr(), + &metadata, + &metrics(), + ); + assert!(pages.should_scan(0)); +} From 5b464f88f01fc1cb0e63e13c3edc615890850421 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Thu, 27 Aug 2026 01:34:40 +0000 Subject: [PATCH 09/19] perf: use compact pruning for large string IN lists (#24526) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Why are the changes needed? ### Which issue does this PR close? Related to #8668 and #8609; follows #24074 and the merged statistics-ordering fix #24525. ### Rationale for this change Queries often select a sparse set of string identifiers. Parquet min/max statistics can make these queries much cheaper by ruling out row groups or pages that cannot contain any requested identifier. For example, this query asks for 21 IDs, spaced ten apart: ```sql SET datafusion.execution.parquet.max_in_list_size = 1024; SELECT * FROM events WHERE customer_id IN ( 'id000', 'id010', 'id020', 'id030', 'id040', 'id050', 'id060', 'id070', 'id080', 'id090', 'id100', 'id110', 'id120', 'id130', 'id140', 'id150', 'id160', 'id170', 'id180', 'id190', 'id200' ); ``` A row group whose values fall between `id003` and `id007` cannot contain a match. The default pruning limit is 20, so this list is not eligible for the `IN` min/max rewrite unless the caller raises the limit. #24074 made that limit configurable. With a raised limit, DataFusion can already reject this row group, but it does so by constructing a growing expression tree resembling: ```text (min <= id000 AND id000 <= max) OR (min <= id010 AND id010 <= max) OR ... ``` For hundreds or thousands of identifiers, building and evaluating that tree can become expensive in its own right. Replacing the list with one enclosing range, `[id000, id200]`, would be cheaper, but would lose the gaps: that broad range overlaps `[id003, id007]` even though none of the requested IDs is present there. The aim is to keep the useful pruning precision of the existing per-value checks while making large lists cheaper to represent and evaluate. In the included local microbenchmark, evaluating 1,024 values against 4,096 intervals falls from 68.8 ms to 0.198 ms. This measures pruning work only, not end-to-end query speedup. ## What changes were proposed in this PR? ### What changes are included in this PR? Eligible large string lists are stored as a sorted, deduplicated set of values inside the pruning predicate, rather than expanded into one comparison branch per value. For each inclusive statistics interval, DataFusion finds the first requested value at or after the interval's minimum, then checks whether that value is also at or before its maximum. In the example, the first requested ID at or after `id003` is `id010`. Since `id010 > id007`, the row group can be skipped. An interval such as `[id019, id021]` must be kept because it could contain `id020`. This takes a binary search per interval after sorting the values once per constructed predicate, and the expression tree no longer grows with the number of IDs. The result remains a conservative pruning decision. An overlapping interval means only that a match is possible; the original `IN` expression still performs exact row filtering. The original literal information also remains available to other pruning mechanisms, including Bloom filters. Inverted or unusable bounds yield UNKNOWN. A single known bound can still exclude the entire domain; otherwise, incomplete bounds yield UNKNOWN and keep the container eligible for reading. **Limits and scope** For a nonempty input list of `N` values, before this optimization deduplicates it, and configured cap `C = datafusion.execution.parquet.max_in_list_size`: | Condition | Min/max pruning behavior | | --- | --- | | `N <= min(20, C)` | Existing per-value rewrite | | `20 < N <= C`, positive, non-null literal strings on a direct string column | Compact sorted-domain lookup | | `20 < N <= C`, other lists | Existing per-value rewrite, including non-string lists, `NOT IN`, and NULL-containing lists | | `N > C` | This `IN` min/max rewrite is skipped; it does not disable all pruning | The compact path supports `Utf8`, `LargeUtf8`, `Utf8View`, and dictionary-encoded string columns. Unsupported expressions keep their existing handling. The default cap remains 20, and `C = 0` disables this `IN` min/max rewrite. Other predicates and literal/containment pruning, including Bloom filters, remain available. Raising the cap can still produce expensive comparison trees for lists that do not use the compact path. The lower threshold of 20 preserves the existing expression shape and limits this PR's scope; it is not a measured performance crossover. Page-index pruning now receives the same configured cap as row-group pruning, so raising the limit can benefit both. An interval search requires bounds in the same comparison order as the query. The compact path reads the same statistics adapters as the existing per-value path. The already-merged #24525 masks untrusted Parquet byte-array bounds. For other statistics sources, `PruningStatistics` requires conservative bounds in Arrow's comparison order; arbitrary provider bounds are trusted rather than validated by this expression. ### Are there any user-facing changes? Users who raise `datafusion.execution.parquet.max_in_list_size` get cheaper min/max pruning for eligible large string lists, and page pruning now honors that setting. The configuration default, exact query results, and existing public APIs are unchanged. This is a focused optimization for literal string lists, not a general rewrite of every large `IN` expression. ## How was this PR tested? ### Are these changes tested? The pruning-crate suite passed 95 tests. The standalone Parquet regressions use lists of 20, 21, 256, and 1,024 values and check both exact query results and scan/pruning metrics. They cover gaps inside the list's enclosing range, row-group pruning, page-only pruning, and the default and zero-cap controls. Two correctness regressions exercise the less obvious interactions. A direct physical-source test uses `NOT IN (..., NULL)`, row-filter pushdown, and `LIMIT 1`, so logical optimizer folding cannot hide an incorrect decision to bypass filtering. A real-file test combines this PR with #24525 and verifies that compact page pruning cannot lose a matching row when the footer's ordering is missing or unknown. Against unchanged Apache `f1f0449a`, the positive row-group/page tests fail as expected; the `NOT IN (..., NULL)` control passes. The benchmark compares the actual raised-cap `IN` path on Apache `f1f0449a` and this patch, using separate build directories and checking that both return the same nontrivial pruning results. Local results on an Apple M5 Max (18 CPUs, 128 GiB), Rust 1.97.0, `release-nonlto`, 20 samples: | Values | Predicate construction, main → PR | Evaluate 4,096 intervals, main → PR | | --- | --- | --- | | 20 | 31.1 → 27.4 µs | 177.6 → 165.7 µs | | 21 | 30.5 → 3.52 µs | 183.8 → 98.1 µs | | 256 | 356.6 → 21.4 µs | 5.53 → 0.144 ms | | 1,024 | 1.466 → 0.081 ms | 68.8 → 0.198 ms | A balanced explicit OR tree is included as another comparison: at 1,024 values it takes 8.03 ms to evaluate the same intervals. These measurements isolate pruning overhead; no end-to-end workload improvement is claimed. The latest documentation/comment update passed configuration-doc generation, formatting, all-targets/all-features Clippy with warnings denied, and `./dev/rust_lint.sh` (including the workspace Rust documentation check). Focused validation also passed all 3 Parquet IN-list integration tests and all 9 Parquet statistics-ordering tests. The CI follow-up in [1a2b38c](https://github.com/apache/datafusion/pull/24526/commits/1a2b38c4dc4c5730e451ee281609daf552503ccf) updates the `SHOW ALL VERBOSE` snapshot to match the revised `max_in_list_size` help text. The stale row caused both the cargo test and benchmark verification jobs to fail. After the correction, the targeted `information_schema.slt` test and full extended suite passed: 10,796 Rust tests (8 ignored) and all 505 SQL logic files. Formatting, Clippy with all targets and features and warnings denied, and `./dev/rust_lint.sh` also passed.
Validation commands ```sh cargo test --locked --profile ci -p datafusion-pruning cargo test --locked --profile ci -p datafusion \ --test parquet_integration string_in_list_pruning cargo bench --locked --profile release-nonlto -p datafusion-pruning \ --bench string_in_list_pruning -- \ --sample-size 20 --warm-up-time 0.5 --measurement-time 1 --noplot RUST_BACKTRACE=1 cargo test --locked --profile ci \ --exclude datafusion-examples --exclude datafusion-benchmarks --exclude datafusion-cli \ --workspace --lib --tests --bins \ --features avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption ```
Backport provenance: - Upstream PR: https://github.com/apache/datafusion/pull/24526 - Upstream commit: 9ad1217c43888820d7519bfcb9bcdb09c00e506e - Upstream self-tests and all additions/removals retained; only unavoidable fork patch context resolved. - Any fork-specific adaptation is kept in a separate following commit. Signed-off-by: discord9 --- Cargo.lock | 1 + datafusion/common/src/config.rs | 14 +- datafusion/common/src/pruning.rs | 11 + datafusion/core/tests/parquet/mod.rs | 1 + .../tests/parquet/string_in_list_pruning.rs | 363 +++++++++++++++ .../datasource-parquet/src/opener/mod.rs | 19 +- .../datasource-parquet/src/page_filter.rs | 16 +- datafusion/datasource-parquet/src/source.rs | 5 +- .../src/statistics_order_tests.rs | 68 +++ datafusion/pruning/Cargo.toml | 5 + .../pruning/benches/string_in_list_pruning.rs | 241 ++++++++++ datafusion/pruning/src/lib.rs | 1 + datafusion/pruning/src/pruning_predicate.rs | 423 +++++++++++++++++- datafusion/pruning/src/string_in_list.rs | 239 ++++++++++ .../test_files/information_schema.slt | 2 +- docs/source/user-guide/configs.md | 2 +- 16 files changed, 1374 insertions(+), 37 deletions(-) create mode 100644 datafusion/core/tests/parquet/string_in_list_pruning.rs create mode 100644 datafusion/pruning/benches/string_in_list_pruning.rs create mode 100644 datafusion/pruning/src/string_in_list.rs diff --git a/Cargo.lock b/Cargo.lock index e5c5691f5d03c..8d0821fb8a0e0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2588,6 +2588,7 @@ name = "datafusion-pruning" version = "55.1.0" dependencies = [ "arrow", + "criterion", "datafusion-common", "datafusion-datasource", "datafusion-expr", diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index f5742f09f9b08..561445cb565e3 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1282,13 +1282,15 @@ config_namespace! { /// parquet reader setting. 0 means no caching. pub max_predicate_cache_size: Option, default = None - /// Maximum number of values in an `IN (...)` list for which pruning will - /// occur. Longer lists will not be used to prune files, row groups, or - /// data pages. + /// Maximum number of input values in an `IN (...)` list eligible for + /// min/max pruning. Lists above this cap, or a cap of 0, skip this + /// rewrite; other predicates and Bloom-filter pruning remain available. /// - /// Higher values help in cases such as filtering on a list of - /// ~25-100 identifiers, but also make the predicate more expensive to - /// evaluate. Set to 0 to disable `IN (...)` list pruning entirely. + /// Within the cap, nonempty lists of at most 20 values use the existing + /// per-value rewrite. Larger positive, non-null literal string lists + /// on a string column use a compact sorted domain. Other lists retain + /// the existing per-value rewrite, so raising the cap can make those + /// predicates expensive to build and evaluate. /// /// Defaults to 20. pub max_in_list_size: usize, default = 20 diff --git a/datafusion/common/src/pruning.rs b/datafusion/common/src/pruning.rs index a36ac9f795b95..9f6a95c0978d4 100644 --- a/datafusion/common/src/pruning.rs +++ b/datafusion/common/src/pruning.rs @@ -67,6 +67,13 @@ pub trait PruningStatistics { /// returned array should have `null` in that row. If the minimum value is /// not known for any row, return `None`. /// + /// Each non-null entry must be a conservative lower bound for all non-null + /// values in its container, using Arrow's comparison order for the column + /// type. For strings, this is unsigned lexicographic UTF-8 byte order. + /// Bounds that cannot satisfy this contract must be reported as unknown. + /// Pruning relies on providers to uphold this contract; it does not validate + /// the ordering or bound guarantees of arbitrary statistics sources. + /// /// Note: the returned array must contain [`Self::num_containers`] rows fn min_values(&self, column: &Column) -> Option; @@ -74,6 +81,10 @@ pub trait PruningStatistics { /// /// See [`Self::min_values`] for when to return `None` and null values. /// + /// Each non-null entry must be a conservative upper bound for all non-null + /// values in its container, using the comparison order described in + /// [`Self::min_values`]. + /// /// Note: the returned array must contain [`Self::num_containers`] rows fn max_values(&self, column: &Column) -> Option; diff --git a/datafusion/core/tests/parquet/mod.rs b/datafusion/core/tests/parquet/mod.rs index 7066a4147c017..cf2e557abad5f 100644 --- a/datafusion/core/tests/parquet/mod.rs +++ b/datafusion/core/tests/parquet/mod.rs @@ -58,6 +58,7 @@ mod page_pruning; mod row_group_pruning; mod schema; mod schema_coercion; +mod string_in_list_pruning; mod utils; #[cfg(test)] diff --git a/datafusion/core/tests/parquet/string_in_list_pruning.rs b/datafusion/core/tests/parquet/string_in_list_pruning.rs new file mode 100644 index 0000000000000..e1005daaa1a38 --- /dev/null +++ b/datafusion/core/tests/parquet/string_in_list_pruning.rs @@ -0,0 +1,363 @@ +// 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. + +//! End-to-end coverage for compact, large string IN-list pruning. The positive +//! IN-list cases disable the row and Bloom filters to isolate min/max pruning. + +use std::sync::Arc; + +use arrow::array::StringArray; +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::record_batch::RecordBatch; +use arrow::util::pretty::pretty_format_batches; +use datafusion::datasource::listing::PartitionedFile; +use datafusion::datasource::object_store::ObjectStoreUrl; +use datafusion::datasource::physical_plan::ParquetSource; +use datafusion::datasource::source::DataSourceExec; +use datafusion::physical_plan::{ExecutionPlan, collect, displayable}; +use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; +use datafusion_common::config::TableParquetOptions; +use datafusion_common::{ScalarValue, assert_batches_eq}; +use datafusion_datasource::file_scan_config::FileScanConfigBuilder; +use datafusion_physical_expr::expressions::{col, in_list, lit}; +use datafusion_physical_plan::metrics::{MetricValue, MetricsSet}; +use object_store::path::Path; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::{EnabledStatistics, WriterProperties}; +use tempfile::NamedTempFile; + +use super::utils::MetricsFinder; + +const ROWS_PER_UNIT: usize = 16; +const UNITS: usize = 4; +const TOTAL_ROWS: usize = ROWS_PER_UNIT * UNITS; +const MATCHING_ROWS: usize = ROWS_PER_UNIT * 2; + +/// Write either four row groups or four pages in one row group. The second +/// unit lies in a gap between two members of every test IN list; an enclosing +/// min/max range for the list cannot prune it. +fn make_file(page_pruning: bool) -> NamedTempFile { + let mut file = tempfile::Builder::new() + .prefix("string_in_list_pruning") + .suffix(".parquet") + .tempfile() + .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Utf8, + false, + )])); + let values = ["v000000", "v000001", "v000010", "v999999"] + .into_iter() + .flat_map(|value| std::iter::repeat_n(value, ROWS_PER_UNIT)) + .collect::>(); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(StringArray::from(values))], + ) + .unwrap(); + let rows_per_group = if page_pruning { + TOTAL_ROWS + } else { + ROWS_PER_UNIT + }; + let properties = WriterProperties::builder() + .set_max_row_group_row_count(Some(rows_per_group)) + .set_data_page_row_count_limit(ROWS_PER_UNIT) + .set_write_batch_size(ROWS_PER_UNIT) + .set_dictionary_enabled(false) + .set_bloom_filter_enabled(false) + .set_statistics_enabled(EnabledStatistics::Page) + .build(); + let mut writer = ArrowWriter::try_new(&mut file, schema, Some(properties)).unwrap(); + writer.write(&batch).unwrap(); + let metadata = writer.close().unwrap(); + assert_eq!(metadata.num_row_groups(), TOTAL_ROWS / rows_per_group); + let offsets = metadata.offset_index().unwrap(); + for row_group in offsets { + assert_eq!( + row_group[0].page_locations().len(), + rows_per_group / ROWS_PER_UNIT + ); + } + file +} + +struct ScanOutput { + batches: Vec, + plan: String, + metrics: MetricsSet, +} + +impl ScanOutput { + fn counter(&self, name: &str) -> usize { + self.metrics + .sum(|metric| metric.value().name() == name) + .unwrap_or_else(|| panic!("missing {name}: {}", self.metrics)) + .as_usize() + } + + fn pruned(&self, name: &str) -> usize { + let value = self + .metrics + .sum(|metric| metric.value().name() == name) + .unwrap_or_else(|| panic!("missing {name}: {}", self.metrics)); + let MetricValue::PruningMetrics { + pruning_metrics, .. + } = value + else { + panic!("expected pruning metric {name}: {}", self.metrics); + }; + pruning_metrics.pruned() + } + + fn fully_matched(&self, name: &str) -> usize { + let value = self + .metrics + .sum(|metric| metric.value().name() == name) + .unwrap_or_else(|| panic!("missing {name}: {}", self.metrics)); + let MetricValue::PruningMetrics { + pruning_metrics, .. + } = value + else { + panic!("expected pruning metric {name}: {}", self.metrics); + }; + pruning_metrics.fully_matched() + } + + fn assert_results(&self) { + assert_batches_eq!( + [ + "+---------+----+", + "| value | n |", + "+---------+----+", + "| v000000 | 16 |", + "| v000010 | 16 |", + "+---------+----+", + ], + &self.batches + ); + assert_eq!(self.counter("predicate_evaluation_errors"), 0); + assert_eq!(self.counter("pushdown_rows_pruned"), 0); + assert_eq!(self.pruned("row_groups_pruned_bloom_filter"), 0); + } +} + +async fn scan( + file: &NamedTempFile, + list_size: usize, + max_in_list_size: Option, + page_pruning: bool, +) -> ScanOutput { + let mut config = SessionConfig::new() + .with_target_partitions(1) + .with_parquet_bloom_filter_pruning(false) + .with_parquet_page_index_pruning(page_pruning); + config.options_mut().execution.parquet.pushdown_filters = false; + if let Some(max_in_list_size) = max_in_list_size { + config.options_mut().execution.parquet.max_in_list_size = max_in_list_size; + } + let ctx = SessionContext::new_with_config(config); + ctx.register_parquet( + "t", + file.path().to_str().unwrap(), + ParquetReadOptions::default(), + ) + .await + .unwrap(); + let values = (0..list_size) + .map(|index| format!("'v{:06}'", index * 10)) + .collect::>() + .join(", "); + let sql = format!( + "SELECT value, count(*) AS n FROM t \ + WHERE value IN ({values}) GROUP BY value ORDER BY value" + ); + let plan = ctx + .sql(&sql) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let plan_text = displayable(plan.as_ref()).indent(true).to_string(); + let batches = collect(Arc::clone(&plan), ctx.task_ctx()).await.unwrap(); + let metrics = MetricsFinder::find_metrics(plan.as_ref()).unwrap(); + ScanOutput { + batches, + plan: plan_text, + metrics, + } +} + +async fn check_string_in_list_pruning(page_pruning: bool) { + let file = make_file(page_pruning); + for list_size in [20, 21, 256, 1024] { + // A zero cap provides a result-equivalence control that cannot use + // min/max IN-list pruning at either granularity. + let unpruned = scan(&file, list_size, Some(0), page_pruning).await; + unpruned.assert_results(); + assert!(!unpruned.plan.contains("IN_SET_INTERSECTS")); + assert_eq!(unpruned.pruned("row_groups_pruned_statistics"), 0); + assert_eq!(unpruned.pruned("page_index_rows_pruned"), 0); + assert_eq!(unpruned.counter("output_rows"), TOTAL_ROWS); + + let output = scan(&file, list_size, Some(list_size), page_pruning).await; + output.assert_results(); + assert_eq!( + pretty_format_batches(&output.batches).unwrap().to_string(), + pretty_format_batches(&unpruned.batches) + .unwrap() + .to_string() + ); + assert_eq!( + output.plan.contains("IN_SET_INTERSECTS"), + list_size > 20, + "list_size={list_size}, plan={}", + output.plan + ); + assert_eq!( + output.pruned("row_groups_pruned_statistics"), + if page_pruning { 0 } else { 2 }, + "list_size={list_size}, metrics={}", + output.metrics + ); + assert_eq!( + output.pruned("page_index_rows_pruned"), + if page_pruning { MATCHING_ROWS } else { 0 }, + "list_size={list_size}, metrics={}", + output.metrics + ); + assert_eq!(output.counter("output_rows"), MATCHING_ROWS); + } + + // The default remains 20: enabling the compact representation must not + // silently change the public cap's meaning. + let default = scan(&file, 21, None, page_pruning).await; + default.assert_results(); + assert!(!default.plan.contains("IN_SET_INTERSECTS")); + assert_eq!(default.pruned("row_groups_pruned_statistics"), 0); + assert_eq!(default.pruned("page_index_rows_pruned"), 0); + assert_eq!(default.counter("output_rows"), TOTAL_ROWS); +} + +#[tokio::test] +async fn string_in_list_row_group_pruning() { + check_string_in_list_pruning(false).await; +} + +#[tokio::test] +async fn string_in_list_page_pruning() { + check_string_in_list_pruning(true).await; +} + +#[tokio::test] +async fn string_not_in_list_with_null_does_not_bypass_row_filter() { + let mut file = tempfile::Builder::new() + .prefix("string_not_in_list_pruning") + .suffix(".parquet") + .tempfile() + .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Utf8, true)])); + // The first row group has a known zero null count, and every value lies + // in a gap in the IN list. Dropping the NULL list member while inverting + // NOT IN would incorrectly prove that this entire row group matches. + let values = vec![ + Some("v000001"), + Some("v000001"), + Some("v000001"), + Some("v000001"), + Some("v000000"), + Some("v000001"), + None, + Some("v999999"), + ]; + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(StringArray::from(values))], + ) + .unwrap(); + let properties = WriterProperties::builder() + .set_max_row_group_row_count(Some(4)) + .set_bloom_filter_enabled(false) + .build(); + let mut writer = + ArrowWriter::try_new(&mut file, Arc::clone(&schema), Some(properties)).unwrap(); + writer.write(&batch).unwrap(); + assert_eq!(writer.close().unwrap().num_row_groups(), 2); + + // Build the physical source directly so a logical optimizer cannot fold + // the SQL NOT IN (..., NULL) filter to an empty relation before the scan. + let mut list = (0..21) + .map(|index| lit(format!("v{:06}", index * 10))) + .collect::>(); + list.push(lit(ScalarValue::Utf8(None))); + let predicate = + in_list(col("value", &schema).unwrap(), list, &true, &schema).unwrap(); + let location = Path::from_filesystem_path(file.path()).unwrap(); + let partitioned_file = PartitionedFile::new( + location.to_string(), + file.as_file().metadata().unwrap().len(), + ); + let ctx = + SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); + + for max_in_list_size in [0, 32] { + let mut options = TableParquetOptions::default(); + options.global.max_in_list_size = max_in_list_size; + let source = Arc::new( + ParquetSource::new(Arc::clone(&schema)) + .with_table_parquet_options(options) + .with_predicate(Arc::clone(&predicate)) + .with_pushdown_filters(true) + .with_enable_page_index(false) + .with_bloom_filter_on_read(false), + ); + let config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + .with_file(partitioned_file.clone()) + .with_limit(Some(1)) + .build(); + let plan: Arc = + Arc::new(DataSourceExec::new(Arc::new(config))); + let plan_text = displayable(plan.as_ref()).indent(true).to_string(); + assert!(plan_text.contains("NOT IN"), "{plan_text}"); + let batches = collect(Arc::clone(&plan), ctx.task_ctx()).await.unwrap(); + let output = ScanOutput { + batches, + plan: plan_text, + metrics: MetricsFinder::find_metrics(plan.as_ref()).unwrap(), + }; + + assert_eq!( + output + .batches + .iter() + .map(RecordBatch::num_rows) + .sum::(), + 0, + "cap={max_in_list_size}, plan={}, metrics={}", + output.plan, + output.metrics + ); + assert_eq!(output.fully_matched("row_groups_pruned_statistics"), 0); + assert_eq!(output.pruned("row_groups_pruned_statistics"), 0); + assert_eq!(output.pruned("limit_pruned_row_groups"), 0); + assert_eq!(output.counter("pushdown_rows_pruned"), 8); + assert_eq!(output.counter("predicate_evaluation_errors"), 0); + } +} diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index 1c59f345e30bf..3f61069b3c0a3 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -289,9 +289,8 @@ pub(super) struct ParquetMorselizer { /// Maximum size of the predicate cache, in bytes. If none, uses /// the arrow-rs default. pub max_predicate_cache_size: Option, - /// Maximum `IN (...)` list size that the pruning predicate will rewrite - /// into per-value statistics checks. Lists longer than this skip - /// container-level pruning. Sourced from + /// Maximum `IN (...)` list size eligible for statistics pruning. Longer + /// lists skip container-level pruning. Sourced from /// `datafusion.execution.parquet.max_in_list_size`. pub max_in_list_size: usize, /// Whether to read row groups in reverse order @@ -1062,7 +1061,11 @@ impl MetadataLoadedParquetOpen { // Only build page pruning predicate if page index is enabled let page_pruning_predicate = if prepared.enable_page_index { prepared.predicate.as_ref().and_then(|predicate| { - let p = build_page_pruning_predicate(predicate, &physical_file_schema); + let p = build_page_pruning_predicate( + predicate, + &physical_file_schema, + prepared.max_in_list_size, + ); (p.filter_number() > 0).then_some(p) }) } else { @@ -1688,10 +1691,12 @@ fn create_initial_plan( pub(crate) fn build_page_pruning_predicate( predicate: &Arc, file_schema: &SchemaRef, + max_in_list_size: usize, ) -> Arc { - Arc::new(PagePruningAccessPlanFilter::new( + Arc::new(PagePruningAccessPlanFilter::new_with_max_in_list_size( predicate, - Arc::clone(file_schema), + file_schema, + max_in_list_size, )) } @@ -1965,7 +1970,7 @@ mod test { ); let page_pruning_predicate = predicate.map(|expr| { let predicate = logical2physical(&expr, &arrow_schema); - build_page_pruning_predicate(&predicate, &arrow_schema) + build_page_pruning_predicate(&predicate, &arrow_schema, MAX_IN_LIST_SIZE) }); let store: Arc = Arc::new(InMemory::new()); diff --git a/datafusion/datasource-parquet/src/page_filter.rs b/datafusion/datasource-parquet/src/page_filter.rs index 5372ae0a0007d..236a08c5608af 100644 --- a/datafusion/datasource-parquet/src/page_filter.rs +++ b/datafusion/datasource-parquet/src/page_filter.rs @@ -32,7 +32,7 @@ use arrow::{ use datafusion_common::ScalarValue; use datafusion_common::pruning::PruningStatistics; use datafusion_physical_expr::{PhysicalExpr, split_conjunction}; -use datafusion_pruning::{PruningPredicate, PruningPredicateBuilder}; +use datafusion_pruning::{MAX_IN_LIST_SIZE, PruningPredicate, PruningPredicateBuilder}; use log::{debug, trace}; use parquet::arrow::arrow_reader::statistics::StatisticsConverter; @@ -138,15 +138,25 @@ impl PagePruningResult { impl PagePruningAccessPlanFilter { /// Create a new [`PagePruningAccessPlanFilter`] from a physical - /// expression. + /// expression, using the default `IN (...)` pruning limit. #[expect(clippy::needless_pass_by_value)] pub fn new(expr: &Arc, schema: SchemaRef) -> Self { + Self::new_with_max_in_list_size(expr, &schema, MAX_IN_LIST_SIZE) + } + + /// Create a page filter using the same `IN (...)` limit as row-group pruning. + pub(crate) fn new_with_max_in_list_size( + expr: &Arc, + schema: &SchemaRef, + max_in_list_size: usize, + ) -> Self { // extract any single column predicates let predicates = split_conjunction(expr) .into_iter() .filter_map(|predicate| { let pp = match PruningPredicateBuilder::new() - .with_file_schema(Arc::clone(&schema)) + .with_file_schema(Arc::clone(schema)) + .with_max_in_list_size(max_in_list_size) .try_build(Arc::clone(predicate)) { Ok(pp) => pp, diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index 097b4563af5df..b29480f58379b 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -486,9 +486,8 @@ impl ParquetSource { self.table_parquet_options.global.max_predicate_cache_size } - /// Return the maximum size of an `IN (...)` list that the pruning - /// predicate will rewrite into per-value statistics checks. Lists - /// longer than this skip container-level pruning. Reads from + /// Return the maximum size of an `IN (...)` list eligible for statistics + /// pruning. Longer lists skip container-level pruning. Reads from /// `datafusion.execution.parquet.max_in_list_size`. pub fn max_in_list_size(&self) -> usize { self.table_parquet_options.global.max_in_list_size diff --git a/datafusion/datasource-parquet/src/statistics_order_tests.rs b/datafusion/datasource-parquet/src/statistics_order_tests.rs index bf9df532f1b58..5081abf2cb47b 100644 --- a/datafusion/datasource-parquet/src/statistics_order_tests.rs +++ b/datafusion/datasource-parquet/src/statistics_order_tests.rs @@ -359,6 +359,74 @@ fn byte_array_order_preserves_matching_rows_at_every_pruning_level() { } } +#[test] +fn large_string_in_list_preserves_rows_with_untrusted_page_order() { + let max_in_list_size = MAX_IN_LIST_SIZE + 2; + for order in [ + StatisticsOrder::Modern, + StatisticsOrder::Missing, + StatisticsOrder::Unknown, + ] { + let file = TestFile::new(order); + // Only "az" occurs in the file. All other list members are above + // even the unsafe ["aé", "b"] interval in the first page's index. + let mut values = (0..=MAX_IN_LIST_SIZE) + .map(|index| lit(format!("z{index:03}"))) + .collect::>(); + values.push(lit("az")); + let physical = logical2physical(&col("s").in_list(values, false), &file.schema); + let predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&file.schema)) + .with_max_in_list_size(max_in_list_size) + .try_build(Arc::clone(&physical)) + .unwrap(); + assert!( + predicate + .predicate_expr() + .to_string() + .contains("IN_SET_INTERSECTS") + ); + + // Exercise the opener's configured-cap path, not the page filter's + // compatibility constructor that retains the default limit of 20. + let page_filter = crate::opener::build_page_pruning_predicate( + &physical, + &file.schema, + max_in_list_size, + ); + assert_eq!(page_filter.filter_number(), 1); + let all = ParquetAccessPlan::new_all(file.metadata.num_row_groups()); + assert_eq!(file.matching_rows(&physical, all.clone()), 1); + let file_metrics = metrics(); + let pages = page_filter.prune_plan_with_page_index( + all, + &file.schema, + file.metadata.file_metadata().schema_descr(), + &file.metadata, + &file_metrics, + ); + assert_eq!( + pages.row_group_indexes(), + if order == StatisticsOrder::Modern { + vec![0] + } else { + vec![0, 2] + }, + "order={order:?}", + ); + assert_eq!(file.matching_rows(&physical, pages), 1, "order={order:?}",); + assert_eq!( + file_metrics.page_index_rows_pruned.pruned(), + if order == StatisticsOrder::Modern { + 6 + } else { + 3 + }, + "order={order:?}", + ); + } +} + #[test] fn byte_array_order_keeps_null_counts_and_unrelated_column_bounds() { for order in [ diff --git a/datafusion/pruning/Cargo.toml b/datafusion/pruning/Cargo.toml index e6f4bb6f273c9..a703f68222f38 100644 --- a/datafusion/pruning/Cargo.toml +++ b/datafusion/pruning/Cargo.toml @@ -26,7 +26,12 @@ datafusion-physical-plan = { workspace = true } log = { workspace = true } [dev-dependencies] +criterion = { workspace = true } datafusion-expr = { workspace = true } datafusion-functions-nested = { workspace = true } insta = { workspace = true } itertools = { workspace = true } + +[[bench]] +harness = false +name = "string_in_list_pruning" diff --git a/datafusion/pruning/benches/string_in_list_pruning.rs b/datafusion/pruning/benches/string_in_list_pruning.rs new file mode 100644 index 0000000000000..d26112611bf3c --- /dev/null +++ b/datafusion/pruning/benches/string_in_list_pruning.rs @@ -0,0 +1,241 @@ +// 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. + +//! Compare compact string IN-list pruning with per-value min/max expansion. +//! +//! Both cases raise `max_in_list_size` to the domain size. On a baseline +//! without compact pruning, `in_list` measures the ordinary raised-cap path. +//! The explicit `expanded_or` is a balanced tree of equalities, which produces +//! the same per-value statistics checks without making the baseline depend on +//! a deeply nested expression. Half of the statistics intervals hit a domain +//! member and half fall in a sparse gap. Bloom filters are not involved. +//! +//! Run with `cargo bench -p datafusion-pruning --bench string_in_list_pruning`. +//! The construction benchmarks reuse their input physical expressions; the +//! evaluation benchmarks reuse their already-built pruning predicates. + +use std::collections::HashSet; +use std::hint::black_box; +use std::sync::Arc; + +use arrow::array::{ArrayRef, BooleanArray, StringViewArray, UInt64Array}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use datafusion_common::{Column, ScalarValue}; +use datafusion_expr_common::operator::Operator; +use datafusion_physical_expr::PhysicalExprRef; +use datafusion_physical_expr::expressions::{BinaryExpr, col, in_list, lit}; +use datafusion_pruning::{PruningPredicate, PruningPredicateBuilder, PruningStatistics}; + +const DOMAIN_SIZES: [usize; 4] = [20, 21, 256, 1024]; +const CONTAINERS: usize = 4096; + +fn value(index: usize) -> String { + format!("key{index:08}") +} + +fn balanced_or(expressions: &[PhysicalExprRef]) -> PhysicalExprRef { + if expressions.len() == 1 { + return Arc::clone(&expressions[0]); + } + let middle = expressions.len() / 2; + Arc::new(BinaryExpr::new( + balanced_or(&expressions[..middle]), + Operator::Or, + balanced_or(&expressions[middle..]), + )) +} + +fn build_predicate( + expression: &PhysicalExprRef, + schema: &SchemaRef, + max_in_list_size: usize, +) -> PruningPredicate { + PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(schema)) + .with_max_in_list_size(max_in_list_size) + .try_build(Arc::clone(expression)) + .unwrap() +} + +struct IntervalStatistics { + min: ArrayRef, + max: ArrayRef, + null_counts: ArrayRef, + row_counts: ArrayRef, +} + +impl IntervalStatistics { + fn new(domain_size: usize) -> Self { + let min = StringViewArray::from_iter_values((0..CONTAINERS).map(|index| { + let start = (index / 2 % domain_size) * 10; + value(start + if index % 2 == 0 { 0 } else { 3 }) + })); + let max = StringViewArray::from_iter_values((0..CONTAINERS).map(|index| { + let start = (index / 2 % domain_size) * 10; + value(start + if index % 2 == 0 { 0 } else { 7 }) + })); + Self { + min: Arc::new(min), + max: Arc::new(max), + null_counts: Arc::new(UInt64Array::from(vec![0; CONTAINERS])), + row_counts: Arc::new(UInt64Array::from(vec![128; CONTAINERS])), + } + } +} + +impl PruningStatistics for IntervalStatistics { + fn min_values(&self, column: &Column) -> Option { + (column.name == "value").then(|| Arc::clone(&self.min)) + } + + fn max_values(&self, column: &Column) -> Option { + (column.name == "value").then(|| Arc::clone(&self.max)) + } + + fn num_containers(&self) -> usize { + CONTAINERS + } + + fn null_counts(&self, column: &Column) -> Option { + (column.name == "value").then(|| Arc::clone(&self.null_counts)) + } + + fn row_counts(&self) -> Option { + Some(Arc::clone(&self.row_counts)) + } + + fn contained( + &self, + _column: &Column, + _values: &HashSet, + ) -> Option { + None + } +} + +struct BenchmarkCase { + size: usize, + schema: SchemaRef, + in_list: PhysicalExprRef, + expanded_or: PhysicalExprRef, + in_list_predicate: PruningPredicate, + expanded_or_predicate: PruningPredicate, + statistics: IntervalStatistics, +} + +impl BenchmarkCase { + fn new(size: usize) -> Self { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Utf8View, + false, + )])); + let column = col("value", &schema).unwrap(); + let values = (0..size) + .map(|index| lit(ScalarValue::new_utf8view(value(index * 10)))) + .collect::>(); + let in_list = + in_list(Arc::clone(&column), values.clone(), &false, &schema).unwrap(); + let equalities = values + .into_iter() + .map(|value| { + Arc::new(BinaryExpr::new(Arc::clone(&column), Operator::Eq, value)) + as PhysicalExprRef + }) + .collect::>(); + let expanded_or = balanced_or(&equalities); + let in_list_predicate = build_predicate(&in_list, &schema, size); + let expanded_or_predicate = build_predicate(&expanded_or, &schema, size); + eprintln!( + "string_in_list_pruning: {size} values, compact={}", + in_list_predicate + .predicate_expr() + .to_string() + .contains("IN_SET_INTERSECTS") + ); + let statistics = IntervalStatistics::new(size); + + // Check that both benchmark paths do the same useful work, rather + // than comparing compact pruning with an always-true fallback. + let expected = (0..CONTAINERS) + .map(|index| index % 2 == 0) + .collect::>(); + assert_eq!(in_list_predicate.prune(&statistics).unwrap(), expected); + assert_eq!(expanded_or_predicate.prune(&statistics).unwrap(), expected); + + Self { + size, + schema, + in_list, + expanded_or, + in_list_predicate, + expanded_or_predicate, + statistics, + } + } +} + +fn criterion_benchmark(criterion: &mut Criterion) { + let cases = DOMAIN_SIZES.map(BenchmarkCase::new); + let mut construction = criterion.benchmark_group("string_in_list_pruning/construct"); + for case in &cases { + construction.throughput(Throughput::Elements(case.size as u64)); + for (name, expression) in [ + ("in_list", &case.in_list), + ("expanded_or", &case.expanded_or), + ] { + construction.bench_with_input( + BenchmarkId::new(name, case.size), + expression, + |bencher, expression| { + bencher.iter(|| { + black_box(build_predicate( + black_box(expression), + &case.schema, + case.size, + )) + }); + }, + ); + } + } + construction.finish(); + + let mut evaluation = criterion.benchmark_group("string_in_list_pruning/evaluate"); + evaluation.throughput(Throughput::Elements(CONTAINERS as u64)); + for case in &cases { + for (name, predicate) in [ + ("in_list", &case.in_list_predicate), + ("expanded_or", &case.expanded_or_predicate), + ] { + evaluation.bench_with_input( + BenchmarkId::new(name, case.size), + predicate, + |bencher, predicate| { + bencher.iter(|| { + black_box(predicate.prune(black_box(&case.statistics)).unwrap()) + }); + }, + ); + } + } + evaluation.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/pruning/src/lib.rs b/datafusion/pruning/src/lib.rs index 2b334d2847980..6bf1815900aa8 100644 --- a/datafusion/pruning/src/lib.rs +++ b/datafusion/pruning/src/lib.rs @@ -19,6 +19,7 @@ mod file_pruner; mod pruning_predicate; +mod string_in_list; pub use file_pruner::FilePruner; pub use pruning_predicate::{ diff --git a/datafusion/pruning/src/pruning_predicate.rs b/datafusion/pruning/src/pruning_predicate.rs index 3a63451495e4c..f9b19c3959265 100644 --- a/datafusion/pruning/src/pruning_predicate.rs +++ b/datafusion/pruning/src/pruning_predicate.rs @@ -22,6 +22,8 @@ use std::collections::HashSet; use std::sync::Arc; +use crate::string_in_list::StringInListPruningExpr; + use arrow::array::AsArray; use arrow::{ array::{ArrayRef, BooleanArray, new_null_array}, @@ -441,9 +443,24 @@ impl<'a> PruningPredicateBuilder<'a> { self } - /// Cap on the size of `IN (...)` lists that will be rewritten into per- - /// value min/max statistics checks. Lists longer than this fall back to - /// the unhandled-predicate hook (typically "keep the container"). + /// Cap on the input length of `IN (...)` lists eligible for min/max pruning. + /// + /// For a nonempty list of `N` values (before deduplication) and cap `C`: + /// + /// | Condition | Pruning representation | + /// | --- | --- | + /// | `N <= min(20, C)` | Existing per-value rewrite | + /// | `20 < N <= C`, positive, non-null literal strings on a string column | Compact sorted domain | + /// | `20 < N <= C`, other lists | Existing per-value rewrite | + /// | `N > C` | Unhandled-predicate hook, normally "keep the container" | + /// + /// Empty lists also use the unhandled-predicate hook. A cap of zero disables + /// this `IN` min/max rewrite, not other predicates or literal/containment + /// pruning (such as Bloom filters). The default cap is [`MAX_IN_LIST_SIZE`] + /// (20), so the compact path requires an explicitly raised cap. + /// + /// Raising the cap can still build large comparison trees for non-string + /// lists, `NOT IN`, or lists containing NULL; their handling is unchanged. /// /// Query engines typically pass /// `datafusion.execution.parquet.max_in_list_size` here. @@ -1458,10 +1475,51 @@ fn build_is_null_column_expr( } } -/// Default maximum number of entries in an `IN (...)` list that will be -/// rewritten into a chain of per-value min/max checks by -/// `build_predicate_expression`. Callers threading a [`PredicateRewriter`] -/// can override this via [`PredicateRewriter::with_max_in_list_size`], and +/// Keep large literal string domains compact instead of building an OR tree. +fn build_string_in_list_expr( + in_list: &phys_expr::InListExpr, + schema: &Schema, + required_columns: &mut RequiredColumns, +) -> Option> { + if in_list.negated() { + return None; + } + let column = in_list.expr().downcast_ref::()?; + let field = schema.fields().get(column.index())?; + let data_type = match field.data_type() { + DataType::Dictionary(_, value) => value.as_ref(), + data_type => data_type, + }; + if field.name() != column.name() || !data_type.is_string() { + return None; + } + // NULLs must remain unhandled: the inverse predicate is also used to prove + // that every row matches, and IN (..., NULL) can evaluate to UNKNOWN. + let values = in_list + .list() + .iter() + .map(|expr| extract_string_literal(expr).map(str::to_owned)) + .collect::>>()?; + let min = required_columns + .min_column_expr(column, in_list.expr(), field) + .ok()?; + let max = required_columns + .max_column_expr(column, in_list.expr(), field) + .ok()?; + let non_null = + build_is_null_column_expr(in_list.expr(), schema, required_columns, true)?; + let intersects = Arc::new(StringInListPruningExpr::new(min, max, values)); + Some(Arc::new(phys_expr::BinaryExpr::new( + non_null, + Operator::And, + intersects, + ))) +} + +/// Default maximum number of entries in an `IN (...)` list eligible for +/// statistics pruning. Eligible positive literal string lists above this +/// threshold use a compact sorted domain instead of per-value min/max checks. +/// Callers can raise the cap via [`PredicateRewriter::with_max_in_list_size`], and /// query engines can wire it from the /// `datafusion.execution.parquet.max_in_list_size` config option. pub const MAX_IN_LIST_SIZE: usize = 20; @@ -1497,14 +1555,15 @@ impl PredicateRewriter { self } - /// Set the maximum size of an `IN (...)` list that will be rewritten into a - /// chain of per-value statistics checks. Lists longer than this fall back - /// to the unhandled-predicate hook (typically "keep the container"), - /// effectively skipping container-level pruning for large IN lists. + /// Set the maximum input length of an `IN (...)` list eligible for min/max + /// pruning. See [`PruningPredicateBuilder::with_max_in_list_size`] for the + /// exact boundaries and representations. Longer lists fall back to the + /// unhandled-predicate hook (typically "keep the container"); this skips + /// only that `IN` min/max rewrite, not all container-level pruning. /// /// The default (see [`MAX_IN_LIST_SIZE`]) preserves the /// historical behaviour. Callers wiring config through can override via - /// `datafusion.execution.max_in_list_size`. + /// `datafusion.execution.parquet.max_in_list_size`. pub fn with_max_in_list_size(mut self, max_in_list_size: usize) -> Self { self.max_in_list_size = max_in_list_size; self @@ -1544,8 +1603,9 @@ impl PredicateRewriter { /// /// Returns the pruning predicate as an [`PhysicalExpr`] /// -/// `max_in_list_size` is the largest `IN (...)` list that will be rewritten -/// into a chain of per-value statistics checks; longer lists fall back to +/// `max_in_list_size` is the largest `IN (...)` list eligible for statistics +/// pruning. Large positive literal string lists use a compact sorted domain; +/// other eligible lists use per-value checks. Longer lists fall back to /// `unhandled_hook`. fn build_predicate_expression( expr: &Arc, @@ -1587,6 +1647,16 @@ fn build_predicate_expression( } } if let Some(in_list) = expr.downcast_ref::() { + // Keep the existing expression shape for lists of at most 20 values. + // This lower bound is a scope/compatibility choice, not a measured + // performance threshold; compact pruning is opt-in via a raised cap. + if in_list.list().len() > MAX_IN_LIST_SIZE + && in_list.list().len() <= max_in_list_size + && let Some(pruning_expr) = + build_string_in_list_expr(in_list, schema, required_columns) + { + return pruning_expr; + } if !in_list.list().is_empty() && in_list.list().len() <= max_in_list_size { let eq_op = if in_list.negated() { Operator::NotEq @@ -2187,8 +2257,12 @@ mod tests { use arrow::array::Decimal128Array; use arrow::{ - array::{BinaryArray, Int32Array, Int64Array, StringArray, UInt64Array}, - datatypes::TimeUnit, + array::{ + BinaryArray, DictionaryArray, Int32Array, Int64Array, StringArray, + UInt64Array, + }, + buffer::NullBuffer, + datatypes::{Int32Type, TimeUnit}, }; use datafusion_expr::expr::InList; use datafusion_expr::{BinaryExpr, Expr, cast, is_null, try_cast}; @@ -3580,6 +3654,323 @@ mod tests { Ok(()) } + fn large_string_pruning_predicate( + expr: PhysicalExprRef, + schema: SchemaRef, + ) -> Result { + PruningPredicateBuilder::new() + .with_file_schema(schema) + .with_max_in_list_size(10_000) + .try_build(expr) + } + + #[test] + fn large_string_in_list_prunes_exact_intervals() -> Result<()> { + let types = [ + DataType::Utf8, + DataType::LargeUtf8, + DataType::Utf8View, + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + ]; + for count in [20, 21, 256, 10_000] { + for data_type in &types { + let schema = Arc::new(Schema::new(vec![Field::new( + "c1", + data_type.clone(), + true, + )])); + let values = (0..count) + .map(|i| { + let value = ScalarValue::from(format!("k{:06}", i * 10)) + .cast_to(data_type)?; + Ok(Arc::new(phys_expr::Literal::new(value)) as PhysicalExprRef) + }) + .collect::>>()?; + let expr = phys_expr::in_list( + Arc::new(phys_expr::Column::new("c1", 0)), + values, + &false, + &schema, + )?; + let predicate = + large_string_pruning_predicate(Arc::clone(&expr), schema)?; + let last_value = format!("k{:06}", (count - 1) * 10); + let stats = TestStatistics::new().with( + "c1", + ContainerStats::new_utf8( + [ + Some("k000000"), + Some("k000001"), + Some("k000010"), + Some("j"), + Some("z"), + None, + Some("k000020"), + Some("k000005"), + // Single bounds exclude the domain, but equality + // with either endpoint must remain eligible. + Some("z"), + None, + Some(last_value.as_str()), + None, + ], + [ + Some("k000000"), + Some("k000009"), + Some("k000010"), + Some("j9"), + Some("z9"), + Some("k000010"), + None, + Some("k000015"), + None, + Some("j9"), + None, + Some("k000000"), + ], + ) + .with_null_counts([Some(0); 12]) + .with_row_counts([Some(1); 12]), + ); + assert_eq!( + predicate.prune(&stats)?, + [ + true, false, true, false, false, true, true, true, false, false, + true, true, + ], + "count={count}, type={data_type:?}" + ); + assert!(Arc::ptr_eq(predicate.orig_expr(), &expr)); + assert_eq!(predicate.literal_guarantees().len(), 1); + assert_eq!(predicate.literal_guarantees()[0].literals.len(), count); + assert_eq!( + predicate.required_columns().single_column().unwrap().name(), + "c1" + ); + if count > MAX_IN_LIST_SIZE { + // The expression and statistics schema do not grow with the domain. + assert_eq!(predicate.required_columns.columns.len(), 4); + let mut nodes = 0; + predicate.predicate_expr().apply(|_| { + nodes += 1; + Ok(TreeNodeRecursion::Continue) + })?; + assert_eq!(nodes, 7); + } + } + } + Ok(()) + } + + #[test] + fn large_string_in_list_preserves_dictionary_nulls() -> Result<()> { + let data_type = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)); + let schema = + Arc::new(Schema::new(vec![Field::new("c1", data_type.clone(), true)])); + let values = (0..21) + .map(|i| { + let value = ScalarValue::from(format!("k{i:06}")).cast_to(&data_type)?; + Ok(Arc::new(phys_expr::Literal::new(value)) as PhysicalExprRef) + }) + .collect::>>()?; + let expr = phys_expr::in_list( + Arc::new(phys_expr::Column::new("c1", 0)), + values, + &false, + &schema, + )?; + let predicate = large_string_pruning_predicate(expr, schema)?; + // NULL dictionary values can have nonempty payloads. Neither payload + // may become a bound when a valid key references the NULL value. + let (offsets, values, _) = + StringArray::from(vec!["z", "zz", "", "", "k000010", "x"]).into_parts(); + let dictionary_values: ArrayRef = Arc::new(StringArray::new( + offsets, + values, + Some(NullBuffer::from(vec![false, true, false, true, true, true])), + )); + let bounds = |keys: Vec>| -> Result { + Ok(Arc::new(DictionaryArray::::try_new( + Int32Array::from(keys), + Arc::clone(&dictionary_values), + )?)) + }; + let stats = TestStatistics::new().with( + "c1", + ContainerStats::new() + .with_min(bounds(vec![ + Some(0), // NULL with payload "z" + Some(3), // empty string + Some(4), // matching singleton + Some(5), // disjoint singleton + None, // NULL keys + Some(2), // NULL values in both bounds + Some(1), // inverted known bounds + ])?) + .with_max(bounds(vec![ + Some(1), + Some(2), + Some(4), + Some(5), + None, + Some(2), + Some(3), + ])?) + .with_null_counts([Some(0); 7]) + .with_row_counts([Some(1); 7]), + ); + assert_eq!( + predicate.prune(&stats)?, + [true, true, true, false, true, true, true] + ); + Ok(()) + } + + #[test] + fn large_string_in_list_handles_unicode_and_unknown_bounds() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)])); + let mut values = (0..21).map(|i| lit(format!("m{i:03}"))).collect::>(); + values.extend([ + lit(""), + lit("az"), + lit("aé"), + lit("aé"), + lit("é"), + lit("🦀"), + ]); + let expr = col("c1").in_list(values, false); + let predicate = + large_string_pruning_predicate(logical2physical(&expr, &schema), schema)?; + let stats = TestStatistics::new().with( + "c1", + ContainerStats::new_utf8( + [ + Some(""), + Some("az"), + Some("a{"), + Some("aé"), + Some("ê"), + Some("z"), + None, + None, + ], + [ + Some(""), + Some("az"), + Some("aè"), + Some("aé"), + Some("🦀"), + Some("m"), + Some("z"), + None, + ], + ) + .with_null_counts([ + Some(0), + Some(0), + Some(0), + Some(0), + Some(0), + Some(0), + Some(0), + Some(1), + ]) + .with_row_counts([Some(1); 8]), + ); + assert_eq!( + predicate.prune(&stats)?, + [true, true, false, true, true, true, true, false] + ); + Ok(()) + } + + #[test] + fn large_string_in_list_respects_configured_limit() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)])); + let values = (0..21).map(|i| lit(format!("a{i:03}"))).collect::>(); + let expr = logical2physical(&col("c1").in_list(values, false), &schema); + + for (limit, compact) in [(0, false), (20, false), (21, true), (32, true)] { + let predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .with_max_in_list_size(limit) + .try_build(Arc::clone(&expr))?; + assert_eq!( + predicate + .predicate_expr() + .to_string() + .contains("IN_SET_INTERSECTS"), + compact, + "limit={limit}" + ); + assert_eq!(is_always_true(predicate.predicate_expr()), !compact); + } + + let default = PruningPredicateBuilder::new() + .with_file_schema(schema) + .try_build(expr)?; + assert!(is_always_true(default.predicate_expr())); + Ok(()) + } + + #[test] + fn large_string_in_list_keeps_null_and_not_in_semantics() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)])); + let values = (0..21).map(|i| lit(format!("a{i:03}"))).collect::>(); + let stats = TestStatistics::new().with( + "c1", + ContainerStats::new_utf8( + [Some("middle"), Some("other")], + [Some("middle"), Some("other")], + ) + .with_null_counts([Some(0); 2]) + .with_row_counts([Some(1); 2]), + ); + let positive = col("c1").in_list(values.clone(), false); + let predicate = large_string_pruning_predicate( + logical2physical(&positive.or(col("c1").eq(lit("middle"))), &schema), + Arc::clone(&schema), + )?; + assert_eq!(predicate.prune(&stats)?, [true, false]); + + let mut with_null = values.clone(); + with_null.push(lit(ScalarValue::Utf8(None))); + for expr in [ + col("c1").in_list(values, true), + col("c1").in_list(with_null.clone(), false), + col("c1").in_list(with_null.clone(), true), + ] { + let physical = logical2physical(&expr, &schema); + let default = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(Arc::clone(&physical))?; + assert!(is_always_true(default.predicate_expr()), "{expr}"); + assert_eq!(default.prune(&stats)?, [true, true]); + + // Raising the cap retains the existing per-value rewrite for + // NOT IN and lists containing NULL; neither uses the new path. + let raised = large_string_pruning_predicate(physical, Arc::clone(&schema))?; + assert!( + !raised + .predicate_expr() + .to_string() + .contains("IN_SET_INTERSECTS"), + "{expr}" + ); + } + + // Inverting NOT IN (..., NULL) must not prove a full match and bypass + // the original row filter, which returns UNKNOWN for both rows. + let not_in = logical2physical(&col("c1").in_list(with_null, true), &schema); + let batch = RecordBatch::try_new( + schema, + vec![Arc::new(StringArray::from(vec!["middle", "other"]))], + )?; + assert_eq!(not_in.evaluate(&batch)?.into_array(2)?.null_count(), 2); + Ok(()) + } + #[test] fn row_group_predicate_cast_int_int() -> Result<()> { let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]); diff --git a/datafusion/pruning/src/string_in_list.rs b/datafusion/pruning/src/string_in_list.rs new file mode 100644 index 0000000000000..bc5c8fa908e70 --- /dev/null +++ b/datafusion/pruning/src/string_in_list.rs @@ -0,0 +1,239 @@ +// 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 std::fmt::{self, Display, Formatter}; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +use arrow::array::{Array, AsArray, BooleanArray}; +use arrow::compute::cast; +use arrow::datatypes::{DataType, Schema}; +use arrow::record_batch::RecordBatch; +use datafusion_common::{Result, assert_eq_or_internal_err}; +use datafusion_physical_expr::{PhysicalExpr, PhysicalExprRef}; +use datafusion_physical_plan::ColumnarValue; + +/// Tests whether a sorted string domain intersects an inclusive statistics interval. +/// +/// [`PhysicalExpr::evaluate`] returns one nullable Boolean per min/max interval: +/// * `true`: the interval intersects the domain, so matching rows may exist. +/// * `false`: the available bounds prove the interval disjoint from the domain. +/// * `NULL`: incomplete, invalid, or unusable bounds prevent a safe decision. +/// +/// A single known bound can still prove disjointness. Otherwise, unknown results +/// keep the container eligible for reading. +/// +/// This expression is used only for pruning; the original IN remains the row filter. +#[derive(Debug, Eq)] +pub(crate) struct StringInListPruningExpr { + min: PhysicalExprRef, + max: PhysicalExprRef, + values: Arc<[String]>, +} + +impl StringInListPruningExpr { + pub(crate) fn new( + min: PhysicalExprRef, + max: PhysicalExprRef, + mut values: Vec, + ) -> Self { + values.sort_unstable(); + values.dedup(); + Self { + min, + max, + values: values.into(), + } + } +} + +impl PartialEq for StringInListPruningExpr { + fn eq(&self, other: &Self) -> bool { + self.min.eq(&other.min) && self.max.eq(&other.max) && self.values == other.values + } +} + +impl Hash for StringInListPruningExpr { + fn hash(&self, state: &mut H) { + self.min.hash(state); + self.max.hash(state); + self.values.hash(state); + } +} + +impl Display for StringInListPruningExpr { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!( + f, + "IN_SET_INTERSECTS({}, {}, {} values)", + self.min, + self.max, + self.values.len() + ) + } +} + +fn has_oversized_string_buffer(array: &dyn Array, limit: usize) -> bool { + match array.data_type() { + DataType::Utf8 => array.as_string::().values().len() >= limit, + DataType::LargeUtf8 => array.as_string::().values().len() >= limit, + DataType::Dictionary(_, _) => has_oversized_string_buffer( + array.as_any_dictionary().values().as_ref(), + limit, + ), + _ => false, + } +} + +impl PhysicalExpr for StringInListPruningExpr { + fn data_type(&self, _input_schema: &Schema) -> Result { + Ok(DataType::Boolean) + } + + fn nullable(&self, _input_schema: &Schema) -> Result { + Ok(true) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + // Normalize Utf8, LargeUtf8, Utf8View, and dictionary-encoded statistics. + let min = self.min.evaluate(batch)?.into_array(batch.num_rows())?; + let max = self.max.evaluate(batch)?.into_array(batch.num_rows())?; + // A short string slice can retain a buffer too large for Utf8View's + // u32 offsets. Avoid a panic in the cast and keep pruning conservative. + if has_oversized_string_buffer(min.as_ref(), u32::MAX as usize) + || has_oversized_string_buffer(max.as_ref(), u32::MAX as usize) + { + return Ok(ColumnarValue::Array(Arc::new(BooleanArray::new_null( + batch.num_rows(), + )))); + } + // Dictionary values can be NULL behind valid keys. Preserve their + // validity even if the view cast only carries the key nulls. + // TODO: Revisit this workaround once the Arrow dependency includes + // https://github.com/apache/arrow-rs/pull/10510. + let min_nulls = min.logical_nulls(); + let max_nulls = max.logical_nulls(); + let min = cast(&min, &DataType::Utf8View)?; + let max = cast(&max, &DataType::Utf8View)?; + let min = min.as_string_view(); + let max = max.as_string_view(); + let matches: BooleanArray = (0..batch.num_rows()) + .map(|i| { + let min = (min.is_valid(i) + && min_nulls.as_ref().is_none_or(|nulls| nulls.is_valid(i))) + .then(|| min.value(i).as_bytes()); + let max = (max.is_valid(i) + && max_nulls.as_ref().is_none_or(|nulls| nulls.is_valid(i))) + .then(|| max.value(i).as_bytes()); + match (min, max) { + (Some(min), Some(max)) => { + if min > max { + return None; + } + // Rust string ordering and these byte comparisons both use + // unsigned lexicographic UTF-8 order, as required by the + // PruningStatistics min/max contract. Parquet adapters mask + // bounds with unusable ordering; PartitionPruningStatistics + // uses actual Arrow partition values. PrunableStatistics + // trusts file providers' bounds: there is no ordering gate + // for arbitrary statistics providers here. + let index = self.values.partition_point(|v| v.as_bytes() < min); + Some(self.values.get(index).is_some_and(|v| v.as_bytes() <= max)) + } + // A missing bound makes that end of the interval unbounded. + // Exclude only when the whole domain lies beyond the known bound; + // gaps within the domain and equality cannot prove disjointness. + (Some(min), None) + if self.values.last().is_some_and(|v| v.as_bytes() < min) => + { + Some(false) + } + (None, Some(max)) + if self.values.first().is_some_and(|v| v.as_bytes() > max) => + { + Some(false) + } + _ => None, + } + }) + .collect(); + Ok(ColumnarValue::Array(Arc::new(matches))) + } + + fn children(&self) -> Vec<&PhysicalExprRef> { + vec![&self.min, &self.max] + } + + fn with_new_children( + self: Arc, + children: Vec, + ) -> Result { + assert_eq_or_internal_err!(children.len(), 2); + Ok(Arc::new(Self { + min: Arc::clone(&children[0]), + max: Arc::clone(&children[1]), + values: Arc::clone(&self.values), + })) + } + + fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "{self}") + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ArrayRef, StringArray}; + + #[test] + fn oversized_buffers_check_retained_data_not_visible_offsets() -> Result<()> { + // Exercise the size boundary without allocating a 4 GiB buffer. + let limit = 32; + let padding = "p".repeat(limit - 1); + let array: ArrayRef = Arc::new(StringArray::from(vec!["a", padding.as_str()])); + + for value_type in [DataType::Utf8, DataType::LargeUtf8] { + for data_type in [ + value_type.clone(), + DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(value_type.clone()), + ), + DataType::Dictionary( + Box::new(DataType::UInt64), + Box::new(value_type.clone()), + ), + ] { + let slice = cast(&array, &data_type)?.slice(0, 1); + assert!(has_oversized_string_buffer(slice.as_ref(), limit - 1)); + assert!(has_oversized_string_buffer(slice.as_ref(), limit)); + assert!(!has_oversized_string_buffer(slice.as_ref(), limit + 1)); + } + } + + // Already-normalized views do not have the byte-array cast limitation. + for data_type in [ + DataType::Utf8View, + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8View)), + ] { + let slice = cast(&array, &data_type)?.slice(0, 1); + assert!(!has_oversized_string_buffer(slice.as_ref(), limit)); + } + Ok(()) + } +} diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 573fb04b3451b..34ba45962d041 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -413,7 +413,7 @@ datafusion.execution.parquet.dictionary_page_size_limit 1048576 (writing) Sets b datafusion.execution.parquet.enable_page_index true (reading) If true, reads the Parquet data page level metadata (the Page Index), if present, to reduce the I/O and number of rows decoded. datafusion.execution.parquet.encoding NULL (writing) Sets default encoding for any column. Valid values are: plain, plain_dictionary, rle, bit_packed, delta_binary_packed, delta_length_byte_array, delta_byte_array, rle_dictionary, and byte_stream_split. These values are not case sensitive. If NULL, uses default parquet writer setting datafusion.execution.parquet.force_filter_selections false (reading) Force the use of RowSelections for filter results, when pushdown_filters is enabled. If false, the reader will automatically choose between a RowSelection and a Bitmap based on the number and pattern of selected rows. -datafusion.execution.parquet.max_in_list_size 20 Maximum number of values in an `IN (...)` list for which pruning will occur. Longer lists will not be used to prune files, row groups, or data pages. Higher values help in cases such as filtering on a list of ~25-100 identifiers, but also make the predicate more expensive to evaluate. Set to 0 to disable `IN (...)` list pruning entirely. Defaults to 20. +datafusion.execution.parquet.max_in_list_size 20 Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger positive, non-null literal string lists on a string column use a compact sorted domain. Other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. datafusion.execution.parquet.max_predicate_cache_size NULL (reading) The maximum predicate cache size, in bytes. When `pushdown_filters` is enabled, sets the maximum memory used to cache the results of predicate evaluation between filter evaluation and output generation. Decreasing this value will reduce memory usage, but may increase IO and CPU usage. None means use the default parquet reader setting. 0 means no caching. datafusion.execution.parquet.max_row_group_bytes NULL (writing) Target maximum size of each row group in bytes. When set, the writer flushes whenever either this limit or `max_row_group_size` is reached, whichever comes first. Useful for bounding writer memory on wide schemas where a row-count limit can map to very different byte sizes. Matches the behavior of `parquet.block.size` in parquet-mr. If `None` (the default), only the row-count limit applies. Currently only honored when `allow_single_file_parallelism` is `false`; by default the parallel file writer ignores this limit. datafusion.execution.parquet.max_row_group_size 1048576 (writing) Target maximum number of rows in each row group (defaults to 1M rows). Writing larger row groups requires more memory to write, but can get better compression and be faster to read. When `max_row_group_bytes` is also set, the writer flushes a row group when either limit is reached, whichever comes first. diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index c750a3e689951..badfd1ffdba06 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -93,7 +93,7 @@ The following configuration settings are available: | datafusion.execution.parquet.coerce_int96_tz | NULL | (reading) Optional timezone applied to INT96 columns when `coerce_int96` is set. When `Some`, INT96 columns coerce to `Timestamp(, Some())` instead of the default `Timestamp(, None)`. Spark and other systems write INT96 values as UTC-adjusted instants, so callers that need the resulting Arrow type to be timezone-aware (e.g. for Spark `TimestampType` semantics) should set this to `"UTC"`. No effect when `coerce_int96` is `None`. | | datafusion.execution.parquet.bloom_filter_on_read | true | (reading) Use any available bloom filters when reading parquet files | | datafusion.execution.parquet.max_predicate_cache_size | NULL | (reading) The maximum predicate cache size, in bytes. When `pushdown_filters` is enabled, sets the maximum memory used to cache the results of predicate evaluation between filter evaluation and output generation. Decreasing this value will reduce memory usage, but may increase IO and CPU usage. None means use the default parquet reader setting. 0 means no caching. | -| datafusion.execution.parquet.max_in_list_size | 20 | Maximum number of values in an `IN (...)` list for which pruning will occur. Longer lists will not be used to prune files, row groups, or data pages. Higher values help in cases such as filtering on a list of ~25-100 identifiers, but also make the predicate more expensive to evaluate. Set to 0 to disable `IN (...)` list pruning entirely. Defaults to 20. | +| datafusion.execution.parquet.max_in_list_size | 20 | Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger positive, non-null literal string lists on a string column use a compact sorted domain. Other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. | | datafusion.execution.parquet.data_pagesize_limit | 1048576 | (writing) Sets best effort maximum size of data page in bytes | | datafusion.execution.parquet.write_batch_size | 1024 | (writing) Sets write_batch_size in rows | | datafusion.execution.parquet.writer_version | 1.0 | (writing) Sets parquet writer version valid values are "1.0" and "2.0" | From 57f27e90301699422bd24824f36f8b544e41e0b2 Mon Sep 17 00:00:00 2001 From: Kumar Ujjawal Date: Fri, 4 Sep 2026 15:17:27 +0000 Subject: [PATCH 10/19] perf: use compact pruning for large string NOT IN lists (#24781) ## Which issue does this PR close? - Part of #24711. ## Rationale for this change Large string NOT IN lists create long expression chains for Parquet pruning. These expressions are expensive to build and evaluate. ## What changes are included in this PR? - Use a compact sorted domain for large, non-null string NOT IN lists. - Preserve the existing behavior for lists containing NULL. - Add unit tests and Parquet row-group and page-pruning tests. - Add a mixed-container regression test. - Extend the existing benchmark and configuration documentation. - In a local benchmark with 1,024 values, construction was about 17 times faster and evaluation was about 27 times faster. ## Are these changes tested? Yes ## Are there any user-facing changes? No public API changes Backport provenance: - Upstream PR: https://github.com/apache/datafusion/pull/24781 - Upstream commit: eaf7fffbdc7a184cffc81b97e76ab5365546f585 - Upstream self-tests and all additions/removals retained; only unavoidable fork patch context resolved. - Any fork-specific adaptation is kept in a separate following commit. Signed-off-by: discord9 --- datafusion/common/src/config.rs | 8 +- .../tests/parquet/string_in_list_pruning.rs | 210 +++++++++- .../pruning/benches/string_in_list_pruning.rs | 229 +++++++++-- datafusion/pruning/src/pruning_predicate.rs | 381 +++++++++++++++--- datafusion/pruning/src/string_in_list.rs | 92 ++++- .../test_files/information_schema.slt | 2 +- docs/source/user-guide/configs.md | 2 +- 7 files changed, 823 insertions(+), 101 deletions(-) diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 561445cb565e3..8342af7b4a08f 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1287,10 +1287,10 @@ config_namespace! { /// rewrite; other predicates and Bloom-filter pruning remain available. /// /// Within the cap, nonempty lists of at most 20 values use the existing - /// per-value rewrite. Larger positive, non-null literal string lists - /// on a string column use a compact sorted domain. Other lists retain - /// the existing per-value rewrite, so raising the cap can make those - /// predicates expensive to build and evaluate. + /// per-value rewrite. Larger non-null literal string lists on a string + /// column use a compact sorted domain, for both `IN` and `NOT IN`. + /// Other lists retain the existing per-value rewrite, so raising the cap + /// can make those predicates expensive to build and evaluate. /// /// Defaults to 20. pub max_in_list_size: usize, default = 20 diff --git a/datafusion/core/tests/parquet/string_in_list_pruning.rs b/datafusion/core/tests/parquet/string_in_list_pruning.rs index e1005daaa1a38..56a83449a1320 100644 --- a/datafusion/core/tests/parquet/string_in_list_pruning.rs +++ b/datafusion/core/tests/parquet/string_in_list_pruning.rs @@ -15,8 +15,8 @@ // specific language governing permissions and limitations // under the License. -//! End-to-end coverage for compact, large string IN-list pruning. The positive -//! IN-list cases disable the row and Bloom filters to isolate min/max pruning. +//! End-to-end coverage for compact, large string IN-list pruning. The `IN` and +//! `NOT IN` cases disable the row and Bloom filters to isolate min/max pruning. use std::sync::Arc; @@ -47,10 +47,37 @@ const UNITS: usize = 4; const TOTAL_ROWS: usize = ROWS_PER_UNIT * UNITS; const MATCHING_ROWS: usize = ROWS_PER_UNIT * 2; -/// Write either four row groups or four pages in one row group. The second -/// unit lies in a gap between two members of every test IN list; an enclosing -/// min/max range for the list cannot prune it. +/// Write either four row groups or four pages in one row group. Each unit holds +/// a single repeated value, so `NOT IN` can exclude the two units whose value is +/// a list member. The second unit lies in a gap between two members of every +/// test IN list; an enclosing min/max range for the list cannot prune it. fn make_file(page_pruning: bool) -> NamedTempFile { + let values = ["v000000", "v000001", "v000010", "v999999"] + .into_iter() + .flat_map(|value| std::iter::repeat_n(value, ROWS_PER_UNIT)) + .collect::>(); + write_file(page_pruning, values) +} + +/// Write one mixed unit and one single-valued unit. The mixed unit contains +/// both a list member and a value that satisfies `NOT IN`. +fn make_mixed_not_in_file(page_pruning: bool) -> NamedTempFile { + let values = std::iter::repeat_n("v000000", ROWS_PER_UNIT / 2) + .chain(std::iter::repeat_n("v000001", ROWS_PER_UNIT / 2)) + .chain(std::iter::repeat_n("v000010", ROWS_PER_UNIT)) + .collect::>(); + write_file(page_pruning, values) +} + +fn write_file(page_pruning: bool, values: Vec<&str>) -> NamedTempFile { + write_file_with_truncation(page_pruning, values, None) +} + +fn write_file_with_truncation( + page_pruning: bool, + values: Vec<&str>, + truncation_length: Option, +) -> NamedTempFile { let mut file = tempfile::Builder::new() .prefix("string_in_list_pruning") .suffix(".parquet") @@ -61,32 +88,35 @@ fn make_file(page_pruning: bool) -> NamedTempFile { DataType::Utf8, false, )])); - let values = ["v000000", "v000001", "v000010", "v999999"] - .into_iter() - .flat_map(|value| std::iter::repeat_n(value, ROWS_PER_UNIT)) - .collect::>(); + let total_rows = values.len(); + assert_eq!(total_rows % ROWS_PER_UNIT, 0); let batch = RecordBatch::try_new( Arc::clone(&schema), vec![Arc::new(StringArray::from(values))], ) .unwrap(); let rows_per_group = if page_pruning { - TOTAL_ROWS + total_rows } else { ROWS_PER_UNIT }; - let properties = WriterProperties::builder() + let mut properties = WriterProperties::builder() .set_max_row_group_row_count(Some(rows_per_group)) .set_data_page_row_count_limit(ROWS_PER_UNIT) .set_write_batch_size(ROWS_PER_UNIT) .set_dictionary_enabled(false) .set_bloom_filter_enabled(false) - .set_statistics_enabled(EnabledStatistics::Page) - .build(); + .set_statistics_enabled(EnabledStatistics::Page); + if let Some(truncation_length) = truncation_length { + properties = properties + .set_statistics_truncate_length(Some(truncation_length)) + .set_column_index_truncate_length(Some(truncation_length)); + } + let properties = properties.build(); let mut writer = ArrowWriter::try_new(&mut file, schema, Some(properties)).unwrap(); writer.write(&batch).unwrap(); let metadata = writer.close().unwrap(); - assert_eq!(metadata.num_row_groups(), TOTAL_ROWS / rows_per_group); + assert_eq!(metadata.num_row_groups(), total_rows / rows_per_group); let offsets = metadata.offset_index().unwrap(); for row_group in offsets { assert_eq!( @@ -151,6 +181,27 @@ impl ScanOutput { ], &self.batches ); + self.assert_no_filter_interference(); + } + + /// The complement of [`Self::assert_results`]: the two units whose value is + /// not a list member. + fn assert_negated_results(&self) { + assert_batches_eq!( + [ + "+---------+----+", + "| value | n |", + "+---------+----+", + "| v000001 | 16 |", + "| v999999 | 16 |", + "+---------+----+", + ], + &self.batches + ); + self.assert_no_filter_interference(); + } + + fn assert_no_filter_interference(&self) { assert_eq!(self.counter("predicate_evaluation_errors"), 0); assert_eq!(self.counter("pushdown_rows_pruned"), 0); assert_eq!(self.pruned("row_groups_pruned_bloom_filter"), 0); @@ -162,6 +213,7 @@ async fn scan( list_size: usize, max_in_list_size: Option, page_pruning: bool, + negated: bool, ) -> ScanOutput { let mut config = SessionConfig::new() .with_target_partitions(1) @@ -183,9 +235,10 @@ async fn scan( .map(|index| format!("'v{:06}'", index * 10)) .collect::>() .join(", "); + let op = if negated { "NOT IN" } else { "IN" }; let sql = format!( "SELECT value, count(*) AS n FROM t \ - WHERE value IN ({values}) GROUP BY value ORDER BY value" + WHERE value {op} ({values}) GROUP BY value ORDER BY value" ); let plan = ctx .sql(&sql) @@ -209,14 +262,14 @@ async fn check_string_in_list_pruning(page_pruning: bool) { for list_size in [20, 21, 256, 1024] { // A zero cap provides a result-equivalence control that cannot use // min/max IN-list pruning at either granularity. - let unpruned = scan(&file, list_size, Some(0), page_pruning).await; + let unpruned = scan(&file, list_size, Some(0), page_pruning, false).await; unpruned.assert_results(); assert!(!unpruned.plan.contains("IN_SET_INTERSECTS")); assert_eq!(unpruned.pruned("row_groups_pruned_statistics"), 0); assert_eq!(unpruned.pruned("page_index_rows_pruned"), 0); assert_eq!(unpruned.counter("output_rows"), TOTAL_ROWS); - let output = scan(&file, list_size, Some(list_size), page_pruning).await; + let output = scan(&file, list_size, Some(list_size), page_pruning, false).await; output.assert_results(); assert_eq!( pretty_format_batches(&output.batches).unwrap().to_string(), @@ -247,7 +300,7 @@ async fn check_string_in_list_pruning(page_pruning: bool) { // The default remains 20: enabling the compact representation must not // silently change the public cap's meaning. - let default = scan(&file, 21, None, page_pruning).await; + let default = scan(&file, 21, None, page_pruning, false).await; default.assert_results(); assert!(!default.plan.contains("IN_SET_INTERSECTS")); assert_eq!(default.pruned("row_groups_pruned_statistics"), 0); @@ -255,6 +308,111 @@ async fn check_string_in_list_pruning(page_pruning: bool) { assert_eq!(default.counter("output_rows"), TOTAL_ROWS); } +/// The compact `NOT IN` form must prune exactly the units whose single repeated +/// value is a list member, and nothing else. Overlapping a list member is not +/// enough: units 1 and 3 each sit inside the list's enclosing range. +async fn check_string_not_in_list_pruning(page_pruning: bool) { + let file = make_file(page_pruning); + for list_size in [20, 21, 256, 1024] { + let unpruned = scan(&file, list_size, Some(0), page_pruning, true).await; + unpruned.assert_negated_results(); + assert!(!unpruned.plan.contains("NOT_IN_SET_MAY_MATCH")); + assert_eq!(unpruned.pruned("row_groups_pruned_statistics"), 0); + assert_eq!(unpruned.pruned("page_index_rows_pruned"), 0); + assert_eq!(unpruned.counter("output_rows"), TOTAL_ROWS); + + let output = scan(&file, list_size, Some(list_size), page_pruning, true).await; + output.assert_negated_results(); + assert_eq!( + pretty_format_batches(&output.batches).unwrap().to_string(), + pretty_format_batches(&unpruned.batches) + .unwrap() + .to_string() + ); + assert_eq!( + output.plan.contains("NOT_IN_SET_MAY_MATCH"), + list_size > 20, + "list_size={list_size}, plan={}", + output.plan + ); + // The compact form and the per-value AND chain it replaces prune the + // same units, so the counts do not depend on which one ran. + assert_eq!( + output.pruned("row_groups_pruned_statistics"), + if page_pruning { 0 } else { 2 }, + "list_size={list_size}, metrics={}", + output.metrics + ); + assert_eq!( + output.pruned("page_index_rows_pruned"), + if page_pruning { MATCHING_ROWS } else { 0 }, + "list_size={list_size}, metrics={}", + output.metrics + ); + assert_eq!(output.counter("output_rows"), MATCHING_ROWS); + } + + let default = scan(&file, 21, None, page_pruning, true).await; + default.assert_negated_results(); + assert!(!default.plan.contains("NOT_IN_SET_MAY_MATCH")); + assert_eq!(default.pruned("row_groups_pruned_statistics"), 0); + assert_eq!(default.pruned("page_index_rows_pruned"), 0); + assert_eq!(default.counter("output_rows"), TOTAL_ROWS); + + // A mixed interval that overlaps a list member can still contain matching + // rows. Only the adjacent single-valued unit can be excluded. + let mixed_file = make_mixed_not_in_file(page_pruning); + let unpruned = scan(&mixed_file, 21, Some(0), page_pruning, true).await; + assert_batches_eq!( + [ + "+---------+---+", + "| value | n |", + "+---------+---+", + "| v000001 | 8 |", + "+---------+---+", + ], + &unpruned.batches + ); + unpruned.assert_no_filter_interference(); + assert_eq!(unpruned.pruned("row_groups_pruned_statistics"), 0); + assert_eq!(unpruned.pruned("page_index_rows_pruned"), 0); + assert_eq!(unpruned.counter("output_rows"), ROWS_PER_UNIT * 2); + + let output = scan(&mixed_file, 21, Some(21), page_pruning, true).await; + assert_eq!( + pretty_format_batches(&output.batches).unwrap().to_string(), + pretty_format_batches(&unpruned.batches) + .unwrap() + .to_string() + ); + output.assert_no_filter_interference(); + assert!(output.plan.contains("NOT_IN_SET_MAY_MATCH")); + assert_eq!( + output.pruned("row_groups_pruned_statistics"), + usize::from(!page_pruning) + ); + assert_eq!( + output.pruned("page_index_rows_pruned"), + if page_pruning { ROWS_PER_UNIT } else { 0 } + ); + assert_eq!(output.counter("output_rows"), ROWS_PER_UNIT); +} + +async fn check_string_not_in_list_with_truncated_bounds(page_pruning: bool) { + // Exact bounds for this singleton are equal and can exclude the unit. With + // four-byte truncation, Parquet lowers the min and raises the max instead. + let values = std::iter::repeat_n("v000000", ROWS_PER_UNIT).collect(); + let file = write_file_with_truncation(page_pruning, values, Some(4)); + let output = scan(&file, 21, Some(21), page_pruning, true).await; + + assert!(output.batches.iter().all(|batch| batch.num_rows() == 0)); + output.assert_no_filter_interference(); + assert!(output.plan.contains("NOT_IN_SET_MAY_MATCH")); + assert_eq!(output.pruned("row_groups_pruned_statistics"), 0); + assert_eq!(output.pruned("page_index_rows_pruned"), 0); + assert_eq!(output.counter("output_rows"), ROWS_PER_UNIT); +} + #[tokio::test] async fn string_in_list_row_group_pruning() { check_string_in_list_pruning(false).await; @@ -265,6 +423,22 @@ async fn string_in_list_page_pruning() { check_string_in_list_pruning(true).await; } +#[tokio::test] +async fn string_not_in_list_row_group_pruning() { + check_string_not_in_list_pruning(false).await; +} + +#[tokio::test] +async fn string_not_in_list_page_pruning() { + check_string_not_in_list_pruning(true).await; +} + +#[tokio::test] +async fn string_not_in_list_with_truncated_bounds() { + check_string_not_in_list_with_truncated_bounds(false).await; + check_string_not_in_list_with_truncated_bounds(true).await; +} + #[tokio::test] async fn string_not_in_list_with_null_does_not_bypass_row_filter() { let mut file = tempfile::Builder::new() diff --git a/datafusion/pruning/benches/string_in_list_pruning.rs b/datafusion/pruning/benches/string_in_list_pruning.rs index d26112611bf3c..7431d08d25041 100644 --- a/datafusion/pruning/benches/string_in_list_pruning.rs +++ b/datafusion/pruning/benches/string_in_list_pruning.rs @@ -17,12 +17,16 @@ //! Compare compact string IN-list pruning with per-value min/max expansion. //! -//! Both cases raise `max_in_list_size` to the domain size. On a baseline -//! without compact pruning, `in_list` measures the ordinary raised-cap path. -//! The explicit `expanded_or` is a balanced tree of equalities, which produces -//! the same per-value statistics checks without making the baseline depend on -//! a deeply nested expression. Half of the statistics intervals hit a domain -//! member and half fall in a sparse gap. Bloom filters are not involved. +//! Every case raises `max_in_list_size` to the domain size. On a baseline +//! without compact pruning, `in_list` and `not_in_list` measure the ordinary +//! raised-cap path. The explicit `expanded_or` and `expanded_and` are balanced +//! trees of (in)equalities, which produce the same per-value statistics checks +//! without making the baseline depend on a deeply nested expression. +//! +//! The main cases alternate intervals pinned to a domain member with intervals +//! that span a sparse gap. Supplemental `NOT IN` cases measure uniform singleton +//! containers, long bounds, and long literals that share bounds' prefixes. Bloom +//! filters are not involved. //! //! Run with `cargo bench -p datafusion-pruning --bench string_in_list_pruning`. //! The construction benchmarks reuse their input physical expressions; the @@ -48,18 +52,35 @@ fn value(index: usize) -> String { format!("key{index:08}") } -fn balanced_or(expressions: &[PhysicalExprRef]) -> PhysicalExprRef { +fn balanced(expressions: &[PhysicalExprRef], op: Operator) -> PhysicalExprRef { if expressions.len() == 1 { return Arc::clone(&expressions[0]); } let middle = expressions.len() / 2; Arc::new(BinaryExpr::new( - balanced_or(&expressions[..middle]), - Operator::Or, - balanced_or(&expressions[middle..]), + balanced(&expressions[..middle], op), + op, + balanced(&expressions[middle..], op), )) } +/// A balanced tree of `column value`, one branch per domain member. +fn expanded( + column: &PhysicalExprRef, + values: &[PhysicalExprRef], + op: Operator, + combine: Operator, +) -> PhysicalExprRef { + let comparisons = values + .iter() + .map(|value| { + Arc::new(BinaryExpr::new(Arc::clone(column), op, Arc::clone(value))) + as PhysicalExprRef + }) + .collect::>(); + balanced(&comparisons, combine) +} + fn build_predicate( expression: &PhysicalExprRef, schema: &SchemaRef, @@ -96,6 +117,41 @@ impl IntervalStatistics { row_counts: Arc::new(UInt64Array::from(vec![128; CONTAINERS])), } } + + fn uniform_singleton() -> Self { + let value = value(0); + Self::singleton(&value) + } + + fn singleton(value: &str) -> Self { + let min = + StringViewArray::from_iter_values(std::iter::repeat_n(value, CONTAINERS)); + let max = min.clone(); + Self { + min: Arc::new(min), + max: Arc::new(max), + null_counts: Arc::new(UInt64Array::from(vec![0; CONTAINERS])), + row_counts: Arc::new(UInt64Array::from(vec![128; CONTAINERS])), + } + } + + fn long_common_prefix() -> Self { + let prefix = "z".repeat(16 * 1024); + let min = format!("{prefix}a"); + let max = format!("{prefix}z"); + Self { + min: Arc::new(StringViewArray::from_iter_values(std::iter::repeat_n( + min.as_str(), + CONTAINERS, + ))), + max: Arc::new(StringViewArray::from_iter_values(std::iter::repeat_n( + max.as_str(), + CONTAINERS, + ))), + null_counts: Arc::new(UInt64Array::from(vec![0; CONTAINERS])), + row_counts: Arc::new(UInt64Array::from(vec![128; CONTAINERS])), + } + } } impl PruningStatistics for IntervalStatistics { @@ -133,8 +189,12 @@ struct BenchmarkCase { schema: SchemaRef, in_list: PhysicalExprRef, expanded_or: PhysicalExprRef, + not_in_list: PhysicalExprRef, + expanded_and: PhysicalExprRef, in_list_predicate: PruningPredicate, expanded_or_predicate: PruningPredicate, + not_in_list_predicate: PruningPredicate, + expanded_and_predicate: PruningPredicate, statistics: IntervalStatistics, } @@ -149,42 +209,57 @@ impl BenchmarkCase { let values = (0..size) .map(|index| lit(ScalarValue::new_utf8view(value(index * 10)))) .collect::>(); + let not_in_list = + in_list(Arc::clone(&column), values.clone(), &true, &schema).unwrap(); let in_list = in_list(Arc::clone(&column), values.clone(), &false, &schema).unwrap(); - let equalities = values - .into_iter() - .map(|value| { - Arc::new(BinaryExpr::new(Arc::clone(&column), Operator::Eq, value)) - as PhysicalExprRef - }) - .collect::>(); - let expanded_or = balanced_or(&equalities); + let expanded_or = expanded(&column, &values, Operator::Eq, Operator::Or); + let expanded_and = expanded(&column, &values, Operator::NotEq, Operator::And); let in_list_predicate = build_predicate(&in_list, &schema, size); let expanded_or_predicate = build_predicate(&expanded_or, &schema, size); + let not_in_list_predicate = build_predicate(¬_in_list, &schema, size); + let expanded_and_predicate = build_predicate(&expanded_and, &schema, size); eprintln!( - "string_in_list_pruning: {size} values, compact={}", + "string_in_list_pruning: {size} values, compact in={}, compact not in={}", in_list_predicate .predicate_expr() .to_string() - .contains("IN_SET_INTERSECTS") + .contains("IN_SET_INTERSECTS"), + not_in_list_predicate + .predicate_expr() + .to_string() + .contains("NOT_IN_SET_MAY_MATCH") ); let statistics = IntervalStatistics::new(size); - // Check that both benchmark paths do the same useful work, rather + // Check that every benchmark path does the same useful work, rather // than comparing compact pruning with an always-true fallback. - let expected = (0..CONTAINERS) + let kept = (0..CONTAINERS) .map(|index| index % 2 == 0) .collect::>(); - assert_eq!(in_list_predicate.prune(&statistics).unwrap(), expected); - assert_eq!(expanded_or_predicate.prune(&statistics).unwrap(), expected); + let negated_kept = kept.iter().map(|keep| !keep).collect::>(); + assert_eq!(in_list_predicate.prune(&statistics).unwrap(), kept); + assert_eq!(expanded_or_predicate.prune(&statistics).unwrap(), kept); + assert_eq!( + not_in_list_predicate.prune(&statistics).unwrap(), + negated_kept + ); + assert_eq!( + expanded_and_predicate.prune(&statistics).unwrap(), + negated_kept + ); Self { size, schema, in_list, expanded_or, + not_in_list, + expanded_and, in_list_predicate, expanded_or_predicate, + not_in_list_predicate, + expanded_and_predicate, statistics, } } @@ -198,6 +273,8 @@ fn criterion_benchmark(criterion: &mut Criterion) { for (name, expression) in [ ("in_list", &case.in_list), ("expanded_or", &case.expanded_or), + ("not_in_list", &case.not_in_list), + ("expanded_and", &case.expanded_and), ] { construction.bench_with_input( BenchmarkId::new(name, case.size), @@ -222,6 +299,8 @@ fn criterion_benchmark(criterion: &mut Criterion) { for (name, predicate) in [ ("in_list", &case.in_list_predicate), ("expanded_or", &case.expanded_or_predicate), + ("not_in_list", &case.not_in_list_predicate), + ("expanded_and", &case.expanded_and_predicate), ] { evaluation.bench_with_input( BenchmarkId::new(name, case.size), @@ -235,6 +314,108 @@ fn criterion_benchmark(criterion: &mut Criterion) { } } evaluation.finish(); + + let mut distributions = + criterion.benchmark_group("string_in_list_pruning/not_in_distributions"); + distributions.throughput(Throughput::Elements(CONTAINERS as u64)); + let singleton = IntervalStatistics::uniform_singleton(); + for case in cases.iter().filter(|case| case.size > 20) { + let compact = case.not_in_list_predicate.prune(&singleton).unwrap(); + assert!(compact.iter().all(|keep| !keep)); + assert_eq!( + compact, + case.expanded_and_predicate.prune(&singleton).unwrap() + ); + for (name, predicate) in [ + ("uniform_singleton/not_in_list", &case.not_in_list_predicate), + ( + "uniform_singleton/expanded_and", + &case.expanded_and_predicate, + ), + ] { + distributions.bench_with_input( + BenchmarkId::new(name, case.size), + predicate, + |bencher, predicate| { + bencher.iter(|| { + black_box(predicate.prune(black_box(&singleton)).unwrap()) + }); + }, + ); + } + } + + let long_bounds = IntervalStatistics::long_common_prefix(); + let case = cases + .iter() + .find(|case| case.size == 21) + .expect("DOMAIN_SIZES contains 21"); + let compact = case.not_in_list_predicate.prune(&long_bounds).unwrap(); + assert!(compact.iter().all(|keep| *keep)); + assert_eq!( + compact, + case.expanded_and_predicate.prune(&long_bounds).unwrap() + ); + for (name, predicate) in [ + ( + "long_common_prefix/not_in_list", + &case.not_in_list_predicate, + ), + ( + "long_common_prefix/expanded_and", + &case.expanded_and_predicate, + ), + ] { + distributions.bench_with_input( + BenchmarkId::new(name, case.size), + predicate, + |bencher, predicate| { + bencher.iter(|| { + black_box(predicate.prune(black_box(&long_bounds)).unwrap()) + }); + }, + ); + } + + // The bound and literals share a long prefix but have different lengths, + // so exact membership can reject the bound without comparing their bytes. + let prefix = "z".repeat(16 * 1024); + let bound = format!("{prefix}a"); + let long_literal_bounds = IntervalStatistics::singleton(&bound); + let values = (0..case.size) + .map(|index| { + lit(ScalarValue::new_utf8view(format!( + "{prefix}domain{index:08}" + ))) + }) + .collect::>(); + let column = col("value", &case.schema).unwrap(); + let not_in_list = + in_list(Arc::clone(&column), values.clone(), &true, &case.schema).unwrap(); + let expanded_and = expanded(&column, &values, Operator::NotEq, Operator::And); + let not_in_list_predicate = build_predicate(¬_in_list, &case.schema, case.size); + let expanded_and_predicate = build_predicate(&expanded_and, &case.schema, case.size); + let compact = not_in_list_predicate.prune(&long_literal_bounds).unwrap(); + assert!(compact.iter().all(|keep| *keep)); + assert_eq!( + compact, + expanded_and_predicate.prune(&long_literal_bounds).unwrap() + ); + for (name, predicate) in [ + ("long_literal_prefix/not_in_list", ¬_in_list_predicate), + ("long_literal_prefix/expanded_and", &expanded_and_predicate), + ] { + distributions.bench_with_input( + BenchmarkId::new(name, case.size), + predicate, + |bencher, predicate| { + bencher.iter(|| { + black_box(predicate.prune(black_box(&long_literal_bounds)).unwrap()) + }); + }, + ); + } + distributions.finish(); } criterion_group!(benches, criterion_benchmark); diff --git a/datafusion/pruning/src/pruning_predicate.rs b/datafusion/pruning/src/pruning_predicate.rs index f9b19c3959265..cf0631cdb663c 100644 --- a/datafusion/pruning/src/pruning_predicate.rs +++ b/datafusion/pruning/src/pruning_predicate.rs @@ -22,7 +22,7 @@ use std::collections::HashSet; use std::sync::Arc; -use crate::string_in_list::StringInListPruningExpr; +use crate::string_in_list::{SetMembership, StringInListPruningExpr}; use arrow::array::AsArray; use arrow::{ @@ -450,7 +450,7 @@ impl<'a> PruningPredicateBuilder<'a> { /// | Condition | Pruning representation | /// | --- | --- | /// | `N <= min(20, C)` | Existing per-value rewrite | - /// | `20 < N <= C`, positive, non-null literal strings on a string column | Compact sorted domain | + /// | `20 < N <= C`, non-null literal strings on a string column | Compact sorted domain | /// | `20 < N <= C`, other lists | Existing per-value rewrite | /// | `N > C` | Unhandled-predicate hook, normally "keep the container" | /// @@ -459,8 +459,9 @@ impl<'a> PruningPredicateBuilder<'a> { /// pruning (such as Bloom filters). The default cap is [`MAX_IN_LIST_SIZE`] /// (20), so the compact path requires an explicitly raised cap. /// - /// Raising the cap can still build large comparison trees for non-string - /// lists, `NOT IN`, or lists containing NULL; their handling is unchanged. + /// The compact form covers `IN` and `NOT IN` alike. Raising the cap can + /// still build large comparison trees for non-string lists or lists + /// containing NULL; their handling is unchanged. /// /// Query engines typically pass /// `datafusion.execution.parquet.max_in_list_size` here. @@ -1475,15 +1476,27 @@ fn build_is_null_column_expr( } } -/// Keep large literal string domains compact instead of building an OR tree. +/// Keep large literal string domains compact instead of building a per-value +/// tree: an OR tree for `IN`, an AND chain for `NOT IN`. +/// +/// `IN` excludes a container whose interval is disjoint from the domain. That +/// matches the per-value OR tree except for inverted bounds, which the compact +/// form reports as UNKNOWN rather than excluding. +/// +/// `NOT IN` excludes a container only where `min` and `max` agree on a single +/// value the domain holds, which is exactly what makes the per-value +/// `min != v OR v != max` chain false. Its decisions match that chain +/// everywhere, including absent and inverted bounds. fn build_string_in_list_expr( in_list: &phys_expr::InListExpr, schema: &Schema, required_columns: &mut RequiredColumns, ) -> Option> { - if in_list.negated() { - return None; - } + let membership = if in_list.negated() { + SetMembership::NotIn + } else { + SetMembership::In + }; let column = in_list.expr().downcast_ref::()?; let field = schema.fields().get(column.index())?; let data_type = match field.data_type() { @@ -1495,6 +1508,7 @@ fn build_string_in_list_expr( } // NULLs must remain unhandled: the inverse predicate is also used to prove // that every row matches, and IN (..., NULL) can evaluate to UNKNOWN. + // Tracked in https://github.com/apache/datafusion/issues/24711. let values = in_list .list() .iter() @@ -1508,17 +1522,18 @@ fn build_string_in_list_expr( .ok()?; let non_null = build_is_null_column_expr(in_list.expr(), schema, required_columns, true)?; - let intersects = Arc::new(StringInListPruningExpr::new(min, max, values)); + let may_match = Arc::new(StringInListPruningExpr::new(membership, min, max, values)); Some(Arc::new(phys_expr::BinaryExpr::new( non_null, Operator::And, - intersects, + may_match, ))) } /// Default maximum number of entries in an `IN (...)` list eligible for -/// statistics pruning. Eligible positive literal string lists above this -/// threshold use a compact sorted domain instead of per-value min/max checks. +/// statistics pruning. Eligible literal string lists above this threshold use a +/// compact sorted domain instead of per-value min/max checks, for both `IN` and +/// `NOT IN`. /// Callers can raise the cap via [`PredicateRewriter::with_max_in_list_size`], and /// query engines can wire it from the /// `datafusion.execution.parquet.max_in_list_size` config option. @@ -1604,9 +1619,9 @@ impl PredicateRewriter { /// Returns the pruning predicate as an [`PhysicalExpr`] /// /// `max_in_list_size` is the largest `IN (...)` list eligible for statistics -/// pruning. Large positive literal string lists use a compact sorted domain; -/// other eligible lists use per-value checks. Longer lists fall back to -/// `unhandled_hook`. +/// pruning. Large literal string lists use a compact sorted domain, for both +/// `IN` and `NOT IN`; other eligible lists use per-value checks. Longer lists +/// fall back to `unhandled_hook`. fn build_predicate_expression( expr: &Arc, schema: &SchemaRef, @@ -1650,6 +1665,7 @@ fn build_predicate_expression( // Keep the existing expression shape for lists of at most 20 values. // This lower bound is a scope/compatibility choice, not a measured // performance threshold; compact pruning is opt-in via a raised cap. + // The compact form covers both `IN` and `NOT IN`. if in_list.list().len() > MAX_IN_LIST_SIZE && in_list.list().len() <= max_in_list_size && let Some(pruning_expr) = @@ -3889,33 +3905,43 @@ mod tests { fn large_string_in_list_respects_configured_limit() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)])); let values = (0..21).map(|i| lit(format!("a{i:03}"))).collect::>(); - let expr = logical2physical(&col("c1").in_list(values, false), &schema); - for (limit, compact) in [(0, false), (20, false), (21, true), (32, true)] { - let predicate = PruningPredicateBuilder::new() + for (negated, marker) in + [(false, "IN_SET_INTERSECTS"), (true, "NOT_IN_SET_MAY_MATCH")] + { + let expr = + logical2physical(&col("c1").in_list(values.clone(), negated), &schema); + + for (limit, compact) in [(0, false), (20, false), (21, true), (32, true)] { + let predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .with_max_in_list_size(limit) + .try_build(Arc::clone(&expr))?; + assert_eq!( + predicate.predicate_expr().to_string().contains(marker), + compact, + "negated={negated}, limit={limit}" + ); + assert_eq!( + is_always_true(predicate.predicate_expr()), + !compact, + "negated={negated}, limit={limit}" + ); + } + + let default = PruningPredicateBuilder::new() .with_file_schema(Arc::clone(&schema)) - .with_max_in_list_size(limit) - .try_build(Arc::clone(&expr))?; - assert_eq!( - predicate - .predicate_expr() - .to_string() - .contains("IN_SET_INTERSECTS"), - compact, - "limit={limit}" + .try_build(expr)?; + assert!( + is_always_true(default.predicate_expr()), + "negated={negated}" ); - assert_eq!(is_always_true(predicate.predicate_expr()), !compact); } - - let default = PruningPredicateBuilder::new() - .with_file_schema(schema) - .try_build(expr)?; - assert!(is_always_true(default.predicate_expr())); Ok(()) } #[test] - fn large_string_in_list_keeps_null_and_not_in_semantics() -> Result<()> { + fn large_string_in_list_keeps_null_semantics() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)])); let values = (0..21).map(|i| lit(format!("a{i:03}"))).collect::>(); let stats = TestStatistics::new().with( @@ -3934,10 +3960,9 @@ mod tests { )?; assert_eq!(predicate.prune(&stats)?, [true, false]); - let mut with_null = values.clone(); + let mut with_null = values; with_null.push(lit(ScalarValue::Utf8(None))); for expr in [ - col("c1").in_list(values, true), col("c1").in_list(with_null.clone(), false), col("c1").in_list(with_null.clone(), true), ] { @@ -3948,16 +3973,15 @@ mod tests { assert!(is_always_true(default.predicate_expr()), "{expr}"); assert_eq!(default.prune(&stats)?, [true, true]); - // Raising the cap retains the existing per-value rewrite for - // NOT IN and lists containing NULL; neither uses the new path. + // Raising the cap retains the existing per-value rewrite for lists + // containing NULL, in either direction. Dropping the NULL literal + // would turn UNKNOWN into FALSE, which the inverse-predicate proof + // in identify_fully_matched_row_groups cannot absorb. + // See https://github.com/apache/datafusion/issues/24711. let raised = large_string_pruning_predicate(physical, Arc::clone(&schema))?; - assert!( - !raised - .predicate_expr() - .to_string() - .contains("IN_SET_INTERSECTS"), - "{expr}" - ); + let raised = raised.predicate_expr().to_string(); + assert!(!raised.contains("IN_SET_INTERSECTS"), "{expr}"); + assert!(!raised.contains("NOT_IN_SET_MAY_MATCH"), "{expr}"); } // Inverting NOT IN (..., NULL) must not prove a full match and bypass @@ -3971,6 +3995,273 @@ mod tests { Ok(()) } + /// Statistics shared by the compact `NOT IN` tests, one row per case. + fn not_in_container_stats() -> TestStatistics { + TestStatistics::new().with( + "c1", + ContainerStats::new_utf8( + [ + Some("a005"), // single value inside the domain + Some("a005"), // same, with some NULL rows + Some("zzz"), // single value outside the domain + Some("a005"), // two adjacent domain values + Some("a000"), // spans the whole domain + None, // every row NULL + None, // no bounds at all + Some("z"), // inverted bounds + Some("a005"), // unbounded above + None, // unbounded below + ], + [ + Some("a005"), + Some("a005"), + Some("zzz"), + Some("a006"), + Some("a020"), + None, + None, + Some("m"), + None, + Some("a005"), + ], + ) + .with_null_counts([ + Some(0), + Some(3), + Some(0), + Some(0), + Some(0), + Some(10), + Some(0), + Some(0), + Some(0), + Some(0), + ]) + .with_row_counts([Some(10); 10]), + ) + } + + #[test] + fn large_string_not_in_list_prunes_only_single_valued_containers() -> Result<()> { + let types = [ + DataType::Utf8, + DataType::LargeUtf8, + DataType::Utf8View, + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + ]; + for data_type in &types { + let schema = + Arc::new(Schema::new(vec![Field::new("c1", data_type.clone(), true)])); + let values = (0..21) + .map(|i| { + let value = + ScalarValue::from(format!("a{i:03}")).cast_to(data_type)?; + Ok(Arc::new(phys_expr::Literal::new(value)) as PhysicalExprRef) + }) + .collect::>>()?; + let expr = phys_expr::in_list( + Arc::new(phys_expr::Column::new("c1", 0)), + values, + &true, + &schema, + )?; + let predicate = large_string_pruning_predicate(expr, schema)?; + assert_eq!( + predicate.predicate_expr().to_string(), + "c1_null_count@3 != row_count@2 AND NOT_IN_SET_MAY_MATCH(c1_min@0, c1_max@1, 21 values)", + "type={data_type:?}" + ); + + // Overlap never excludes a NOT IN container: only an interval pinned + // to one value the domain holds, or a container that is entirely NULL. + let stats = not_in_container_stats(); + assert_eq!( + predicate.prune(&stats)?, + [ + false, false, true, true, true, false, true, true, true, true + ], + "type={data_type:?}" + ); + + // A one-sided bound in the domain cannot prove whether another + // value satisfies NOT IN, so containers 8 and 9 remain unknown. + let batch = + build_statistics_record_batch(&stats, predicate.required_columns())?; + let result = predicate + .predicate_expr() + .evaluate(&batch)? + .into_array(batch.num_rows())?; + assert!(result.is_null(8), "type={data_type:?}"); + assert!(result.is_null(9), "type={data_type:?}"); + + // The expression and statistics schema do not grow with the domain. + assert_eq!(predicate.required_columns.columns.len(), 4); + let mut nodes = 0; + predicate.predicate_expr().apply(|_| { + nodes += 1; + Ok(TreeNodeRecursion::Continue) + })?; + assert_eq!(nodes, 7, "type={data_type:?}"); + } + Ok(()) + } + + #[test] + fn large_string_not_in_list_preserves_or_short_circuit() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("c1", DataType::Utf8, true), + Field::new("c2", DataType::Int32, true), + ])); + let not_in = + col("c1").in_list((0..21).map(|i| lit(format!("a{i:03}"))).collect(), true); + let other = col("c2").in_list((0..21).map(|i| lit(i * 10)).collect(), false); + let predicate = large_string_pruning_predicate( + logical2physical(¬_in.or(other), &schema), + Arc::clone(&schema), + )?; + let stats = TestStatistics::new() + .with( + "c1", + ContainerStats::new_utf8([Some("zzz"), None], [None, Some("zzz")]) + .with_null_counts([Some(0); 2]) + .with_row_counts([Some(128); 2]), + ) + .with( + "c2", + ContainerStats::new_i32([Some(999); 2], [Some(999); 2]) + .with_null_counts([Some(0); 2]), + ); + + // Both c1 results are true, so the enclosing OR skips its large + // per-value right-hand side for the whole statistics batch. + let batch = build_statistics_record_batch(&stats, predicate.required_columns())?; + let result = predicate + .predicate_expr() + .evaluate(&batch)? + .into_array(batch.num_rows())?; + assert_eq!( + result.as_boolean().iter().collect::>(), + [Some(true), Some(true)] + ); + Ok(()) + } + + /// `identify_fully_matched_row_groups` proves "every row matches" by showing + /// the pruning predicate for `NOT P OR IsNull(col)` excludes the container. + /// That inference needs `P` to be two-valued, so the compact `NOT IN` form + /// and the compact `IN` it inverts to must stay exact. Both sides use a + /// raised cap here, which is what wiring the configured cap into the + /// inverted builder, or removing the lower bound, would produce. + #[test] + fn large_string_not_in_list_inverts_without_false_full_match() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)])); + let values = (0..21).map(|i| lit(format!("a{i:03}"))).collect::>(); + let forward = large_string_pruning_predicate( + logical2physical(&col("c1").in_list(values.clone(), true), &schema), + Arc::clone(&schema), + )?; + // NOT(c1 NOT IN (...)) OR c1 IS NULL, the shape row_group_filter builds + // once PhysicalExprSimplifier turns NOT(NOT IN) back into IN. + let inverted = large_string_pruning_predicate( + logical2physical( + &col("c1").in_list(values, false).or(col("c1").is_null()), + &schema, + ), + schema, + )?; + let inverted_expr = inverted.predicate_expr().to_string(); + assert!( + inverted_expr.contains("IN_SET_INTERSECTS"), + "{inverted_expr}" + ); + + let stats = not_in_container_stats(); + let kept = forward.prune(&stats)?; + let fully_matched = inverted + .prune(&stats)? + .into_iter() + .map(|keep| !keep) + .collect::>(); + + // Only container 2 qualifies: every row holds the single value "zzz", + // which is outside the domain, so every row satisfies NOT IN. Container + // 0 repeats a domain member, and container 5 is entirely NULL, so + // neither may be claimed even though both are single-valued. + assert_eq!( + fully_matched, + [ + false, false, true, false, false, false, false, false, false, false + ] + ); + for (index, (keep, matched)) in kept.iter().zip(&fully_matched).enumerate() { + assert!(*keep || !matched, "container {index} claimed after pruning"); + } + Ok(()) + } + + #[test] + fn large_string_not_in_list_matches_per_value_chain() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)])); + let values = (0..21).map(|i| format!("a{i:03}")).collect::>(); + + let not_in = col("c1").in_list( + values.iter().map(|value| lit(value.clone())).collect(), + true, + ); + let compact = large_string_pruning_predicate( + logical2physical(¬_in, &schema), + Arc::clone(&schema), + )?; + + // The per-value rewrite this replaces, written out as the AND chain of + // `col != value` that build_predicate_expression would produce. + let chain = values + .iter() + .map(|value| col("c1").not_eq(lit(value.clone()))) + .reduce(Expr::and) + .unwrap(); + let per_value = large_string_pruning_predicate( + logical2physical(&chain, &schema), + Arc::clone(&schema), + )?; + let per_value_expr = per_value.predicate_expr().to_string(); + assert!(!per_value_expr.contains("NOT_IN_SET_MAY_MATCH")); + assert!(per_value_expr.contains("c1_min@0 != a000 OR a000 != c1_max@1")); + + // Containers 7 and 9 carry inverted and absent bounds, where the two + // forms could diverge without the compact side matching the chain. + let stats = not_in_container_stats(); + assert_eq!(compact.prune(&stats)?, per_value.prune(&stats)?); + + // The chain above must be what the per-value NOT IN rewrite actually + // builds, or the comparison checks the wrong baseline. Twenty values + // still take that path under the default cap, so both agree there. + let short = (0..20).map(|i| format!("a{i:03}")).collect::>(); + let rewritten = |expr| -> Result { + Ok(PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(logical2physical(expr, &schema))? + .predicate_expr() + .to_string()) + }; + assert_eq!( + rewritten( + &col("c1").in_list( + short.iter().map(|value| lit(value.clone())).collect(), + true + ) + )?, + rewritten( + &short + .iter() + .map(|value| col("c1").not_eq(lit(value.clone()))) + .reduce(Expr::and) + .unwrap() + )? + ); + Ok(()) + } + #[test] fn row_group_predicate_cast_int_int() -> Result<()> { let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]); diff --git a/datafusion/pruning/src/string_in_list.rs b/datafusion/pruning/src/string_in_list.rs index bc5c8fa908e70..b9d9f68b30511 100644 --- a/datafusion/pruning/src/string_in_list.rs +++ b/datafusion/pruning/src/string_in_list.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::cmp::Ordering; use std::fmt::{self, Display, Formatter}; use std::hash::{Hash, Hasher}; use std::sync::Arc; @@ -27,19 +28,48 @@ use datafusion_common::{Result, assert_eq_or_internal_err}; use datafusion_physical_expr::{PhysicalExpr, PhysicalExprRef}; use datafusion_physical_plan::ColumnarValue; -/// Tests whether a sorted string domain intersects an inclusive statistics interval. +/// Which `IN` form a sorted string domain is pruning for. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub(crate) enum SetMembership { + /// `col IN (...)`. A row matches only where the domain intersects the + /// interval, so a disjoint interval excludes every row. + In, + /// `col NOT IN (...)`. Overlap proves nothing here: values outside the + /// domain still satisfy the predicate. An interval excludes every row only + /// when it holds a single value that the domain contains. + NotIn, +} + +impl SetMembership { + fn compare_values(self, left: &[u8], right: &[u8]) -> Ordering { + match self { + // IN uses this order for interval searches. + Self::In => left.cmp(right), + // NOT IN only needs exact membership. Reject impossible lengths + // before comparing bytes that may have a long common prefix. + Self::NotIn => left.len().cmp(&right.len()).then_with(|| left.cmp(right)), + } + } +} + +/// Tests an inclusive statistics interval against a sorted string domain. /// /// [`PhysicalExpr::evaluate`] returns one nullable Boolean per min/max interval: -/// * `true`: the interval intersects the domain, so matching rows may exist. -/// * `false`: the available bounds prove the interval disjoint from the domain. +/// * `true`: matching rows may exist, so the container must be read. +/// * `false`: the available bounds prove no row can match. /// * `NULL`: incomplete, invalid, or unusable bounds prevent a safe decision. /// -/// A single known bound can still prove disjointness. Otherwise, unknown results -/// keep the container eligible for reading. +/// [`SetMembership`] selects the test. For [`SetMembership::In`] a single known +/// bound can still prove disjointness. For [`SetMembership::NotIn`], one known +/// bound outside the domain proves the container may match, while two equal +/// bounds in the domain prove it cannot. This is the same reach as the per-value +/// `min != v OR v != max` chain it replaces. Otherwise, unknown results keep the +/// container eligible for reading. /// /// This expression is used only for pruning; the original IN remains the row filter. #[derive(Debug, Eq)] pub(crate) struct StringInListPruningExpr { + membership: SetMembership, min: PhysicalExprRef, max: PhysicalExprRef, values: Arc<[String]>, @@ -47,28 +77,45 @@ pub(crate) struct StringInListPruningExpr { impl StringInListPruningExpr { pub(crate) fn new( + membership: SetMembership, min: PhysicalExprRef, max: PhysicalExprRef, mut values: Vec, ) -> Self { - values.sort_unstable(); + values.sort_unstable_by(|left, right| { + membership.compare_values(left.as_bytes(), right.as_bytes()) + }); values.dedup(); Self { + membership, min, max, values: values.into(), } } + + /// Does the sorted, deduplicated domain hold `value`? + fn contains(&self, value: &[u8]) -> bool { + self.values + .binary_search_by(|candidate| { + self.membership.compare_values(candidate.as_bytes(), value) + }) + .is_ok() + } } impl PartialEq for StringInListPruningExpr { fn eq(&self, other: &Self) -> bool { - self.min.eq(&other.min) && self.max.eq(&other.max) && self.values == other.values + self.membership == other.membership + && self.min.eq(&other.min) + && self.max.eq(&other.max) + && self.values == other.values } } impl Hash for StringInListPruningExpr { fn hash(&self, state: &mut H) { + self.membership.hash(state); self.min.hash(state); self.max.hash(state); self.values.hash(state); @@ -77,9 +124,13 @@ impl Hash for StringInListPruningExpr { impl Display for StringInListPruningExpr { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + let name = match self.membership { + SetMembership::In => "IN_SET_INTERSECTS", + SetMembership::NotIn => "NOT_IN_SET_MAY_MATCH", + }; write!( f, - "IN_SET_INTERSECTS({}, {}, {} values)", + "{name}({}, {}, {} values)", self.min, self.max, self.values.len() @@ -139,6 +190,30 @@ impl PhysicalExpr for StringInListPruningExpr { let max = (max.is_valid(i) && max_nulls.as_ref().is_none_or(|nulls| nulls.is_valid(i))) .then(|| max.value(i).as_bytes()); + if self.membership == SetMembership::NotIn { + return match (min, max) { + // Check membership before comparing the bounds. NOT IN + // ordering rejects different byte lengths first, avoiding + // scans of long common prefixes. + (Some(min), Some(max)) if self.contains(min) => { + // A wider interval can hold a value outside the domain, + // which satisfies NOT IN. Only an interval pinned to one + // domain value rules out every row. Truncated Parquet + // bounds cannot fake that: min truncates downward and max + // upward, so equal bounds mean the true values were equal. + Some(min != max) + } + (Some(_), Some(_)) => Some(true), + // One known bound outside the domain is enough to preserve + // the true result that lets an enclosing OR short-circuit. + (Some(bound), None) | (None, Some(bound)) + if !self.contains(bound) => + { + Some(true) + } + _ => None, + }; + } match (min, max) { (Some(min), Some(max)) => { if min > max { @@ -184,6 +259,7 @@ impl PhysicalExpr for StringInListPruningExpr { ) -> Result { assert_eq_or_internal_err!(children.len(), 2); Ok(Arc::new(Self { + membership: self.membership, min: Arc::clone(&children[0]), max: Arc::clone(&children[1]), values: Arc::clone(&self.values), diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 34ba45962d041..d99cf69f63339 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -413,7 +413,7 @@ datafusion.execution.parquet.dictionary_page_size_limit 1048576 (writing) Sets b datafusion.execution.parquet.enable_page_index true (reading) If true, reads the Parquet data page level metadata (the Page Index), if present, to reduce the I/O and number of rows decoded. datafusion.execution.parquet.encoding NULL (writing) Sets default encoding for any column. Valid values are: plain, plain_dictionary, rle, bit_packed, delta_binary_packed, delta_length_byte_array, delta_byte_array, rle_dictionary, and byte_stream_split. These values are not case sensitive. If NULL, uses default parquet writer setting datafusion.execution.parquet.force_filter_selections false (reading) Force the use of RowSelections for filter results, when pushdown_filters is enabled. If false, the reader will automatically choose between a RowSelection and a Bitmap based on the number and pattern of selected rows. -datafusion.execution.parquet.max_in_list_size 20 Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger positive, non-null literal string lists on a string column use a compact sorted domain. Other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. +datafusion.execution.parquet.max_in_list_size 20 Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger non-null literal string lists on a string column use a compact sorted domain, for both `IN` and `NOT IN`. Other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. datafusion.execution.parquet.max_predicate_cache_size NULL (reading) The maximum predicate cache size, in bytes. When `pushdown_filters` is enabled, sets the maximum memory used to cache the results of predicate evaluation between filter evaluation and output generation. Decreasing this value will reduce memory usage, but may increase IO and CPU usage. None means use the default parquet reader setting. 0 means no caching. datafusion.execution.parquet.max_row_group_bytes NULL (writing) Target maximum size of each row group in bytes. When set, the writer flushes whenever either this limit or `max_row_group_size` is reached, whichever comes first. Useful for bounding writer memory on wide schemas where a row-count limit can map to very different byte sizes. Matches the behavior of `parquet.block.size` in parquet-mr. If `None` (the default), only the row-count limit applies. Currently only honored when `allow_single_file_parallelism` is `false`; by default the parallel file writer ignores this limit. datafusion.execution.parquet.max_row_group_size 1048576 (writing) Target maximum number of rows in each row group (defaults to 1M rows). Writing larger row groups requires more memory to write, but can get better compression and be faster to read. When `max_row_group_bytes` is also set, the writer flushes a row group when either limit is reached, whichever comes first. diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index badfd1ffdba06..50a580a3fed64 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -93,7 +93,7 @@ The following configuration settings are available: | datafusion.execution.parquet.coerce_int96_tz | NULL | (reading) Optional timezone applied to INT96 columns when `coerce_int96` is set. When `Some`, INT96 columns coerce to `Timestamp(, Some())` instead of the default `Timestamp(, None)`. Spark and other systems write INT96 values as UTC-adjusted instants, so callers that need the resulting Arrow type to be timezone-aware (e.g. for Spark `TimestampType` semantics) should set this to `"UTC"`. No effect when `coerce_int96` is `None`. | | datafusion.execution.parquet.bloom_filter_on_read | true | (reading) Use any available bloom filters when reading parquet files | | datafusion.execution.parquet.max_predicate_cache_size | NULL | (reading) The maximum predicate cache size, in bytes. When `pushdown_filters` is enabled, sets the maximum memory used to cache the results of predicate evaluation between filter evaluation and output generation. Decreasing this value will reduce memory usage, but may increase IO and CPU usage. None means use the default parquet reader setting. 0 means no caching. | -| datafusion.execution.parquet.max_in_list_size | 20 | Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger positive, non-null literal string lists on a string column use a compact sorted domain. Other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. | +| datafusion.execution.parquet.max_in_list_size | 20 | Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger non-null literal string lists on a string column use a compact sorted domain, for both `IN` and `NOT IN`. Other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. | | datafusion.execution.parquet.data_pagesize_limit | 1048576 | (writing) Sets best effort maximum size of data page in bytes | | datafusion.execution.parquet.write_batch_size | 1024 | (writing) Sets write_batch_size in rows | | datafusion.execution.parquet.writer_version | 1.0 | (writing) Sets parquet writer version valid values are "1.0" and "2.0" | From fc503ee2273ec8885b724d9f45dfabb4cc944374 Mon Sep 17 00:00:00 2001 From: Kumar Ujjawal Date: Sat, 5 Sep 2026 11:15:01 +0000 Subject: [PATCH 11/19] Optimize NULL-containing string IN list pruning (#24954) ## Which issue does this PR close? - Closes #24711. ## Rationale for this change #24781 added compact pruning for non-NULL `NOT IN` lists. Issue #24711 still needs support for NULL-containing lists. Large string lists can create expensive per-value pruning expressions. The compact path reduces this cost when `max_in_list_size` permits the rewrite. NULL values require special handling. Filter pruning treats FALSE and UNKNOWN equally, but full-match detection must keep these results separate. ## What changes are included in this PR? This PR makes these changes: - Supports compact pruning for large string `IN` lists that contain NULL. - Treats `NOT IN` with NULL as unable to match rows. - Treats all-NULL `IN` lists as unable to match rows. - Records whether a pruning predicate is safe for full-match inversion. - Uses the same configured list-size cap when it builds the inverse predicate. - Adds unit tests, Parquet integration tests, documentation, and benchmarks. ## What is the testing strategy for this PR? The tests cover positive and negative lists, all-NULL lists, nested `OR` expressions, configured caps, and Parquet full-match detection. These commands passed: ```shell cargo test -p datafusion-pruning cargo test -p datafusion-datasource-parquet cargo test -p datafusion --test parquet_integration string_in_list_pruning cargo fmt --all -- --check cargo clippy --all-targets --all-features -- -D warnings ``` Focused benchmarks also cover compact `IN` and constant-false `NOT IN` expressions with NULL. ## Are there any user-facing changes? Large string `IN` and `NOT IN` lists that contain NULL can now use compact pruning when the cap permits it. Query results and configuration defaults do not change. This PR adds two non-breaking `PruningPredicate` accessors: `max_in_list_size` and `can_be_inverted_for_full_match`. Backport provenance: - Upstream PR: https://github.com/apache/datafusion/pull/24954 - Upstream commit: 33028d5543ed2fd2ad19bd247888cfb3c55cec3b - Upstream self-tests and all additions/removals retained; only unavoidable fork patch context resolved. - Any fork-specific adaptation is kept in a separate following commit. Signed-off-by: discord9 --- datafusion/common/src/config.rs | 10 +- .../tests/parquet/string_in_list_pruning.rs | 132 ++++++++----- .../src/row_group_filter.rs | 6 + .../pruning/benches/string_in_list_pruning.rs | 85 +++++++-- datafusion/pruning/src/pruning_predicate.rs | 178 ++++++++++++++---- .../test_files/information_schema.slt | 2 +- docs/source/user-guide/configs.md | 2 +- 7 files changed, 302 insertions(+), 113 deletions(-) diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 8342af7b4a08f..3e9bdb313870d 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1287,10 +1287,12 @@ config_namespace! { /// rewrite; other predicates and Bloom-filter pruning remain available. /// /// Within the cap, nonempty lists of at most 20 values use the existing - /// per-value rewrite. Larger non-null literal string lists on a string - /// column use a compact sorted domain, for both `IN` and `NOT IN`. - /// Other lists retain the existing per-value rewrite, so raising the cap - /// can make those predicates expensive to build and evaluate. + /// per-value rewrite. Larger literal string lists on a string column use + /// a compact representation, for both `IN` and `NOT IN`, including lists + /// with NULL members. `NOT IN` with NULL and all-NULL `IN` lists cannot + /// match any rows. Other lists retain the existing per-value rewrite, so + /// raising the cap can make those predicates expensive to build and + /// evaluate. /// /// Defaults to 20. pub max_in_list_size: usize, default = 20 diff --git a/datafusion/core/tests/parquet/string_in_list_pruning.rs b/datafusion/core/tests/parquet/string_in_list_pruning.rs index 56a83449a1320..52d2c58bdd943 100644 --- a/datafusion/core/tests/parquet/string_in_list_pruning.rs +++ b/datafusion/core/tests/parquet/string_in_list_pruning.rs @@ -440,23 +440,24 @@ async fn string_not_in_list_with_truncated_bounds() { } #[tokio::test] -async fn string_not_in_list_with_null_does_not_bypass_row_filter() { +async fn string_in_list_with_null_preserves_filter_semantics() { let mut file = tempfile::Builder::new() - .prefix("string_not_in_list_pruning") + .prefix("string_in_list_null_pruning") .suffix(".parquet") .tempfile() .unwrap(); let schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Utf8, true)])); // The first row group has a known zero null count, and every value lies - // in a gap in the IN list. Dropping the NULL list member while inverting - // NOT IN would incorrectly prove that this entire row group matches. + // in a gap in the IN list. The matching value in the second row group is + // deliberately not first, so incorrectly bypassing the row filter changes + // the result when the scan has a limit. let values = vec![ Some("v000001"), Some("v000001"), Some("v000001"), Some("v000001"), - Some("v000000"), Some("v000001"), + Some("v000000"), None, Some("v999999"), ]; @@ -475,13 +476,11 @@ async fn string_not_in_list_with_null_does_not_bypass_row_filter() { assert_eq!(writer.close().unwrap().num_row_groups(), 2); // Build the physical source directly so a logical optimizer cannot fold - // the SQL NOT IN (..., NULL) filter to an empty relation before the scan. + // NOT IN (..., NULL) to an empty relation before the scan. let mut list = (0..21) .map(|index| lit(format!("v{:06}", index * 10))) .collect::>(); list.push(lit(ScalarValue::Utf8(None))); - let predicate = - in_list(col("value", &schema).unwrap(), list, &true, &schema).unwrap(); let location = Path::from_filesystem_path(file.path()).unwrap(); let partitioned_file = PartitionedFile::new( location.to_string(), @@ -490,48 +489,79 @@ async fn string_not_in_list_with_null_does_not_bypass_row_filter() { let ctx = SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); - for max_in_list_size in [0, 32] { - let mut options = TableParquetOptions::default(); - options.global.max_in_list_size = max_in_list_size; - let source = Arc::new( - ParquetSource::new(Arc::clone(&schema)) - .with_table_parquet_options(options) - .with_predicate(Arc::clone(&predicate)) - .with_pushdown_filters(true) - .with_enable_page_index(false) - .with_bloom_filter_on_read(false), - ); - let config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) - .with_file(partitioned_file.clone()) - .with_limit(Some(1)) - .build(); - let plan: Arc = - Arc::new(DataSourceExec::new(Arc::new(config))); - let plan_text = displayable(plan.as_ref()).indent(true).to_string(); - assert!(plan_text.contains("NOT IN"), "{plan_text}"); - let batches = collect(Arc::clone(&plan), ctx.task_ctx()).await.unwrap(); - let output = ScanOutput { - batches, - plan: plan_text, - metrics: MetricsFinder::find_metrics(plan.as_ref()).unwrap(), - }; - - assert_eq!( - output - .batches - .iter() - .map(RecordBatch::num_rows) - .sum::(), - 0, - "cap={max_in_list_size}, plan={}, metrics={}", - output.plan, - output.metrics - ); - assert_eq!(output.fully_matched("row_groups_pruned_statistics"), 0); - assert_eq!(output.pruned("row_groups_pruned_statistics"), 0); - assert_eq!(output.pruned("limit_pruned_row_groups"), 0); - assert_eq!(output.counter("pushdown_rows_pruned"), 8); - assert_eq!(output.counter("predicate_evaluation_errors"), 0); + for negated in [false, true] { + let predicate = in_list( + col("value", &schema).unwrap(), + list.clone(), + &negated, + &schema, + ) + .unwrap(); + for max_in_list_size in [0, 32] { + let mut options = TableParquetOptions::default(); + options.global.max_in_list_size = max_in_list_size; + let source = Arc::new( + ParquetSource::new(Arc::clone(&schema)) + .with_table_parquet_options(options) + .with_predicate(Arc::clone(&predicate)) + .with_pushdown_filters(true) + .with_enable_page_index(false) + .with_bloom_filter_on_read(false), + ); + let config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + .with_file(partitioned_file.clone()) + .with_limit(Some(1)) + .build(); + let plan: Arc = + Arc::new(DataSourceExec::new(Arc::new(config))); + let plan_text = displayable(plan.as_ref()).indent(true).to_string(); + assert!(plan_text.contains("IN"), "{plan_text}"); + let batches = collect(Arc::clone(&plan), ctx.task_ctx()).await.unwrap(); + let output = ScanOutput { + batches, + plan: plan_text, + metrics: MetricsFinder::find_metrics(plan.as_ref()).unwrap(), + }; + + if negated { + assert!(output.batches.iter().all(|batch| batch.num_rows() == 0)); + } else { + assert_batches_eq!( + [ + "+---------+", + "| value |", + "+---------+", + "| v000000 |", + "+---------+", + ], + &output.batches + ); + } + assert_eq!( + output.counter("pushdown_rows_pruned"), + match (negated, max_in_list_size) { + (false, 0) => 7, + (false, _) => 3, + (true, 0) => 8, + (true, _) => 0, + } + ); + assert_eq!(output.fully_matched("row_groups_pruned_statistics"), 0); + assert_eq!( + output.pruned("row_groups_pruned_statistics"), + if max_in_list_size == 0 { + 0 + } else if negated { + 2 + } else { + 1 + }, + "negated={negated}, cap={max_in_list_size}, metrics={}", + output.metrics + ); + assert_eq!(output.pruned("limit_pruned_row_groups"), 0); + assert_eq!(output.counter("predicate_evaluation_errors"), 0); + } } } diff --git a/datafusion/datasource-parquet/src/row_group_filter.rs b/datafusion/datasource-parquet/src/row_group_filter.rs index 80db13edcda81..2b487ec80f944 100644 --- a/datafusion/datasource-parquet/src/row_group_filter.rs +++ b/datafusion/datasource-parquet/src/row_group_filter.rs @@ -388,6 +388,11 @@ impl RowGroupAccessPlanFilter { if candidate_row_group_indices.is_empty() { return; } + // Some pruning rewrites preserve only whether rows can be TRUE, while + // full-match inference must distinguish FALSE from UNKNOWN. + if !predicate.can_be_inverted_for_full_match() { + return; + } let arrow_schema = pruning_stats.arrow_schema; let mut inverted_expr: Arc = @@ -425,6 +430,7 @@ impl RowGroupAccessPlanFilter { let Ok(inverted_predicate) = PruningPredicateBuilder::new() .with_file_schema(Arc::clone(predicate.schema())) + .with_max_in_list_size(predicate.max_in_list_size()) .try_build(inverted_expr) else { return; diff --git a/datafusion/pruning/benches/string_in_list_pruning.rs b/datafusion/pruning/benches/string_in_list_pruning.rs index 7431d08d25041..54d30922d7300 100644 --- a/datafusion/pruning/benches/string_in_list_pruning.rs +++ b/datafusion/pruning/benches/string_in_list_pruning.rs @@ -24,9 +24,10 @@ //! without making the baseline depend on a deeply nested expression. //! //! The main cases alternate intervals pinned to a domain member with intervals -//! that span a sparse gap. Supplemental `NOT IN` cases measure uniform singleton -//! containers, long bounds, and long literals that share bounds' prefixes. Bloom -//! filters are not involved. +//! that span a sparse gap. NULL-containing variants cover the optimized +//! filter-semantics paths, including the constant-false `NOT IN` rewrite. +//! Supplemental `NOT IN` cases measure uniform singleton containers, long bounds, +//! and long literals that share bounds' prefixes. Bloom filters are not involved. //! //! Run with `cargo bench -p datafusion-pruning --bench string_in_list_pruning`. //! The construction benchmarks reuse their input physical expressions; the @@ -188,12 +189,16 @@ struct BenchmarkCase { size: usize, schema: SchemaRef, in_list: PhysicalExprRef, + in_list_with_null: PhysicalExprRef, expanded_or: PhysicalExprRef, not_in_list: PhysicalExprRef, + not_in_list_with_null: PhysicalExprRef, expanded_and: PhysicalExprRef, in_list_predicate: PruningPredicate, + in_list_with_null_predicate: PruningPredicate, expanded_or_predicate: PruningPredicate, not_in_list_predicate: PruningPredicate, + not_in_list_with_null_predicate: PruningPredicate, expanded_and_predicate: PruningPredicate, statistics: IntervalStatistics, } @@ -209,15 +214,30 @@ impl BenchmarkCase { let values = (0..size) .map(|index| lit(ScalarValue::new_utf8view(value(index * 10)))) .collect::>(); + let mut values_with_null = values.clone(); + values_with_null.push(lit(ScalarValue::Utf8View(None))); let not_in_list = in_list(Arc::clone(&column), values.clone(), &true, &schema).unwrap(); + let not_in_list_with_null = in_list( + Arc::clone(&column), + values_with_null.clone(), + &true, + &schema, + ) + .unwrap(); + let in_list_with_null = + in_list(Arc::clone(&column), values_with_null, &false, &schema).unwrap(); let in_list = in_list(Arc::clone(&column), values.clone(), &false, &schema).unwrap(); let expanded_or = expanded(&column, &values, Operator::Eq, Operator::Or); let expanded_and = expanded(&column, &values, Operator::NotEq, Operator::And); let in_list_predicate = build_predicate(&in_list, &schema, size); + let in_list_with_null_predicate = + build_predicate(&in_list_with_null, &schema, size + 1); let expanded_or_predicate = build_predicate(&expanded_or, &schema, size); let not_in_list_predicate = build_predicate(¬_in_list, &schema, size); + let not_in_list_with_null_predicate = + build_predicate(¬_in_list_with_null, &schema, size + 1); let expanded_and_predicate = build_predicate(&expanded_and, &schema, size); eprintln!( "string_in_list_pruning: {size} values, compact in={}, compact not in={}", @@ -239,6 +259,10 @@ impl BenchmarkCase { .collect::>(); let negated_kept = kept.iter().map(|keep| !keep).collect::>(); assert_eq!(in_list_predicate.prune(&statistics).unwrap(), kept); + assert_eq!( + in_list_with_null_predicate.prune(&statistics).unwrap(), + kept + ); assert_eq!(expanded_or_predicate.prune(&statistics).unwrap(), kept); assert_eq!( not_in_list_predicate.prune(&statistics).unwrap(), @@ -248,17 +272,28 @@ impl BenchmarkCase { expanded_and_predicate.prune(&statistics).unwrap(), negated_kept ); + assert!( + not_in_list_with_null_predicate + .prune(&statistics) + .unwrap() + .iter() + .all(|keep| !keep) + ); Self { size, schema, in_list, + in_list_with_null, expanded_or, not_in_list, + not_in_list_with_null, expanded_and, in_list_predicate, + in_list_with_null_predicate, expanded_or_predicate, not_in_list_predicate, + not_in_list_with_null_predicate, expanded_and_predicate, statistics, } @@ -269,22 +304,28 @@ fn criterion_benchmark(criterion: &mut Criterion) { let cases = DOMAIN_SIZES.map(BenchmarkCase::new); let mut construction = criterion.benchmark_group("string_in_list_pruning/construct"); for case in &cases { - construction.throughput(Throughput::Elements(case.size as u64)); - for (name, expression) in [ - ("in_list", &case.in_list), - ("expanded_or", &case.expanded_or), - ("not_in_list", &case.not_in_list), - ("expanded_and", &case.expanded_and), + for (name, expression, max_in_list_size) in [ + ("in_list", &case.in_list, case.size), + ("in_list_with_null", &case.in_list_with_null, case.size + 1), + ("expanded_or", &case.expanded_or, case.size), + ("not_in_list", &case.not_in_list, case.size), + ( + "not_in_list_with_null_constant_false", + &case.not_in_list_with_null, + case.size + 1, + ), + ("expanded_and", &case.expanded_and, case.size), ] { + construction.throughput(Throughput::Elements(max_in_list_size as u64)); construction.bench_with_input( - BenchmarkId::new(name, case.size), + BenchmarkId::new(name, max_in_list_size), expression, |bencher, expression| { bencher.iter(|| { black_box(build_predicate( black_box(expression), &case.schema, - case.size, + max_in_list_size, )) }); }, @@ -296,14 +337,24 @@ fn criterion_benchmark(criterion: &mut Criterion) { let mut evaluation = criterion.benchmark_group("string_in_list_pruning/evaluate"); evaluation.throughput(Throughput::Elements(CONTAINERS as u64)); for case in &cases { - for (name, predicate) in [ - ("in_list", &case.in_list_predicate), - ("expanded_or", &case.expanded_or_predicate), - ("not_in_list", &case.not_in_list_predicate), - ("expanded_and", &case.expanded_and_predicate), + for (name, predicate, list_size) in [ + ("in_list", &case.in_list_predicate, case.size), + ( + "in_list_with_null", + &case.in_list_with_null_predicate, + case.size + 1, + ), + ("expanded_or", &case.expanded_or_predicate, case.size), + ("not_in_list", &case.not_in_list_predicate, case.size), + ( + "not_in_list_with_null_constant_false", + &case.not_in_list_with_null_predicate, + case.size + 1, + ), + ("expanded_and", &case.expanded_and_predicate, case.size), ] { evaluation.bench_with_input( - BenchmarkId::new(name, case.size), + BenchmarkId::new(name, list_size), predicate, |bencher, predicate| { bencher.iter(|| { diff --git a/datafusion/pruning/src/pruning_predicate.rs b/datafusion/pruning/src/pruning_predicate.rs index cf0631cdb663c..614177546b473 100644 --- a/datafusion/pruning/src/pruning_predicate.rs +++ b/datafusion/pruning/src/pruning_predicate.rs @@ -380,6 +380,17 @@ pub struct PruningPredicate { /// /// See [`PruningPredicate::literal_guarantees`] for more details. literal_guarantees: Vec, + /// Maximum IN-list size used to build this predicate. + max_in_list_size: usize, + /// Whether its logical inverse can safely prove every row matches. + can_be_inverted_for_full_match: bool, +} + +#[derive(Default)] +struct PruningExpressionProperties { + /// The rewrite preserves which rows can be TRUE, but not whether rejected + /// rows are FALSE or UNKNOWN. + has_filter_semantics_only: bool, } /// Build a pruning predicate from an optional predicate expression. @@ -450,7 +461,7 @@ impl<'a> PruningPredicateBuilder<'a> { /// | Condition | Pruning representation | /// | --- | --- | /// | `N <= min(20, C)` | Existing per-value rewrite | - /// | `20 < N <= C`, non-null literal strings on a string column | Compact sorted domain | + /// | `20 < N <= C`, literal strings with optional NULLs on a string column | Compact pruning expression | /// | `20 < N <= C`, other lists | Existing per-value rewrite | /// | `N > C` | Unhandled-predicate hook, normally "keep the container" | /// @@ -459,9 +470,9 @@ impl<'a> PruningPredicateBuilder<'a> { /// pruning (such as Bloom filters). The default cap is [`MAX_IN_LIST_SIZE`] /// (20), so the compact path requires an explicitly raised cap. /// - /// The compact form covers `IN` and `NOT IN` alike. Raising the cap can - /// still build large comparison trees for non-string lists or lists - /// containing NULL; their handling is unchanged. + /// The compact form covers `IN` and `NOT IN` alike, including lists with + /// NULL members. Raising the cap can still build large comparison trees for + /// other eligible lists. /// /// Query engines typically pass /// `datafusion.execution.parquet.max_in_list_size` here. @@ -528,12 +539,14 @@ impl<'a> PruningPredicateBuilder<'a> { // build predicate expression once let mut required_columns = RequiredColumns::new(); + let mut properties = PruningExpressionProperties::default(); let predicate_expr = build_predicate_expression( &predicate, &file_schema, &mut required_columns, &unhandled_hook, self.max_in_list_size, + &mut properties, ); let predicate_schema = required_columns.schema(); // Simplify the newly created predicate to get rid of redundant casts, comparisons, etc. @@ -547,6 +560,8 @@ impl<'a> PruningPredicateBuilder<'a> { required_columns, orig_expr: predicate, literal_guarantees, + max_in_list_size: self.max_in_list_size, + can_be_inverted_for_full_match: !properties.has_filter_semantics_only, }) } } @@ -719,6 +734,17 @@ impl PruningPredicate { is_always_true(&self.predicate_expr) && self.literal_guarantees.is_empty() } + /// Returns the configured maximum IN-list size used to build this predicate. + pub fn max_in_list_size(&self) -> usize { + self.max_in_list_size + } + + /// Returns whether pruning the logical inverse can safely prove that every + /// row in a container satisfies the original predicate. + pub fn can_be_inverted_for_full_match(&self) -> bool { + self.can_be_inverted_for_full_match + } + pub fn required_columns(&self) -> &RequiredColumns { &self.required_columns } @@ -1476,7 +1502,7 @@ fn build_is_null_column_expr( } } -/// Keep large literal string domains compact instead of building a per-value +/// Keep large literal string lists compact instead of building a per-value /// tree: an OR tree for `IN`, an AND chain for `NOT IN`. /// /// `IN` excludes a container whose interval is disjoint from the domain. That @@ -1487,10 +1513,14 @@ fn build_is_null_column_expr( /// value the domain holds, which is exactly what makes the per-value /// `min != v OR v != max` chain false. Its decisions match that chain /// everywhere, including absent and inverted bounds. +/// +/// A NULL list member makes `NOT IN` and an all-NULL `IN` list never TRUE. For +/// other `IN` lists, NULL does not change which rows can make the predicate TRUE. fn build_string_in_list_expr( in_list: &phys_expr::InListExpr, schema: &Schema, required_columns: &mut RequiredColumns, + properties: &mut PruningExpressionProperties, ) -> Option> { let membership = if in_list.negated() { SetMembership::NotIn @@ -1506,14 +1536,30 @@ fn build_string_in_list_expr( if field.name() != column.name() || !data_type.is_string() { return None; } - // NULLs must remain unhandled: the inverse predicate is also used to prove - // that every row matches, and IN (..., NULL) can evaluate to UNKNOWN. - // Tracked in https://github.com/apache/datafusion/issues/24711. - let values = in_list - .list() - .iter() - .map(|expr| extract_string_literal(expr).map(str::to_owned)) - .collect::>>()?; + let mut values = Vec::with_capacity(in_list.list().len()); + let mut contains_null = false; + for expr in in_list.list() { + if let Some(value) = extract_string_literal(expr) { + values.push(value.to_owned()); + } else if expr + .downcast_ref::() + .is_some_and(|literal| literal.value().is_null()) + { + contains_null = true; + } else { + return None; + } + } + + // Pruning asks only whether the predicate can be TRUE. UNKNOWN and FALSE + // both reject a row, so NOT IN with a NULL member and an all-NULL IN list + // can never match. IN can otherwise ignore NULL and search its non-null domain. + if contains_null && (membership == SetMembership::NotIn || values.is_empty()) { + properties.has_filter_semantics_only = true; + return Some(Arc::new(phys_expr::Literal::new(ScalarValue::Boolean( + Some(false), + )))); + } let min = required_columns .min_column_expr(column, in_list.expr(), field) .ok()?; @@ -1523,6 +1569,9 @@ fn build_string_in_list_expr( let non_null = build_is_null_column_expr(in_list.expr(), schema, required_columns, true)?; let may_match = Arc::new(StringInListPruningExpr::new(membership, min, max, values)); + if contains_null { + properties.has_filter_semantics_only = true; + } Some(Arc::new(phys_expr::BinaryExpr::new( non_null, Operator::And, @@ -1600,12 +1649,14 @@ impl PredicateRewriter { schema: &Schema, ) -> Arc { let mut required_columns = RequiredColumns::new(); + let mut properties = PruningExpressionProperties::default(); build_predicate_expression( expr, &Arc::new(schema.clone()), &mut required_columns, &self.unhandled_hook, self.max_in_list_size, + &mut properties, ) } } @@ -1619,7 +1670,7 @@ impl PredicateRewriter { /// Returns the pruning predicate as an [`PhysicalExpr`] /// /// `max_in_list_size` is the largest `IN (...)` list eligible for statistics -/// pruning. Large literal string lists use a compact sorted domain, for both +/// pruning. Large literal string lists use a compact representation, for both /// `IN` and `NOT IN`; other eligible lists use per-value checks. Longer lists /// fall back to `unhandled_hook`. fn build_predicate_expression( @@ -1628,6 +1679,7 @@ fn build_predicate_expression( required_columns: &mut RequiredColumns, unhandled_hook: &Arc, max_in_list_size: usize, + properties: &mut PruningExpressionProperties, ) -> Arc { if is_always_false(expr) { // Shouldn't return `unhandled_hook.handle(expr)` @@ -1669,7 +1721,7 @@ fn build_predicate_expression( if in_list.list().len() > MAX_IN_LIST_SIZE && in_list.list().len() <= max_in_list_size && let Some(pruning_expr) = - build_string_in_list_expr(in_list, schema, required_columns) + build_string_in_list_expr(in_list, schema, required_columns, properties) { return pruning_expr; } @@ -1702,6 +1754,7 @@ fn build_predicate_expression( required_columns, unhandled_hook, max_in_list_size, + properties, ); } else { return unhandled_hook.handle(expr); @@ -1742,6 +1795,7 @@ fn build_predicate_expression( required_columns, unhandled_hook, max_in_list_size, + properties, ); let right_expr = build_predicate_expression( &right, @@ -1749,6 +1803,7 @@ fn build_predicate_expression( required_columns, unhandled_hook, max_in_list_size, + properties, ); // simplify boolean expression if applicable let expr = match (&left_expr, op, &right_expr) { @@ -3927,6 +3982,8 @@ mod tests { !compact, "negated={negated}, limit={limit}" ); + assert!(predicate.can_be_inverted_for_full_match()); + assert_eq!(predicate.max_in_list_size(), limit); } let default = PruningPredicateBuilder::new() @@ -3941,10 +3998,23 @@ mod tests { } #[test] - fn large_string_in_list_keeps_null_semantics() -> Result<()> { + fn large_string_in_list_compacts_null_literals() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)])); let values = (0..21).map(|i| lit(format!("a{i:03}"))).collect::>(); let stats = TestStatistics::new().with( + "c1", + ContainerStats::new_utf8( + [Some("a005"), Some("other")], + [Some("a005"), Some("other")], + ) + .with_null_counts([Some(0); 2]) + .with_row_counts([Some(1); 2]), + ); + + let mut with_null = values; + with_null.push(lit(ScalarValue::Utf8(None))); + + let or_stats = TestStatistics::new().with( "c1", ContainerStats::new_utf8( [Some("middle"), Some("other")], @@ -3953,39 +4023,68 @@ mod tests { .with_null_counts([Some(0); 2]) .with_row_counts([Some(1); 2]), ); - let positive = col("c1").in_list(values.clone(), false); + let positive_or = col("c1") + .in_list(with_null.clone(), false) + .or(col("c1").eq(lit("middle"))); let predicate = large_string_pruning_predicate( - logical2physical(&positive.or(col("c1").eq(lit("middle"))), &schema), + logical2physical(&positive_or, &schema), Arc::clone(&schema), )?; - assert_eq!(predicate.prune(&stats)?, [true, false]); + assert_eq!(predicate.prune(&or_stats)?, [true, false]); + assert!(!predicate.can_be_inverted_for_full_match()); - let mut with_null = values; - with_null.push(lit(ScalarValue::Utf8(None))); - for expr in [ - col("c1").in_list(with_null.clone(), false), - col("c1").in_list(with_null.clone(), true), - ] { + for negated in [false, true] { + let expr = col("c1").in_list(with_null.clone(), negated); let physical = logical2physical(&expr, &schema); let default = PruningPredicateBuilder::new() .with_file_schema(Arc::clone(&schema)) .try_build(Arc::clone(&physical))?; assert!(is_always_true(default.predicate_expr()), "{expr}"); assert_eq!(default.prune(&stats)?, [true, true]); + assert!(default.can_be_inverted_for_full_match()); - // Raising the cap retains the existing per-value rewrite for lists - // containing NULL, in either direction. Dropping the NULL literal - // would turn UNKNOWN into FALSE, which the inverse-predicate proof - // in identify_fully_matched_row_groups cannot absorb. - // See https://github.com/apache/datafusion/issues/24711. let raised = large_string_pruning_predicate(physical, Arc::clone(&schema))?; - let raised = raised.predicate_expr().to_string(); - assert!(!raised.contains("IN_SET_INTERSECTS"), "{expr}"); - assert!(!raised.contains("NOT_IN_SET_MAY_MATCH"), "{expr}"); + assert!(!raised.can_be_inverted_for_full_match()); + if negated { + assert!(is_always_false(raised.predicate_expr()), "{expr}"); + assert_eq!(raised.prune(&stats)?, [false, false]); + } else { + assert!( + raised + .predicate_expr() + .to_string() + .contains("IN_SET_INTERSECTS"), + "{expr}" + ); + assert_eq!(raised.prune(&stats)?, [true, false]); + } } - // Inverting NOT IN (..., NULL) must not prove a full match and bypass - // the original row filter, which returns UNKNOWN for both rows. + let all_null = + std::iter::repeat_n(lit(ScalarValue::Utf8(None)), 21).collect::>(); + for negated in [false, true] { + let expr = col("c1").in_list(all_null.clone(), negated); + let predicate = large_string_pruning_predicate( + logical2physical(&expr, &schema), + Arc::clone(&schema), + )?; + assert!(is_always_false(predicate.predicate_expr()), "{expr}"); + assert_eq!(predicate.prune(&stats)?, [false, false]); + assert!(!predicate.can_be_inverted_for_full_match()); + } + + let integer_schema = + Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, true)])); + let mut integer_values = (0..21).map(lit).collect::>(); + integer_values.push(lit(ScalarValue::Int32(None))); + let integer_predicate = large_string_pruning_predicate( + logical2physical(&col("c1").in_list(integer_values, false), &integer_schema), + integer_schema, + )?; + assert!(integer_predicate.can_be_inverted_for_full_match()); + + // The compact false predicate has the same filter result as the original + // NOT IN expression, which returns UNKNOWN for values outside the list. let not_in = logical2physical(&col("c1").in_list(with_null, true), &schema); let batch = RecordBatch::try_new( schema, @@ -4148,10 +4247,9 @@ mod tests { /// `identify_fully_matched_row_groups` proves "every row matches" by showing /// the pruning predicate for `NOT P OR IsNull(col)` excludes the container. - /// That inference needs `P` to be two-valued, so the compact `NOT IN` form - /// and the compact `IN` it inverts to must stay exact. Both sides use a - /// raised cap here, which is what wiring the configured cap into the - /// inverted builder, or removing the lower bound, would produce. + /// For non-NULL lists, the compact `NOT IN` form and the compact `IN` it + /// inverts to are precise enough for this proof. Both sides use a raised cap + /// here, matching the configuration used by row-group pruning. #[test] fn large_string_not_in_list_inverts_without_false_full_match() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)])); @@ -6697,12 +6795,14 @@ mod tests { ) -> Arc { let expr = logical2physical(expr, schema); let unhandled_hook = Arc::new(ConstantUnhandledPredicateHook::default()) as _; + let mut properties = PruningExpressionProperties::default(); build_predicate_expression( &expr, &Arc::new(schema.clone()), required_columns, &unhandled_hook, MAX_IN_LIST_SIZE, + &mut properties, ) } diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index d99cf69f63339..1acdb2058d122 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -413,7 +413,7 @@ datafusion.execution.parquet.dictionary_page_size_limit 1048576 (writing) Sets b datafusion.execution.parquet.enable_page_index true (reading) If true, reads the Parquet data page level metadata (the Page Index), if present, to reduce the I/O and number of rows decoded. datafusion.execution.parquet.encoding NULL (writing) Sets default encoding for any column. Valid values are: plain, plain_dictionary, rle, bit_packed, delta_binary_packed, delta_length_byte_array, delta_byte_array, rle_dictionary, and byte_stream_split. These values are not case sensitive. If NULL, uses default parquet writer setting datafusion.execution.parquet.force_filter_selections false (reading) Force the use of RowSelections for filter results, when pushdown_filters is enabled. If false, the reader will automatically choose between a RowSelection and a Bitmap based on the number and pattern of selected rows. -datafusion.execution.parquet.max_in_list_size 20 Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger non-null literal string lists on a string column use a compact sorted domain, for both `IN` and `NOT IN`. Other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. +datafusion.execution.parquet.max_in_list_size 20 Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger literal string lists on a string column use a compact representation, for both `IN` and `NOT IN`, including lists with NULL members. `NOT IN` with NULL and all-NULL `IN` lists cannot match any rows. Other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. datafusion.execution.parquet.max_predicate_cache_size NULL (reading) The maximum predicate cache size, in bytes. When `pushdown_filters` is enabled, sets the maximum memory used to cache the results of predicate evaluation between filter evaluation and output generation. Decreasing this value will reduce memory usage, but may increase IO and CPU usage. None means use the default parquet reader setting. 0 means no caching. datafusion.execution.parquet.max_row_group_bytes NULL (writing) Target maximum size of each row group in bytes. When set, the writer flushes whenever either this limit or `max_row_group_size` is reached, whichever comes first. Useful for bounding writer memory on wide schemas where a row-count limit can map to very different byte sizes. Matches the behavior of `parquet.block.size` in parquet-mr. If `None` (the default), only the row-count limit applies. Currently only honored when `allow_single_file_parallelism` is `false`; by default the parallel file writer ignores this limit. datafusion.execution.parquet.max_row_group_size 1048576 (writing) Target maximum number of rows in each row group (defaults to 1M rows). Writing larger row groups requires more memory to write, but can get better compression and be faster to read. When `max_row_group_bytes` is also set, the writer flushes a row group when either limit is reached, whichever comes first. diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 50a580a3fed64..bffc69172a201 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -93,7 +93,7 @@ The following configuration settings are available: | datafusion.execution.parquet.coerce_int96_tz | NULL | (reading) Optional timezone applied to INT96 columns when `coerce_int96` is set. When `Some`, INT96 columns coerce to `Timestamp(, Some())` instead of the default `Timestamp(, None)`. Spark and other systems write INT96 values as UTC-adjusted instants, so callers that need the resulting Arrow type to be timezone-aware (e.g. for Spark `TimestampType` semantics) should set this to `"UTC"`. No effect when `coerce_int96` is `None`. | | datafusion.execution.parquet.bloom_filter_on_read | true | (reading) Use any available bloom filters when reading parquet files | | datafusion.execution.parquet.max_predicate_cache_size | NULL | (reading) The maximum predicate cache size, in bytes. When `pushdown_filters` is enabled, sets the maximum memory used to cache the results of predicate evaluation between filter evaluation and output generation. Decreasing this value will reduce memory usage, but may increase IO and CPU usage. None means use the default parquet reader setting. 0 means no caching. | -| datafusion.execution.parquet.max_in_list_size | 20 | Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger non-null literal string lists on a string column use a compact sorted domain, for both `IN` and `NOT IN`. Other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. | +| datafusion.execution.parquet.max_in_list_size | 20 | Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger literal string lists on a string column use a compact representation, for both `IN` and `NOT IN`, including lists with NULL members. `NOT IN` with NULL and all-NULL `IN` lists cannot match any rows. Other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. | | datafusion.execution.parquet.data_pagesize_limit | 1048576 | (writing) Sets best effort maximum size of data page in bytes | | datafusion.execution.parquet.write_batch_size | 1024 | (writing) Sets write_batch_size in rows | | datafusion.execution.parquet.writer_version | 1.0 | (writing) Sets parquet writer version valid values are "1.0" and "2.0" | From a23cfb1f8acfb9f152f665fc1c8a4046a809a3fd Mon Sep 17 00:00:00 2001 From: Kumar Ujjawal Date: Wed, 9 Sep 2026 03:32:35 +0000 Subject: [PATCH 12/19] perf: extend compact IN-list pruning to ordered types (#25012) ## Which issue does this PR close? - Closes #24709 ## Rationale for this change Compact pruning for large `IN` lists only supports string columns. Other ordered types use one min/max comparison for each literal. Large lists create large expression trees and make pruning more expensive. These types can use the same compact sorted-domain approach. ## What changes are included in this PR? This PR adds compact `IN` and `NOT IN` pruning for: - Variable-length binary types - Signed and unsigned integers - Decimal types - Date and time types - Timestamp and duration types - Dictionary-encoded forms of these types String and binary values share the byte-domain implementation. Primitive types use typed contiguous arrays. The compact path preserves NULL behavior and dictionary value NULLs. Unsupported types keep the existing per-value fallback. Floating-point types remain unsupported because of NaN and signed zero. `FixedSizeBinary` and nested types also remain unsupported. The default `max_in_list_size` value remains 20. ### Benchmarks With 1,024 literals and 4,096 statistics containers: - String `IN` evaluation is 37.8x faster. - String `NOT IN` evaluation is 21.4x faster. - Int64 `IN` evaluation is 118x faster. - Int64 `NOT IN` evaluation is 111x faster. - Predicate construction is 15.9x to 32.1x faster. Commands: ```shell cargo bench -p datafusion-pruning --bench string_in_list_pruning cargo bench -p datafusion-pruning --bench primitive_in_list_pruning ``` ## What is the testing strategy for this PR? The tests cover: - All supported type families - Both `IN` and `NOT IN` - NULL-containing and all-NULL lists - Dictionary and run-end-encoded values - Missing, inverted, and truncated statistics - Parquet row-group and page pruning - Per-value fallback for unsupported types - Result equality between compact pruning and the original filter The following commands passed: ```shell cargo test -p datafusion-pruning cargo test -p datafusion --test parquet_integration string_in_list_pruning cargo test -p datafusion-sqllogictest --test sqllogictests -- information_schema cargo clippy -p datafusion-pruning -p datafusion --all-targets --all-features -- -D warnings cargo fmt --all -- --check ``` ## Are there any user-facing changes? Large `IN` and `NOT IN` lists can use compact pruning for more ordered types. This behavior applies when `datafusion.execution.parquet.max_in_list_size` is set above 20. The default value and SQL results do not change. This PR does not change public APIs. Backport provenance: - Upstream PR: https://github.com/apache/datafusion/pull/25012 - Upstream commit: f3f420b4c28ad044a08a7695466ef831060572eb - Upstream self-tests and all additions/removals retained; only unavoidable fork patch context resolved. - Any fork-specific adaptation is kept in a separate following commit. Signed-off-by: discord9 --- datafusion/common/src/config.rs | 14 +- .../tests/parquet/string_in_list_pruning.rs | 369 ++++++++-- datafusion/pruning/Cargo.toml | 4 + .../benches/primitive_in_list_pruning.rs | 285 ++++++++ datafusion/pruning/src/in_list.rs | 85 +++ datafusion/pruning/src/lib.rs | 2 + datafusion/pruning/src/primitive_in_list.rs | 430 ++++++++++++ datafusion/pruning/src/pruning_predicate.rs | 651 ++++++++++++++++-- datafusion/pruning/src/string_in_list.rs | 526 +++++++------- .../test_files/information_schema.slt | 2 +- docs/source/user-guide/configs.md | 2 +- 11 files changed, 1994 insertions(+), 376 deletions(-) create mode 100644 datafusion/pruning/benches/primitive_in_list_pruning.rs create mode 100644 datafusion/pruning/src/in_list.rs create mode 100644 datafusion/pruning/src/primitive_in_list.rs diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 3e9bdb313870d..9f596ceefb7a7 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1287,12 +1287,14 @@ config_namespace! { /// rewrite; other predicates and Bloom-filter pruning remain available. /// /// Within the cap, nonempty lists of at most 20 values use the existing - /// per-value rewrite. Larger literal string lists on a string column use - /// a compact representation, for both `IN` and `NOT IN`, including lists - /// with NULL members. `NOT IN` with NULL and all-NULL `IN` lists cannot - /// match any rows. Other lists retain the existing per-value rewrite, so - /// raising the cap can make those predicates expensive to build and - /// evaluate. + /// per-value rewrite. Larger literal lists use a compact representation + /// when the column type is string, variable-length binary, integer, + /// decimal, date, time, timestamp, or duration. This applies to both `IN` + /// and `NOT IN`, including lists with NULL members. `NOT IN` with NULL and + /// all-NULL `IN` lists cannot match any rows. Compact lists containing NULL + /// do not use the fully-matched-row-group optimization. Floating-point and + /// other lists retain the existing per-value rewrite, so raising the cap + /// can make those predicates expensive to build and evaluate. /// /// Defaults to 20. pub max_in_list_size: usize, default = 20 diff --git a/datafusion/core/tests/parquet/string_in_list_pruning.rs b/datafusion/core/tests/parquet/string_in_list_pruning.rs index 52d2c58bdd943..f2646855713ce 100644 --- a/datafusion/core/tests/parquet/string_in_list_pruning.rs +++ b/datafusion/core/tests/parquet/string_in_list_pruning.rs @@ -15,12 +15,16 @@ // specific language governing permissions and limitations // under the License. -//! End-to-end coverage for compact, large string IN-list pruning. The `IN` and -//! `NOT IN` cases disable the row and Bloom filters to isolate min/max pruning. +//! End-to-end coverage for compact, large IN-list pruning. The string cases +//! exercise detailed edge conditions; representative ordered types verify the +//! same path against real Parquet row-group and page statistics. use std::sync::Arc; -use arrow::array::StringArray; +use arrow::array::{ + ArrayRef, BinaryArray, Date64Array, Decimal128Array, Int64Array, StringArray, + TimestampMicrosecondArray, +}; use arrow::datatypes::{DataType, Field, Schema}; use arrow::record_batch::RecordBatch; use arrow::util::pretty::pretty_format_batches; @@ -34,9 +38,11 @@ use datafusion_common::config::TableParquetOptions; use datafusion_common::{ScalarValue, assert_batches_eq}; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; use datafusion_physical_expr::expressions::{col, in_list, lit}; +use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::metrics::{MetricValue, MetricsSet}; use object_store::path::Path; use parquet::arrow::ArrowWriter; +use parquet::file::metadata::ParquetMetaData; use parquet::file::properties::{EnabledStatistics, WriterProperties}; use tempfile::NamedTempFile; @@ -47,6 +53,21 @@ const UNITS: usize = 4; const TOTAL_ROWS: usize = ROWS_PER_UNIT * UNITS; const MATCHING_ROWS: usize = ROWS_PER_UNIT * 2; +fn assert_file_layout( + metadata: &ParquetMetaData, + total_rows: usize, + rows_per_group: usize, +) { + assert_eq!(metadata.num_row_groups(), total_rows / rows_per_group); + let offsets = metadata.offset_index().unwrap(); + for row_group in offsets { + assert_eq!( + row_group[0].page_locations().len(), + rows_per_group / ROWS_PER_UNIT + ); + } +} + /// Write either four row groups or four pages in one row group. Each unit holds /// a single repeated value, so `NOT IN` can exclude the two units whose value is /// a list member. The second unit lies in a gap between two members of every @@ -116,14 +137,7 @@ fn write_file_with_truncation( let mut writer = ArrowWriter::try_new(&mut file, schema, Some(properties)).unwrap(); writer.write(&batch).unwrap(); let metadata = writer.close().unwrap(); - assert_eq!(metadata.num_row_groups(), total_rows / rows_per_group); - let offsets = metadata.offset_index().unwrap(); - for row_group in offsets { - assert_eq!( - row_group[0].page_locations().len(), - rows_per_group / ROWS_PER_UNIT - ); - } + assert_file_layout(&metadata, total_rows, rows_per_group); file } @@ -257,6 +271,93 @@ async fn scan( } } +fn repeated_ordered_values(scale: i64) -> Vec { + [0, 1, 2, 100] + .into_iter() + .flat_map(|value| std::iter::repeat_n(value * scale, ROWS_PER_UNIT)) + .collect() +} + +fn write_ordered_file( + schema: &Arc, + values: ArrayRef, + page_pruning: bool, +) -> NamedTempFile { + let mut file = tempfile::Builder::new() + .prefix("ordered_in_list_pruning") + .suffix(".parquet") + .tempfile() + .unwrap(); + let rows_per_group = if page_pruning { + TOTAL_ROWS + } else { + ROWS_PER_UNIT + }; + let properties = WriterProperties::builder() + .set_max_row_group_row_count(Some(rows_per_group)) + .set_data_page_row_count_limit(ROWS_PER_UNIT) + .set_write_batch_size(ROWS_PER_UNIT) + .set_dictionary_enabled(false) + .set_bloom_filter_enabled(false) + .set_statistics_enabled(EnabledStatistics::Page) + .build(); + let batch = RecordBatch::try_new(Arc::clone(schema), vec![values]).unwrap(); + let mut writer = + ArrowWriter::try_new(&mut file, Arc::clone(schema), Some(properties)).unwrap(); + writer.write(&batch).unwrap(); + let metadata = writer.close().unwrap(); + assert_file_layout(&metadata, TOTAL_ROWS, rows_per_group); + file +} + +async fn scan_ordered( + file: &NamedTempFile, + schema: &Arc, + list: &[ScalarValue], + max_in_list_size: usize, + page_pruning: bool, + negated: bool, +) -> ScanOutput { + let predicate = in_list( + col("value", schema).unwrap(), + list.iter().cloned().map(lit).collect(), + &negated, + schema, + ) + .unwrap(); + let mut options = TableParquetOptions::default(); + options.global.max_in_list_size = max_in_list_size; + let source = Arc::new( + ParquetSource::new(Arc::clone(schema)) + .with_table_parquet_options(options) + .with_predicate(Arc::clone(&predicate)) + .with_pushdown_filters(false) + .with_enable_page_index(page_pruning) + .with_bloom_filter_on_read(false), + ); + let location = Path::from_filesystem_path(file.path()).unwrap(); + let partitioned_file = PartitionedFile::new( + location.to_string(), + file.as_file().metadata().unwrap().len(), + ); + let config = FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + .with_file(partitioned_file) + .build(); + let scan: Arc = Arc::new(DataSourceExec::new(Arc::new(config))); + let plan: Arc = + Arc::new(FilterExec::try_new(predicate, scan).unwrap()); + let plan_text = displayable(plan.as_ref()).indent(true).to_string(); + let ctx = + SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); + let batches = collect(Arc::clone(&plan), ctx.task_ctx()).await.unwrap(); + let metrics = MetricsFinder::find_metrics(plan.as_ref()).unwrap(); + ScanOutput { + batches, + plan: plan_text, + metrics, + } +} + async fn check_string_in_list_pruning(page_pruning: bool) { let file = make_file(page_pruning); for list_size in [20, 21, 256, 1024] { @@ -413,6 +514,118 @@ async fn check_string_not_in_list_with_truncated_bounds(page_pruning: bool) { assert_eq!(output.counter("output_rows"), ROWS_PER_UNIT); } +#[tokio::test] +async fn ordered_in_list_parquet_pruning() { + let decimal_values = Decimal128Array::from_iter_values( + repeated_ordered_values(1).into_iter().map(i128::from), + ) + .with_precision_and_scale(18, 2) + .unwrap(); + let cases = [ + ( + "binary", + Arc::new(BinaryArray::from_iter_values( + repeated_ordered_values(1) + .into_iter() + .map(|value| [0xff, value as u8]), + )) as ArrayRef, + (0..21) + .map(|value| ScalarValue::Binary(Some(vec![0xff, value * 2]))) + .collect::>(), + ), + ( + "int64", + Arc::new(Int64Array::from(repeated_ordered_values(1))) as ArrayRef, + (0..21) + .map(|value| ScalarValue::Int64(Some(value * 2))) + .collect(), + ), + ( + "decimal128", + Arc::new(decimal_values) as ArrayRef, + (0..21) + .map(|value| ScalarValue::Decimal128(Some(value * 2), 18, 2)) + .collect(), + ), + ( + "date64", + Arc::new(Date64Array::from(repeated_ordered_values(86_400_000))) as ArrayRef, + (0..21) + .map(|value| ScalarValue::Date64(Some(value * 2 * 86_400_000))) + .collect(), + ), + ( + "timestamp_microsecond_utc", + Arc::new( + TimestampMicrosecondArray::from(repeated_ordered_values(1_000_000)) + .with_timezone("UTC"), + ) as ArrayRef, + (0..21) + .map(|value| { + ScalarValue::TimestampMicrosecond( + Some(value * 2 * 1_000_000), + Some("UTC".into()), + ) + }) + .collect(), + ), + ]; + + for (name, values, list) in cases { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + values.data_type().clone(), + false, + )])); + for page_pruning in [false, true] { + let file = write_ordered_file(&schema, Arc::clone(&values), page_pruning); + for negated in [false, true] { + let control = + scan_ordered(&file, &schema, &list, 0, page_pruning, negated).await; + let output = + scan_ordered(&file, &schema, &list, 32, page_pruning, negated).await; + control.assert_no_filter_interference(); + output.assert_no_filter_interference(); + assert!(!control.plan.contains("IN_SET_INTERSECTS")); + assert!(!control.plan.contains("NOT_IN_SET_MAY_MATCH")); + assert_eq!( + pretty_format_batches(&output.batches).unwrap().to_string(), + pretty_format_batches(&control.batches).unwrap().to_string(), + "type={name}, page_pruning={page_pruning}, negated={negated}" + ); + assert_eq!( + output + .batches + .iter() + .map(RecordBatch::num_rows) + .sum::(), + MATCHING_ROWS, + "type={name}, page_pruning={page_pruning}, negated={negated}" + ); + assert!( + output.plan.contains(if negated { + "NOT_IN_SET_MAY_MATCH" + } else { + "IN_SET_INTERSECTS" + }), + "type={name}, page_pruning={page_pruning}, negated={negated}, plan={}", + output.plan + ); + assert_eq!( + output.pruned("row_groups_pruned_statistics"), + if page_pruning { 0 } else { 2 }, + "type={name}, page_pruning={page_pruning}, negated={negated}" + ); + assert_eq!( + output.pruned("page_index_rows_pruned"), + if page_pruning { MATCHING_ROWS } else { 0 }, + "type={name}, page_pruning={page_pruning}, negated={negated}" + ); + } + } + } +} + #[tokio::test] async fn string_in_list_row_group_pruning() { check_string_in_list_pruning(false).await; @@ -439,33 +652,23 @@ async fn string_not_in_list_with_truncated_bounds() { check_string_not_in_list_with_truncated_bounds(true).await; } -#[tokio::test] -async fn string_in_list_with_null_preserves_filter_semantics() { +async fn check_in_list_with_null_preserves_filter_semantics( + name: &str, + values: ArrayRef, + list: Vec, + expected: ScalarValue, +) { let mut file = tempfile::Builder::new() - .prefix("string_in_list_null_pruning") + .prefix("in_list_null_pruning") .suffix(".parquet") .tempfile() .unwrap(); - let schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Utf8, true)])); - // The first row group has a known zero null count, and every value lies - // in a gap in the IN list. The matching value in the second row group is - // deliberately not first, so incorrectly bypassing the row filter changes - // the result when the scan has a limit. - let values = vec![ - Some("v000001"), - Some("v000001"), - Some("v000001"), - Some("v000001"), - Some("v000001"), - Some("v000000"), - None, - Some("v999999"), - ]; - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(StringArray::from(values))], - ) - .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + values.data_type().clone(), + true, + )])); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![values]).unwrap(); let properties = WriterProperties::builder() .set_max_row_group_row_count(Some(4)) .set_bloom_filter_enabled(false) @@ -477,10 +680,7 @@ async fn string_in_list_with_null_preserves_filter_semantics() { // Build the physical source directly so a logical optimizer cannot fold // NOT IN (..., NULL) to an empty relation before the scan. - let mut list = (0..21) - .map(|index| lit(format!("v{:06}", index * 10))) - .collect::>(); - list.push(lit(ScalarValue::Utf8(None))); + let list = list.into_iter().map(lit).collect::>(); let location = Path::from_filesystem_path(file.path()).unwrap(); let partitioned_file = PartitionedFile::new( location.to_string(), @@ -516,7 +716,12 @@ async fn string_in_list_with_null_preserves_filter_semantics() { let plan: Arc = Arc::new(DataSourceExec::new(Arc::new(config))); let plan_text = displayable(plan.as_ref()).indent(true).to_string(); - assert!(plan_text.contains("IN"), "{plan_text}"); + assert!(plan_text.contains("IN"), "type={name}, plan={plan_text}"); + assert_eq!( + plan_text.contains("IN_SET_INTERSECTS"), + !negated && max_in_list_size == 32, + "type={name}, negated={negated}, cap={max_in_list_size}, plan={plan_text}" + ); let batches = collect(Arc::clone(&plan), ctx.task_ctx()).await.unwrap(); let output = ScanOutput { batches, @@ -525,17 +730,29 @@ async fn string_in_list_with_null_preserves_filter_semantics() { }; if negated { - assert!(output.batches.iter().all(|batch| batch.num_rows() == 0)); + assert!( + output.batches.iter().all(|batch| batch.num_rows() == 0), + "type={name}, cap={max_in_list_size}" + ); } else { - assert_batches_eq!( - [ - "+---------+", - "| value |", - "+---------+", - "| v000000 |", - "+---------+", - ], - &output.batches + assert_eq!( + output + .batches + .iter() + .map(RecordBatch::num_rows) + .sum::(), + 1, + "type={name}, cap={max_in_list_size}" + ); + let batch = output + .batches + .iter() + .find(|batch| batch.num_rows() > 0) + .unwrap(); + assert_eq!( + ScalarValue::try_from_array(batch.column(0), 0).unwrap(), + expected, + "type={name}, cap={max_in_list_size}" ); } assert_eq!( @@ -545,7 +762,8 @@ async fn string_in_list_with_null_preserves_filter_semantics() { (false, _) => 3, (true, 0) => 8, (true, _) => 0, - } + }, + "type={name}, negated={negated}, cap={max_in_list_size}" ); assert_eq!(output.fully_matched("row_groups_pruned_statistics"), 0); assert_eq!( @@ -557,7 +775,7 @@ async fn string_in_list_with_null_preserves_filter_semantics() { } else { 1 }, - "negated={negated}, cap={max_in_list_size}, metrics={}", + "type={name}, negated={negated}, cap={max_in_list_size}, metrics={}", output.metrics ); assert_eq!(output.pruned("limit_pruned_row_groups"), 0); @@ -565,3 +783,52 @@ async fn string_in_list_with_null_preserves_filter_semantics() { } } } + +#[tokio::test] +async fn in_list_with_null_preserves_filter_semantics() { + // The first row group has a known zero null count, and every value lies + // in a gap in the IN list. The matching value in the second row group is + // deliberately not first, so incorrectly bypassing the row filter changes + // the result when the scan has a limit. + let mut string_list = (0..21) + .map(|index| ScalarValue::Utf8(Some(format!("v{:06}", index * 10)))) + .collect::>(); + string_list.push(ScalarValue::Utf8(None)); + check_in_list_with_null_preserves_filter_semantics( + "string", + Arc::new(StringArray::from(vec![ + Some("v000001"), + Some("v000001"), + Some("v000001"), + Some("v000001"), + Some("v000001"), + Some("v000000"), + None, + Some("v999999"), + ])) as ArrayRef, + string_list, + ScalarValue::Utf8(Some("v000000".into())), + ) + .await; + + let mut int64_list = (0..21) + .map(|index| ScalarValue::Int64(Some(index * 10))) + .collect::>(); + int64_list.push(ScalarValue::Int64(None)); + check_in_list_with_null_preserves_filter_semantics( + "int64", + Arc::new(Int64Array::from(vec![ + Some(1), + Some(1), + Some(1), + Some(1), + Some(1), + Some(0), + None, + Some(999_999), + ])) as ArrayRef, + int64_list, + ScalarValue::Int64(Some(0)), + ) + .await; +} diff --git a/datafusion/pruning/Cargo.toml b/datafusion/pruning/Cargo.toml index a703f68222f38..a914a0a079bcd 100644 --- a/datafusion/pruning/Cargo.toml +++ b/datafusion/pruning/Cargo.toml @@ -35,3 +35,7 @@ itertools = { workspace = true } [[bench]] harness = false name = "string_in_list_pruning" + +[[bench]] +harness = false +name = "primitive_in_list_pruning" diff --git a/datafusion/pruning/benches/primitive_in_list_pruning.rs b/datafusion/pruning/benches/primitive_in_list_pruning.rs new file mode 100644 index 0000000000000..830cf4d0b6875 --- /dev/null +++ b/datafusion/pruning/benches/primitive_in_list_pruning.rs @@ -0,0 +1,285 @@ +// 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. + +//! Compare primitive IN-list pruning with per-value min/max expansion. +//! +//! The domain and container matrices match `string_in_list_pruning`. The +//! compact form uses a typed contiguous domain, while `expanded_or` and +//! `expanded_and` measure balanced per-value comparison trees. +//! +//! Run with `cargo bench -p datafusion-pruning --bench primitive_in_list_pruning`. + +use std::collections::HashSet; +use std::hint::black_box; +use std::sync::Arc; + +use arrow::array::{ArrayRef, BooleanArray, Int64Array, UInt64Array}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use datafusion_common::{Column, ScalarValue}; +use datafusion_expr_common::operator::Operator; +use datafusion_physical_expr::PhysicalExprRef; +use datafusion_physical_expr::expressions::{ + BinaryExpr, col, in_list as make_in_list, lit, +}; +use datafusion_pruning::{PruningPredicate, PruningPredicateBuilder, PruningStatistics}; + +const DOMAIN_SIZES: [usize; 9] = [1, 2, 4, 8, 16, 20, 21, 256, 1024]; +const CONTAINER_COUNTS: [usize; 3] = [16, 256, 4096]; + +fn sampled_domain_index( + pair_index: usize, + pair_count: usize, + domain_size: usize, +) -> usize { + let sampled_positions = pair_count.min(domain_size); + let position = pair_index % sampled_positions; + if sampled_positions == 1 { + 0 + } else { + position * (domain_size - 1) / (sampled_positions - 1) + } +} + +fn balanced(expressions: &[PhysicalExprRef], op: Operator) -> PhysicalExprRef { + if expressions.len() == 1 { + return Arc::clone(&expressions[0]); + } + let middle = expressions.len() / 2; + Arc::new(BinaryExpr::new( + balanced(&expressions[..middle], op), + op, + balanced(&expressions[middle..], op), + )) +} + +fn expanded( + column: &PhysicalExprRef, + values: &[PhysicalExprRef], + op: Operator, + combine: Operator, +) -> PhysicalExprRef { + let comparisons = values + .iter() + .map(|value| { + Arc::new(BinaryExpr::new(Arc::clone(column), op, Arc::clone(value))) + as PhysicalExprRef + }) + .collect::>(); + balanced(&comparisons, combine) +} + +fn build_predicate( + expression: &PhysicalExprRef, + schema: &SchemaRef, + max_in_list_size: usize, +) -> PruningPredicate { + PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(schema)) + .with_max_in_list_size(max_in_list_size) + .try_build(Arc::clone(expression)) + .unwrap() +} + +struct IntervalStatistics { + min: ArrayRef, + max: ArrayRef, + null_counts: ArrayRef, + row_counts: ArrayRef, +} + +impl IntervalStatistics { + fn new(domain_size: usize, container_count: usize) -> Self { + let pair_count = container_count.div_ceil(2); + let min = Int64Array::from_iter_values((0..container_count).map(|index| { + let start = + sampled_domain_index(index / 2, pair_count, domain_size) as i64 * 10; + start + if index % 2 == 0 { 0 } else { 3 } + })); + let max = Int64Array::from_iter_values((0..container_count).map(|index| { + let start = + sampled_domain_index(index / 2, pair_count, domain_size) as i64 * 10; + start + if index % 2 == 0 { 0 } else { 7 } + })); + Self { + min: Arc::new(min), + max: Arc::new(max), + null_counts: Arc::new(UInt64Array::from(vec![0; container_count])), + row_counts: Arc::new(UInt64Array::from(vec![128; container_count])), + } + } +} + +impl PruningStatistics for IntervalStatistics { + fn min_values(&self, column: &Column) -> Option { + (column.name == "value").then(|| Arc::clone(&self.min)) + } + + fn max_values(&self, column: &Column) -> Option { + (column.name == "value").then(|| Arc::clone(&self.max)) + } + + fn num_containers(&self) -> usize { + self.min.len() + } + + fn null_counts(&self, column: &Column) -> Option { + (column.name == "value").then(|| Arc::clone(&self.null_counts)) + } + + fn row_counts(&self) -> Option { + Some(Arc::clone(&self.row_counts)) + } + + fn contained( + &self, + _column: &Column, + _values: &HashSet, + ) -> Option { + None + } +} + +struct BenchmarkCase { + size: usize, + schema: SchemaRef, + in_list: PhysicalExprRef, + expanded_or: PhysicalExprRef, + not_in_list: PhysicalExprRef, + expanded_and: PhysicalExprRef, + in_list_predicate: PruningPredicate, + expanded_or_predicate: PruningPredicate, + not_in_list_predicate: PruningPredicate, + expanded_and_predicate: PruningPredicate, +} + +impl BenchmarkCase { + fn new(size: usize) -> Self { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int64, + false, + )])); + let column = col("value", &schema).unwrap(); + let values = (0..size) + .map(|index| lit(ScalarValue::Int64(Some(index as i64 * 10)))) + .collect::>(); + let in_list = + make_in_list(Arc::clone(&column), values.clone(), &false, &schema).unwrap(); + let not_in_list = + make_in_list(Arc::clone(&column), values.clone(), &true, &schema).unwrap(); + let expanded_or = expanded(&column, &values, Operator::Eq, Operator::Or); + let expanded_and = expanded(&column, &values, Operator::NotEq, Operator::And); + let in_list_predicate = build_predicate(&in_list, &schema, size); + let expanded_or_predicate = build_predicate(&expanded_or, &schema, size); + let not_in_list_predicate = build_predicate(¬_in_list, &schema, size); + let expanded_and_predicate = build_predicate(&expanded_and, &schema, size); + Self { + size, + schema, + in_list, + expanded_or, + not_in_list, + expanded_and, + in_list_predicate, + expanded_or_predicate, + not_in_list_predicate, + expanded_and_predicate, + } + } +} + +fn assert_equivalent_results(case: &BenchmarkCase, statistics: &IntervalStatistics) { + let expected = (0..statistics.num_containers()) + .map(|index| index % 2 == 0) + .collect::>(); + assert_eq!(case.in_list_predicate.prune(statistics).unwrap(), expected); + assert_eq!( + case.expanded_or_predicate.prune(statistics).unwrap(), + expected + ); + let negated = expected.into_iter().map(|value| !value).collect::>(); + assert_eq!( + case.not_in_list_predicate.prune(statistics).unwrap(), + negated + ); + assert_eq!( + case.expanded_and_predicate.prune(statistics).unwrap(), + negated + ); +} + +fn criterion_benchmark(criterion: &mut Criterion) { + let cases = DOMAIN_SIZES.map(BenchmarkCase::new); + let mut construction = + criterion.benchmark_group("primitive_in_list_pruning/construct"); + for case in &cases { + for (name, expression) in [ + ("in_list", &case.in_list), + ("expanded_or", &case.expanded_or), + ("not_in_list", &case.not_in_list), + ("expanded_and", &case.expanded_and), + ] { + construction.throughput(Throughput::Elements(case.size as u64)); + construction.bench_with_input( + BenchmarkId::new(name, case.size), + expression, + |bencher, expression| { + bencher.iter(|| { + black_box(build_predicate( + black_box(expression), + &case.schema, + case.size, + )) + }); + }, + ); + } + } + construction.finish(); + + for container_count in CONTAINER_COUNTS { + let mut evaluation = criterion.benchmark_group(format!( + "primitive_in_list_pruning/evaluate/{container_count}_containers" + )); + evaluation.throughput(Throughput::Elements(container_count as u64)); + for case in &cases { + let statistics = IntervalStatistics::new(case.size, container_count); + assert_equivalent_results(case, &statistics); + for (name, predicate) in [ + ("in_list", &case.in_list_predicate), + ("expanded_or", &case.expanded_or_predicate), + ("not_in_list", &case.not_in_list_predicate), + ("expanded_and", &case.expanded_and_predicate), + ] { + evaluation.bench_with_input( + BenchmarkId::new(name, case.size), + predicate, + |bencher, predicate| { + bencher.iter(|| { + black_box(predicate.prune(black_box(&statistics)).unwrap()) + }); + }, + ); + } + } + evaluation.finish(); + } +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/pruning/src/in_list.rs b/datafusion/pruning/src/in_list.rs new file mode 100644 index 0000000000000..7f4b6b5aaf93b --- /dev/null +++ b/datafusion/pruning/src/in_list.rs @@ -0,0 +1,85 @@ +// 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 std::cmp::Ordering; + +use datafusion_common::ScalarValue; + +/// Which `IN` form a sorted domain is pruning for. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub(crate) enum SetMembership { + /// `col IN (...)`. A row matches only where the domain intersects the + /// interval, so a disjoint interval excludes every row. + In, + /// `col NOT IN (...)`. Overlap proves nothing here: values outside the + /// domain still satisfy the predicate. An interval excludes every row only + /// when it holds a single value that the domain contains. + NotIn, +} + +impl SetMembership { + pub(crate) fn display_name(self) -> &'static str { + match self { + Self::In => "IN_SET_INTERSECTS", + Self::NotIn => "NOT_IN_SET_MAY_MATCH", + } + } + + pub(crate) fn compare_bytes(self, left: &[u8], right: &[u8]) -> Ordering { + match self { + // IN uses this order for interval searches. + Self::In => left.cmp(right), + // NOT IN only needs exact membership. Reject impossible lengths + // before comparing bytes that may have a long common prefix. + Self::NotIn => left.len().cmp(&right.len()).then_with(|| left.cmp(right)), + } + } +} + +/// Evaluates `NOT IN` against an inclusive statistics interval. +/// +/// A container can be excluded only when both bounds identify one value in the +/// domain. A known bound outside the domain proves the container may match and +/// lets an enclosing Boolean expression short-circuit. +/// +/// This remains safe when Parquet truncates byte-array statistics. A truncated +/// minimum is no greater than the true minimum, and a truncated maximum is no +/// less than the true maximum. Therefore, equal stored bounds prove that the +/// true minimum and maximum are also equal. +#[inline(always)] +pub(crate) fn not_in_may_match( + min: Option<&T>, + max: Option<&T>, + contains: impl Fn(&T) -> bool, +) -> Option { + match (min, max) { + (Some(min), Some(max)) if contains(min) => Some(min != max), + (Some(_), Some(_)) => Some(true), + (Some(bound), None) | (None, Some(bound)) if !contains(bound) => Some(true), + _ => None, + } +} + +/// Removes scalar wrappers that do not change the represented value. +pub(crate) fn unwrap_scalar(value: &ScalarValue) -> &ScalarValue { + match value { + ScalarValue::Dictionary(_, value) | ScalarValue::RunEndEncoded(_, _, value) => { + unwrap_scalar(value) + } + value => value, + } +} diff --git a/datafusion/pruning/src/lib.rs b/datafusion/pruning/src/lib.rs index 6bf1815900aa8..88a7fdda5e733 100644 --- a/datafusion/pruning/src/lib.rs +++ b/datafusion/pruning/src/lib.rs @@ -18,6 +18,8 @@ #![cfg_attr(test, allow(clippy::needless_pass_by_value))] mod file_pruner; +mod in_list; +mod primitive_in_list; mod pruning_predicate; mod string_in_list; diff --git a/datafusion/pruning/src/primitive_in_list.rs b/datafusion/pruning/src/primitive_in_list.rs new file mode 100644 index 0000000000000..6b21a437862a6 --- /dev/null +++ b/datafusion/pruning/src/primitive_in_list.rs @@ -0,0 +1,430 @@ +// 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 std::fmt::{self, Display, Formatter}; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +use arrow::array::{Array, AsArray, BooleanArray}; +use arrow::compute::cast; +use arrow::datatypes::{ + ArrowPrimitiveType, DataType, Date32Type, Date64Type, Decimal32Type, Decimal64Type, + Decimal128Type, Decimal256Type, DurationMicrosecondType, DurationMillisecondType, + DurationNanosecondType, DurationSecondType, Int8Type, Int16Type, Int32Type, + Int64Type, Schema, Time32MillisecondType, Time32SecondType, Time64MicrosecondType, + Time64NanosecondType, TimeUnit, TimestampMicrosecondType, TimestampMillisecondType, + TimestampNanosecondType, TimestampSecondType, UInt8Type, UInt16Type, UInt32Type, + UInt64Type, +}; +use arrow::record_batch::RecordBatch; +use datafusion_common::{Result, ScalarValue, assert_eq_or_internal_err}; +use datafusion_physical_expr::{PhysicalExpr, PhysicalExprRef}; +use datafusion_physical_plan::ColumnarValue; + +use crate::in_list::{SetMembership, not_in_may_match}; + +/// A typed, non-null primitive IN-list domain under construction. +pub(crate) struct PrimitiveInListDomain { + data_type: DataType, + values: PrimitiveValues, +} + +macro_rules! define_primitive_values { + ($( + $variant:ident, $arrow_type:ty, $data_pattern:pat, + $scalar_pattern:pat, $native:expr $(, if $guard:expr)?; + )+) => { + enum PrimitiveValues { + $( + $variant(Vec<<$arrow_type as ArrowPrimitiveType>::Native>), + )+ + } + + impl PrimitiveInListDomain { + pub(crate) fn new( + data_type: &DataType, + capacity: usize, + ) -> Option { + // Parameter bindings in these patterns are used by `push` guards, + // but construction only selects the matching storage type. + #[expect( + unused_variables, + reason = "parameter bindings are reused by the generated push guards" + )] + let values = match data_type { + $( + $data_pattern => PrimitiveValues::$variant( + Vec::with_capacity(capacity), + ), + )+ + _ => return None, + }; + Some(Self { + data_type: data_type.clone(), + values, + }) + } + + /// Adds one non-null scalar whose logical type matches the domain. + pub(crate) fn push(&mut self, value: &ScalarValue) -> Option<()> { + match (&self.data_type, &mut self.values, value) { + $( + ($data_pattern, PrimitiveValues::$variant(values), $scalar_pattern) + $(if $guard)? => + { + values.push($native); + Some(()) + } + )+ + _ => None, + } + } + + pub(crate) fn is_empty(&self) -> bool { + match &self.values { + $(PrimitiveValues::$variant(values) => values.is_empty(),)+ + } + } + + pub(crate) fn into_expr( + self, + membership: SetMembership, + min: PhysicalExprRef, + max: PhysicalExprRef, + ) -> PhysicalExprRef { + let data_type = self.data_type; + match self.values { + $( + PrimitiveValues::$variant(values) => Arc::new( + PrimitiveInListPruningExpr::<$arrow_type>::new( + membership, + data_type, + min, + max, + values, + ), + ), + )+ + } + } + } + }; +} + +// This single list defines both supported types and expression dispatch. +define_primitive_values! { + Int8, Int8Type, DataType::Int8, + ScalarValue::Int8(Some(value)), *value; + Int16, Int16Type, DataType::Int16, + ScalarValue::Int16(Some(value)), *value; + Int32, Int32Type, DataType::Int32, + ScalarValue::Int32(Some(value)), *value; + Int64, Int64Type, DataType::Int64, + ScalarValue::Int64(Some(value)), *value; + UInt8, UInt8Type, DataType::UInt8, + ScalarValue::UInt8(Some(value)), *value; + UInt16, UInt16Type, DataType::UInt16, + ScalarValue::UInt16(Some(value)), *value; + UInt32, UInt32Type, DataType::UInt32, + ScalarValue::UInt32(Some(value)), *value; + UInt64, UInt64Type, DataType::UInt64, + ScalarValue::UInt64(Some(value)), *value; + Decimal32, Decimal32Type, DataType::Decimal32(_precision, scale), + ScalarValue::Decimal32(Some(value), _value_precision, value_scale), *value, + if scale == value_scale; + Decimal64, Decimal64Type, DataType::Decimal64(_precision, scale), + ScalarValue::Decimal64(Some(value), _value_precision, value_scale), *value, + if scale == value_scale; + Decimal128, Decimal128Type, DataType::Decimal128(_precision, scale), + ScalarValue::Decimal128(Some(value), _value_precision, value_scale), *value, + if scale == value_scale; + Decimal256, Decimal256Type, DataType::Decimal256(_precision, scale), + ScalarValue::Decimal256(Some(value), _value_precision, value_scale), *value, + if scale == value_scale; + Date32, Date32Type, DataType::Date32, + ScalarValue::Date32(Some(value)), *value; + Date64, Date64Type, DataType::Date64, + ScalarValue::Date64(Some(value)), *value; + Time32Second, Time32SecondType, DataType::Time32(TimeUnit::Second), + ScalarValue::Time32Second(Some(value)), *value; + Time32Millisecond, Time32MillisecondType, + DataType::Time32(TimeUnit::Millisecond), + ScalarValue::Time32Millisecond(Some(value)), *value; + Time64Microsecond, Time64MicrosecondType, + DataType::Time64(TimeUnit::Microsecond), + ScalarValue::Time64Microsecond(Some(value)), *value; + Time64Nanosecond, Time64NanosecondType, + DataType::Time64(TimeUnit::Nanosecond), + ScalarValue::Time64Nanosecond(Some(value)), *value; + TimestampSecond, TimestampSecondType, + DataType::Timestamp(TimeUnit::Second, _timezone), + ScalarValue::TimestampSecond(Some(value), _value_timezone), *value; + TimestampMillisecond, TimestampMillisecondType, + DataType::Timestamp(TimeUnit::Millisecond, _timezone), + ScalarValue::TimestampMillisecond(Some(value), _value_timezone), *value; + TimestampMicrosecond, TimestampMicrosecondType, + DataType::Timestamp(TimeUnit::Microsecond, _timezone), + ScalarValue::TimestampMicrosecond(Some(value), _value_timezone), *value; + TimestampNanosecond, TimestampNanosecondType, + DataType::Timestamp(TimeUnit::Nanosecond, _timezone), + ScalarValue::TimestampNanosecond(Some(value), _value_timezone), *value; + DurationSecond, DurationSecondType, DataType::Duration(TimeUnit::Second), + ScalarValue::DurationSecond(Some(value)), *value; + DurationMillisecond, DurationMillisecondType, + DataType::Duration(TimeUnit::Millisecond), + ScalarValue::DurationMillisecond(Some(value)), *value; + DurationMicrosecond, DurationMicrosecondType, + DataType::Duration(TimeUnit::Microsecond), + ScalarValue::DurationMicrosecond(Some(value)), *value; + DurationNanosecond, DurationNanosecondType, + DataType::Duration(TimeUnit::Nanosecond), + ScalarValue::DurationNanosecond(Some(value)), *value; +} + +/// Tests an inclusive statistics interval against a sorted primitive domain. +struct PrimitiveInListPruningExpr +where + T::Native: Eq + Hash + Ord, +{ + membership: SetMembership, + // Retain decimal metadata and timestamp timezone for statistics casts. + data_type: DataType, + min: PhysicalExprRef, + max: PhysicalExprRef, + values: Arc<[T::Native]>, +} + +impl fmt::Debug for PrimitiveInListPruningExpr +where + T::Native: Eq + Hash + Ord, +{ + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("PrimitiveInListPruningExpr") + .field("membership", &self.membership) + .field("data_type", &self.data_type) + .field("min", &self.min) + .field("max", &self.max) + .field("values", &self.values) + .finish() + } +} + +impl PrimitiveInListPruningExpr +where + T::Native: Eq + Hash + Ord, +{ + fn new( + membership: SetMembership, + data_type: DataType, + min: PhysicalExprRef, + max: PhysicalExprRef, + mut values: Vec, + ) -> Self { + values.sort_unstable(); + values.dedup(); + Self { + membership, + data_type, + min, + max, + values: values.into(), + } + } + + fn contains(&self, value: T::Native) -> bool { + self.values.binary_search(&value).is_ok() + } + + fn may_match(&self, min: Option, max: Option) -> Option { + if self.membership == SetMembership::NotIn { + return not_in_may_match(min.as_ref(), max.as_ref(), |value| { + self.contains(*value) + }); + } + match (min, max) { + (Some(min), Some(max)) => { + // Inverted statistics are unusable, not proof that the domain + // and interval are disjoint. + if min > max { + return None; + } + let index = self.values.partition_point(|value| *value < min); + Some(self.values.get(index).is_some_and(|value| *value <= max)) + } + (Some(min), None) if self.values.last().is_some_and(|value| *value < min) => { + Some(false) + } + (None, Some(max)) + if self.values.first().is_some_and(|value| *value > max) => + { + Some(false) + } + _ => None, + } + } +} + +impl PartialEq for PrimitiveInListPruningExpr +where + T::Native: Eq + Hash + Ord, +{ + fn eq(&self, other: &Self) -> bool { + self.membership == other.membership + && self.data_type == other.data_type + && self.min.eq(&other.min) + && self.max.eq(&other.max) + && self.values == other.values + } +} + +impl Eq for PrimitiveInListPruningExpr where + T::Native: Eq + Hash + Ord +{ +} + +impl Hash for PrimitiveInListPruningExpr +where + T::Native: Eq + Hash + Ord, +{ + fn hash(&self, state: &mut H) { + self.membership.hash(state); + self.data_type.hash(state); + self.min.hash(state); + self.max.hash(state); + self.values.hash(state); + } +} + +impl Display for PrimitiveInListPruningExpr +where + T::Native: Eq + Hash + Ord, +{ + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!( + f, + "{}({}, {}, {} values)", + self.membership.display_name(), + self.min, + self.max, + self.values.len() + ) + } +} + +impl PhysicalExpr for PrimitiveInListPruningExpr +where + T::Native: Eq + Hash + Ord, +{ + fn data_type(&self, _input_schema: &Schema) -> Result { + Ok(DataType::Boolean) + } + + fn nullable(&self, _input_schema: &Schema) -> Result { + Ok(true) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + let min = self.min.evaluate(batch)?.into_array(batch.num_rows())?; + let max = self.max.evaluate(batch)?.into_array(batch.num_rows())?; + // Unlike view casts, dictionary-to-primitive casts propagate both key + // and value NULLs, so no logical-null intersection is necessary here. + let min = cast(&min, &self.data_type)?; + let max = cast(&max, &self.data_type)?; + let min = min.as_primitive::(); + let max = max.as_primitive::(); + let matches: BooleanArray = (0..batch.num_rows()) + .map(|index| { + let min = min.is_valid(index).then(|| min.value(index)); + let max = max.is_valid(index).then(|| max.value(index)); + self.may_match(min, max) + }) + .collect(); + Ok(ColumnarValue::Array(Arc::new(matches))) + } + + fn children(&self) -> Vec<&PhysicalExprRef> { + vec![&self.min, &self.max] + } + + fn with_new_children( + self: Arc, + children: Vec, + ) -> Result { + assert_eq_or_internal_err!(children.len(), 2); + Ok(Arc::new(Self { + membership: self.membership, + data_type: self.data_type.clone(), + min: Arc::clone(&children[0]), + max: Arc::clone(&children[1]), + values: Arc::clone(&self.values), + })) + } + + fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "{self}") + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::in_list::unwrap_scalar; + use arrow::datatypes::Field; + + #[test] + fn domain_accepts_compatible_non_null_values() { + let mut domain = PrimitiveInListDomain::new(&DataType::Int64, 4).unwrap(); + assert!(domain.push(&ScalarValue::Int64(None)).is_none()); + assert!(domain.push(&ScalarValue::Int32(Some(1))).is_none()); + assert!(domain.is_empty()); + + let run_end_encoded = ScalarValue::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int16, false)), + Arc::new(Field::new("values", DataType::Int64, true)), + Box::new(ScalarValue::Int64(Some(1))), + ); + assert!(domain.push(unwrap_scalar(&run_end_encoded)).is_some()); + assert!(!domain.is_empty()); + + let mut decimal = + PrimitiveInListDomain::new(&DataType::Decimal128(10, 2), 2).unwrap(); + assert!( + decimal + .push(&ScalarValue::Decimal128(Some(1), 12, 2)) + .is_some() + ); + assert!( + decimal + .push(&ScalarValue::Decimal128(Some(1), 10, 3)) + .is_none() + ); + + let mut timestamp = PrimitiveInListDomain::new( + &DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + 1, + ) + .unwrap(); + assert!( + timestamp + .push(&ScalarValue::TimestampMicrosecond( + Some(1), + Some("Asia/Kolkata".into()), + )) + .is_some() + ); + } +} diff --git a/datafusion/pruning/src/pruning_predicate.rs b/datafusion/pruning/src/pruning_predicate.rs index 614177546b473..2118892d089d2 100644 --- a/datafusion/pruning/src/pruning_predicate.rs +++ b/datafusion/pruning/src/pruning_predicate.rs @@ -22,7 +22,9 @@ use std::collections::HashSet; use std::sync::Arc; -use crate::string_in_list::{SetMembership, StringInListPruningExpr}; +use crate::in_list::{SetMembership, unwrap_scalar}; +use crate::primitive_in_list::PrimitiveInListDomain; +use crate::string_in_list::{BinaryInListPruningExpr, StringInListPruningExpr}; use arrow::array::AsArray; use arrow::{ @@ -461,8 +463,8 @@ impl<'a> PruningPredicateBuilder<'a> { /// | Condition | Pruning representation | /// | --- | --- | /// | `N <= min(20, C)` | Existing per-value rewrite | - /// | `20 < N <= C`, literal strings with optional NULLs on a string column | Compact pruning expression | - /// | `20 < N <= C`, other lists | Existing per-value rewrite | + /// | `20 < N <= C`, supported ordered literals with optional NULLs | Compact pruning expression | + /// | `20 < N <= C`, unsupported lists | Existing per-value rewrite | /// | `N > C` | Unhandled-predicate hook, normally "keep the container" | /// /// Empty lists also use the unhandled-predicate hook. A cap of zero disables @@ -470,9 +472,13 @@ impl<'a> PruningPredicateBuilder<'a> { /// pruning (such as Bloom filters). The default cap is [`MAX_IN_LIST_SIZE`] /// (20), so the compact path requires an explicitly raised cap. /// + /// Compact domains support strings, binary values, integers, decimals, + /// dates, times, timestamps, and durations. Floating-point values are + /// excluded because NaN and signed zero do not follow the required order. + /// /// The compact form covers `IN` and `NOT IN` alike, including lists with /// NULL members. Raising the cap can still build large comparison trees for - /// other eligible lists. + /// unsupported lists. /// /// Query engines typically pass /// `datafusion.execution.parquet.max_in_list_size` here. @@ -1502,8 +1508,25 @@ fn build_is_null_column_expr( } } -/// Keep large literal string lists compact instead of building a per-value -/// tree: an OR tree for `IN`, an AND chain for `NOT IN`. +/// Values collected for one supported compact IN-list representation. +enum CompactInListDomain { + String(Vec), + Binary(Vec>), + Primitive(PrimitiveInListDomain), +} + +impl CompactInListDomain { + fn is_empty(&self) -> bool { + match self { + Self::String(values) => values.is_empty(), + Self::Binary(values) => values.is_empty(), + Self::Primitive(values) => values.is_empty(), + } + } +} + +/// Keep large literal lists of supported ordered types compact instead of +/// building a per-value tree: an OR tree for `IN`, an AND chain for `NOT IN`. /// /// `IN` excludes a container whose interval is disjoint from the domain. That /// matches the per-value OR tree except for inverted bounds, which the compact @@ -1515,8 +1538,9 @@ fn build_is_null_column_expr( /// everywhere, including absent and inverted bounds. /// /// A NULL list member makes `NOT IN` and an all-NULL `IN` list never TRUE. For -/// other `IN` lists, NULL does not change which rows can make the predicate TRUE. -fn build_string_in_list_expr( +/// other `IN` lists, NULL does not change which rows can make the predicate TRUE, +/// but prevents inversion from proving that a container is fully matched. +fn build_compact_in_list_expr( in_list: &phys_expr::InListExpr, schema: &Schema, required_columns: &mut RequiredColumns, @@ -1533,42 +1557,79 @@ fn build_string_in_list_expr( DataType::Dictionary(_, value) => value.as_ref(), data_type => data_type, }; - if field.name() != column.name() || !data_type.is_string() { + if field.name() != column.name() { return None; } - let mut values = Vec::with_capacity(in_list.list().len()); + let mut domain = if data_type.is_string() { + CompactInListDomain::String(Vec::with_capacity(in_list.list().len())) + } else if matches!( + data_type, + DataType::Binary | DataType::LargeBinary | DataType::BinaryView + ) { + CompactInListDomain::Binary(Vec::with_capacity(in_list.list().len())) + } else { + CompactInListDomain::Primitive(PrimitiveInListDomain::new( + data_type, + in_list.list().len(), + )?) + }; let mut contains_null = false; for expr in in_list.list() { - if let Some(value) = extract_string_literal(expr) { - values.push(value.to_owned()); - } else if expr - .downcast_ref::() - .is_some_and(|literal| literal.value().is_null()) - { + let literal = expr.downcast_ref::()?; + let value = unwrap_scalar(literal.value()); + if value.is_null() { contains_null = true; } else { - return None; + match &mut domain { + CompactInListDomain::String(values) => { + values.push(unpack_string(value)?.to_owned()); + } + CompactInListDomain::Binary(values) => { + values.push(extract_binary(value)?.into()); + } + CompactInListDomain::Primitive(values) => values.push(value)?, + } } } // Pruning asks only whether the predicate can be TRUE. UNKNOWN and FALSE // both reject a row, so NOT IN with a NULL member and an all-NULL IN list // can never match. IN can otherwise ignore NULL and search its non-null domain. - if contains_null && (membership == SetMembership::NotIn || values.is_empty()) { + if contains_null && (membership == SetMembership::NotIn || domain.is_empty()) { properties.has_filter_semantics_only = true; return Some(Arc::new(phys_expr::Literal::new(ScalarValue::Boolean( Some(false), )))); } - let min = required_columns - .min_column_expr(column, in_list.expr(), field) - .ok()?; - let max = required_columns - .max_column_expr(column, in_list.expr(), field) - .ok()?; - let non_null = - build_is_null_column_expr(in_list.expr(), schema, required_columns, true)?; - let may_match = Arc::new(StringInListPruningExpr::new(membership, min, max, values)); + // Roll back appended statistics columns if the compact rewrite cannot be + // completed. `RequiredColumns::stat_column_expr` only appends entries. + let required_columns_len = required_columns.columns.len(); + let statistics = (|| { + let min = required_columns + .min_column_expr(column, in_list.expr(), field) + .ok()?; + let max = required_columns + .max_column_expr(column, in_list.expr(), field) + .ok()?; + let non_null = + build_is_null_column_expr(in_list.expr(), schema, required_columns, true)?; + Some((min, max, non_null)) + })(); + let Some((min, max, non_null)) = statistics else { + required_columns.columns.truncate(required_columns_len); + return None; + }; + let may_match = match domain { + CompactInListDomain::String(values) => { + Arc::new(StringInListPruningExpr::new(membership, min, max, values)) + as PhysicalExprRef + } + CompactInListDomain::Binary(values) => { + Arc::new(BinaryInListPruningExpr::new(membership, min, max, values)) + as PhysicalExprRef + } + CompactInListDomain::Primitive(values) => values.into_expr(membership, min, max), + }; if contains_null { properties.has_filter_semantics_only = true; } @@ -1580,9 +1641,9 @@ fn build_string_in_list_expr( } /// Default maximum number of entries in an `IN (...)` list eligible for -/// statistics pruning. Eligible literal string lists above this threshold use a -/// compact sorted domain instead of per-value min/max checks, for both `IN` and -/// `NOT IN`. +/// statistics pruning. Eligible literal lists of supported ordered types above +/// this threshold use a compact sorted domain instead of per-value min/max +/// checks, for both `IN` and `NOT IN`. /// Callers can raise the cap via [`PredicateRewriter::with_max_in_list_size`], and /// query engines can wire it from the /// `datafusion.execution.parquet.max_in_list_size` config option. @@ -1670,9 +1731,9 @@ impl PredicateRewriter { /// Returns the pruning predicate as an [`PhysicalExpr`] /// /// `max_in_list_size` is the largest `IN (...)` list eligible for statistics -/// pruning. Large literal string lists use a compact representation, for both -/// `IN` and `NOT IN`; other eligible lists use per-value checks. Longer lists -/// fall back to `unhandled_hook`. +/// pruning. Large literal lists of supported ordered types use a compact +/// representation for both `IN` and `NOT IN`; unsupported lists use per-value +/// checks. Longer lists fall back to `unhandled_hook`. fn build_predicate_expression( expr: &Arc, schema: &SchemaRef, @@ -1721,7 +1782,7 @@ fn build_predicate_expression( if in_list.list().len() > MAX_IN_LIST_SIZE && in_list.list().len() <= max_in_list_size && let Some(pruning_expr) = - build_string_in_list_expr(in_list, schema, required_columns, properties) + build_compact_in_list_expr(in_list, schema, required_columns, properties) { return pruning_expr; } @@ -2090,6 +2151,15 @@ fn extract_string_literal(expr: &Arc) -> Option<&str> { None } +fn extract_binary(value: &ScalarValue) -> Option<&[u8]> { + match value { + ScalarValue::Binary(value) + | ScalarValue::LargeBinary(value) + | ScalarValue::BinaryView(value) => value.as_deref(), + _ => None, + } +} + /// Wrap a string in a `Literal` whose `ScalarValue` matches `target_type` fn string_literal_as(value: String, target_type: &DataType) -> Arc { let utf8 = ScalarValue::Utf8(Some(value)); @@ -2332,8 +2402,8 @@ mod tests { BinaryArray, DictionaryArray, Int32Array, Int64Array, StringArray, UInt64Array, }, - buffer::NullBuffer, - datatypes::{Int32Type, TimeUnit}, + buffer::{NullBuffer, ScalarBuffer}, + datatypes::{Int32Type, Int64Type, TimeUnit, i256}, }; use datafusion_expr::expr::InList; use datafusion_expr::{BinaryExpr, Expr, cast, is_null, try_cast}; @@ -3608,13 +3678,8 @@ mod tests { Ok(()) } - // With the configurable cap, a caller that raises - // `max_in_list_size` above the default gets the IN list rewritten - // into a per-value min/max chain instead of falling through to `true`. - // This verifies both `PredicateRewriter::with_max_in_list_size` and the - // recursive OR path inside `build_predicate_expression`. #[test] - fn row_group_predicate_in_list_rewritten_at_raised_cap() -> Result<()> { + fn row_group_predicate_in_list_compacted_at_raised_cap() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, false)])); // 25 items — above the default 20, below a raised cap of 32. @@ -3623,19 +3688,9 @@ mod tests { let rewriter = PredicateRewriter::new().with_max_in_list_size(32); let predicate_expr = rewriter.rewrite_predicate_to_statistics_predicate(&physical, &schema); - // At the raised cap, IN is rewritten into per-value min/max checks - // OR'd together; the resulting predicate must not collapse to - // `true` (which is what the default cap produces). - assert_ne!( - predicate_expr.to_string(), - "true", - "IN(25) with raised cap must rewrite into a statistics-based predicate, not fall through to `true`" - ); - // Sanity: the rewritten predicate references per-value literals. assert!( - predicate_expr.to_string().contains(" <= 1 ") - && predicate_expr.to_string().contains(" <= 25 "), - "rewritten predicate should include per-value bounds for each IN entry, got: {predicate_expr}" + predicate_expr.to_string().contains("IN_SET_INTERSECTS"), + "raised-cap predicate should use compact pruning, got: {predicate_expr}" ); Ok(()) } @@ -3663,7 +3718,7 @@ mod tests { // The high-level [`PruningPredicateBuilder`] should thread // `max_in_list_size` all the way through: a 25-item IN with the default // cap must fall through to the unhandled hook (`predicate_expr = true`), - // while a raised cap produces a real per-value statistics predicate. + // while a raised cap produces a compact statistics predicate. #[test] fn pruning_predicate_builder_threads_max_in_list_size() -> Result<()> { let schema = @@ -3683,8 +3738,6 @@ mod tests { "default cap must fall through to `true` for 25-item IN" ); - // Raising the cap produces a real statistics predicate with per- - // value bounds. let raised_pp = PruningPredicateBuilder::new() .with_file_schema(Arc::clone(&schema)) .with_max_in_list_size(32) @@ -3695,8 +3748,8 @@ mod tests { "raised cap must produce a real statistics predicate for 25-item IN" ); assert!( - raised_expr.contains(" <= 1 ") && raised_expr.contains(" <= 25 "), - "raised-cap predicate should include per-value bounds, got: {raised_expr}" + raised_expr.contains("IN_SET_INTERSECTS"), + "raised-cap predicate should use compact pruning, got: {raised_expr}" ); Ok(()) } @@ -3725,7 +3778,7 @@ mod tests { Ok(()) } - fn large_string_pruning_predicate( + fn large_in_list_pruning_predicate( expr: PhysicalExprRef, schema: SchemaRef, ) -> Result { @@ -3735,6 +3788,388 @@ mod tests { .try_build(expr) } + fn assert_per_value_fallback( + predicate: &PruningPredicate, + expected_literals: &[ScalarValue], + ) -> Result<()> { + let expression = predicate.predicate_expr(); + let display = expression.to_string(); + assert_ne!(display, "true"); + assert!(!display.contains("IN_SET_INTERSECTS")); + assert!(!display.contains("NOT_IN_SET_MAY_MATCH")); + assert!(display.contains("_min")); + assert!(display.contains("_max")); + let mut nodes = 0; + let mut literals = vec![]; + expression.apply(|expr| { + nodes += 1; + if let Some(literal) = expr.downcast_ref::() { + literals.push(literal.value().clone()); + } + Ok(TreeNodeRecursion::Continue) + })?; + assert!(nodes > 7, "expected a per-value tree, got {expression}"); + for expected in expected_literals { + assert!( + literals.contains(expected), + "expected literal {expected} in per-value tree {expression}" + ); + } + Ok(()) + } + + #[test] + fn compact_extraction_unwraps_encoded_binary_literals() { + let value = ScalarValue::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int16, false)), + Arc::new(Field::new("values", DataType::Binary, true)), + Box::new(ScalarValue::Binary(Some(vec![1, 2]))), + ); + assert_eq!(extract_binary(unwrap_scalar(&value)), Some(&[1, 2][..])); + } + + fn ordered_scalar(data_type: &DataType, value: Option) -> ScalarValue { + match data_type { + DataType::Int8 => ScalarValue::Int8(value.map(|value| value as i8)), + DataType::Int16 => ScalarValue::Int16(value.map(|value| value as i16)), + DataType::Int32 => ScalarValue::Int32(value.map(|value| value as i32)), + DataType::Int64 => ScalarValue::Int64(value), + DataType::UInt8 => ScalarValue::UInt8(value.map(|value| value as u8)), + DataType::UInt16 => ScalarValue::UInt16(value.map(|value| value as u16)), + DataType::UInt32 => ScalarValue::UInt32(value.map(|value| value as u32)), + DataType::UInt64 => ScalarValue::UInt64(value.map(|value| value as u64)), + DataType::Decimal32(precision, scale) => ScalarValue::Decimal32( + value.map(|value| value as i32), + *precision, + *scale, + ), + DataType::Decimal64(precision, scale) => { + ScalarValue::Decimal64(value, *precision, *scale) + } + DataType::Decimal128(precision, scale) => { + ScalarValue::Decimal128(value.map(i128::from), *precision, *scale) + } + DataType::Decimal256(precision, scale) => { + ScalarValue::Decimal256(value.map(i256::from), *precision, *scale) + } + DataType::Date32 => ScalarValue::Date32(value.map(|value| value as i32)), + DataType::Date64 => ScalarValue::Date64(value), + DataType::Time32(TimeUnit::Second) => { + ScalarValue::Time32Second(value.map(|value| value as i32)) + } + DataType::Time32(TimeUnit::Millisecond) => { + ScalarValue::Time32Millisecond(value.map(|value| value as i32)) + } + DataType::Time64(TimeUnit::Microsecond) => { + ScalarValue::Time64Microsecond(value) + } + DataType::Time64(TimeUnit::Nanosecond) => { + ScalarValue::Time64Nanosecond(value) + } + DataType::Timestamp(TimeUnit::Second, timezone) => { + ScalarValue::TimestampSecond(value, timezone.clone()) + } + DataType::Timestamp(TimeUnit::Millisecond, timezone) => { + ScalarValue::TimestampMillisecond(value, timezone.clone()) + } + DataType::Timestamp(TimeUnit::Microsecond, timezone) => { + ScalarValue::TimestampMicrosecond(value, timezone.clone()) + } + DataType::Timestamp(TimeUnit::Nanosecond, timezone) => { + ScalarValue::TimestampNanosecond(value, timezone.clone()) + } + DataType::Duration(TimeUnit::Second) => ScalarValue::DurationSecond(value), + DataType::Duration(TimeUnit::Millisecond) => { + ScalarValue::DurationMillisecond(value) + } + DataType::Duration(TimeUnit::Microsecond) => { + ScalarValue::DurationMicrosecond(value) + } + DataType::Duration(TimeUnit::Nanosecond) => { + ScalarValue::DurationNanosecond(value) + } + DataType::Binary => { + ScalarValue::Binary(value.map(|value| vec![0xff, value as u8])) + } + DataType::LargeBinary => { + ScalarValue::LargeBinary(value.map(|value| vec![0xff, value as u8])) + } + DataType::BinaryView => { + ScalarValue::BinaryView(value.map(|value| vec![0xff, value as u8])) + } + DataType::Dictionary(key_type, value_type) => ScalarValue::Dictionary( + key_type.clone(), + Box::new(ordered_scalar(value_type, value)), + ), + data_type => panic!("unsupported test type {data_type:?}"), + } + } + + fn ordered_statistics(data_type: &DataType) -> ContainerStats { + let statistics_type = match data_type { + DataType::Dictionary(_, value_type) => value_type.as_ref(), + data_type => data_type, + }; + let min = [Some(0), Some(1), Some(1), Some(50), None, Some(41), Some(4)]; + let max = [Some(0), Some(1), Some(2), Some(51), Some(2), None, Some(2)]; + ContainerStats::new() + .with_min( + ScalarValue::iter_to_array( + min.into_iter() + .map(|value| ordered_scalar(statistics_type, value)), + ) + .unwrap(), + ) + .with_max( + ScalarValue::iter_to_array( + max.into_iter() + .map(|value| ordered_scalar(statistics_type, value)), + ) + .unwrap(), + ) + .with_null_counts([Some(0); 7]) + .with_row_counts([Some(1); 7]) + } + + #[test] + fn large_ordered_in_lists_use_compact_pruning() -> Result<()> { + let types = [ + DataType::Int8, + DataType::Int16, + DataType::Int32, + DataType::Int64, + DataType::UInt8, + DataType::UInt16, + DataType::UInt32, + DataType::UInt64, + DataType::Decimal32(9, 2), + DataType::Decimal64(18, 3), + DataType::Decimal128(30, 4), + DataType::Decimal256(60, 5), + DataType::Date32, + DataType::Date64, + DataType::Time32(TimeUnit::Second), + DataType::Time32(TimeUnit::Millisecond), + DataType::Time64(TimeUnit::Microsecond), + DataType::Time64(TimeUnit::Nanosecond), + DataType::Timestamp(TimeUnit::Second, None), + DataType::Timestamp(TimeUnit::Millisecond, None), + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + DataType::Timestamp(TimeUnit::Nanosecond, None), + DataType::Duration(TimeUnit::Second), + DataType::Duration(TimeUnit::Millisecond), + DataType::Duration(TimeUnit::Microsecond), + DataType::Duration(TimeUnit::Nanosecond), + DataType::Binary, + DataType::LargeBinary, + DataType::BinaryView, + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Int64)), + DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Binary)), + ]; + + for data_type in types { + let schema = + Arc::new(Schema::new(vec![Field::new("c1", data_type.clone(), true)])); + let values = (0..21) + .map(|index| { + Arc::new(phys_expr::Literal::new(ordered_scalar( + &data_type, + Some(index * 2), + ))) as PhysicalExprRef + }) + .collect::>(); + let stats = TestStatistics::new().with("c1", ordered_statistics(&data_type)); + + for (negated, expected) in [ + (false, [true, false, true, false, true, false, true]), + (true, [false, true, true, true, true, true, true]), + ] { + let expr = phys_expr::in_list( + Arc::new(phys_expr::Column::new("c1", 0)), + values.clone(), + &negated, + &schema, + )?; + let predicate = large_in_list_pruning_predicate(expr, schema.clone())?; + assert!( + predicate.predicate_expr().to_string().contains(if negated { + "NOT_IN_SET_MAY_MATCH" + } else { + "IN_SET_INTERSECTS" + }), + "type={data_type:?}, negated={negated}" + ); + assert_eq!( + predicate.prune(&stats)?, + expected, + "type={data_type:?}, negated={negated}" + ); + assert_eq!(predicate.required_columns.columns.len(), 4); + } + } + Ok(()) + } + + #[test] + fn large_float_in_list_keeps_per_value_rewrite() -> Result<()> { + let data_type = DataType::Float64; + let schema = + Arc::new(Schema::new(vec![Field::new("c1", data_type.clone(), true)])); + let values = (0..21) + .map(|value| { + Arc::new(phys_expr::Literal::new(ScalarValue::Float64(Some( + value as f64, + )))) as PhysicalExprRef + }) + .collect::>(); + for negated in [false, true] { + let expr = phys_expr::in_list( + Arc::new(phys_expr::Column::new("c1", 0)), + values.clone(), + &negated, + &schema, + )?; + let predicate = large_in_list_pruning_predicate(expr, Arc::clone(&schema))?; + assert_per_value_fallback( + &predicate, + &[ + ScalarValue::Float64(Some(0.0)), + ScalarValue::Float64(Some(20.0)), + ], + )?; + } + Ok(()) + } + + #[test] + fn large_fixed_size_binary_in_list_keeps_per_value_rewrite() -> Result<()> { + let data_type = DataType::FixedSizeBinary(2); + let schema = + Arc::new(Schema::new(vec![Field::new("c1", data_type.clone(), true)])); + let values = (0..21) + .map(|value| { + Arc::new(phys_expr::Literal::new(ScalarValue::FixedSizeBinary( + 2, + Some(vec![0xff, value]), + ))) as PhysicalExprRef + }) + .collect::>(); + for negated in [false, true] { + let expr = phys_expr::in_list( + Arc::new(phys_expr::Column::new("c1", 0)), + values.clone(), + &negated, + &schema, + )?; + let predicate = large_in_list_pruning_predicate(expr, Arc::clone(&schema))?; + assert_per_value_fallback( + &predicate, + &[ + ScalarValue::FixedSizeBinary(2, Some(vec![0xff, 0])), + ScalarValue::FixedSizeBinary(2, Some(vec![0xff, 20])), + ], + )?; + } + Ok(()) + } + + #[test] + fn large_primitive_in_list_preserves_dictionary_nulls() -> Result<()> { + let data_type = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Int64)); + let schema = + Arc::new(Schema::new(vec![Field::new("c1", data_type.clone(), true)])); + let values = (0..21) + .map(|value| { + Arc::new(phys_expr::Literal::new(ordered_scalar( + &data_type, + Some(value * 2), + ))) as PhysicalExprRef + }) + .collect::>(); + let expr = phys_expr::in_list( + Arc::new(phys_expr::Column::new("c1", 0)), + values, + &false, + &schema, + )?; + let predicate = large_in_list_pruning_predicate(expr, schema)?; + + // The invalid dictionary values have nonzero payloads outside the + // domain. They must remain missing bounds after dictionary casting. + let dictionary_values = Arc::new(arrow::array::PrimitiveArray::::new( + ScalarBuffer::from(vec![50, 51, -1, -1]), + Some(NullBuffer::from(vec![false, true, true, false])), + )); + let dictionary = |keys| -> Result { + Ok(Arc::new(DictionaryArray::::try_new( + Int32Array::from(keys), + dictionary_values.clone(), + )?)) + }; + let stats = TestStatistics::new().with( + "c1", + ContainerStats::new() + .with_min(dictionary(vec![Some(0), Some(2)])?) + .with_max(dictionary(vec![Some(1), Some(3)])?) + .with_null_counts([Some(0); 2]) + .with_row_counts([Some(1); 2]), + ); + assert_eq!(predicate.prune(&stats)?, [true, true]); + Ok(()) + } + + #[test] + fn large_binary_in_list_preserves_dictionary_nulls() -> Result<()> { + let data_type = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Binary)); + let schema = + Arc::new(Schema::new(vec![Field::new("c1", data_type.clone(), true)])); + let values = (0..21) + .map(|value| { + Arc::new(phys_expr::Literal::new(ordered_scalar( + &data_type, + Some(value * 2), + ))) as PhysicalExprRef + }) + .collect::>(); + let expr = phys_expr::in_list( + Arc::new(phys_expr::Column::new("c1", 0)), + values, + &false, + &schema, + )?; + let predicate = large_in_list_pruning_predicate(expr, schema)?; + + // BinaryView casts currently lose NULLs in dictionary values. Use + // nonempty invalid payloads so the test fails if they become bounds. + let payloads = [vec![0xff, 50], vec![0xff, 51], vec![0], vec![0]]; + let (offsets, values, _) = + BinaryArray::from_iter_values(payloads.iter().map(Vec::as_slice)) + .into_parts(); + let dictionary_values: ArrayRef = Arc::new(BinaryArray::new( + offsets, + values, + Some(NullBuffer::from(vec![false, true, true, false])), + )); + let dictionary = |keys| -> Result { + Ok(Arc::new(DictionaryArray::::try_new( + Int32Array::from(keys), + Arc::clone(&dictionary_values), + )?)) + }; + let stats = TestStatistics::new().with( + "c1", + ContainerStats::new() + .with_min(dictionary(vec![Some(0), Some(2)])?) + .with_max(dictionary(vec![Some(1), Some(3)])?) + .with_null_counts([Some(0); 2]) + .with_row_counts([Some(1); 2]), + ); + assert_eq!(predicate.prune(&stats)?, [true, true]); + Ok(()) + } + #[test] fn large_string_in_list_prunes_exact_intervals() -> Result<()> { let types = [ @@ -3764,7 +4199,7 @@ mod tests { &schema, )?; let predicate = - large_string_pruning_predicate(Arc::clone(&expr), schema)?; + large_in_list_pruning_predicate(Arc::clone(&expr), schema)?; let last_value = format!("k{:06}", (count - 1) * 10); let stats = TestStatistics::new().with( "c1", @@ -3851,7 +4286,7 @@ mod tests { &false, &schema, )?; - let predicate = large_string_pruning_predicate(expr, schema)?; + let predicate = large_in_list_pruning_predicate(expr, schema)?; // NULL dictionary values can have nonempty payloads. Neither payload // may become a bound when a valid key references the NULL value. let (offsets, values, _) = @@ -3912,7 +4347,7 @@ mod tests { ]); let expr = col("c1").in_list(values, false); let predicate = - large_string_pruning_predicate(logical2physical(&expr, &schema), schema)?; + large_in_list_pruning_predicate(logical2physical(&expr, &schema), schema)?; let stats = TestStatistics::new().with( "c1", ContainerStats::new_utf8( @@ -3998,7 +4433,7 @@ mod tests { } #[test] - fn large_string_in_list_compacts_null_literals() -> Result<()> { + fn large_in_list_compacts_null_literals() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)])); let values = (0..21).map(|i| lit(format!("a{i:03}"))).collect::>(); let stats = TestStatistics::new().with( @@ -4026,7 +4461,7 @@ mod tests { let positive_or = col("c1") .in_list(with_null.clone(), false) .or(col("c1").eq(lit("middle"))); - let predicate = large_string_pruning_predicate( + let predicate = large_in_list_pruning_predicate( logical2physical(&positive_or, &schema), Arc::clone(&schema), )?; @@ -4043,7 +4478,7 @@ mod tests { assert_eq!(default.prune(&stats)?, [true, true]); assert!(default.can_be_inverted_for_full_match()); - let raised = large_string_pruning_predicate(physical, Arc::clone(&schema))?; + let raised = large_in_list_pruning_predicate(physical, Arc::clone(&schema))?; assert!(!raised.can_be_inverted_for_full_match()); if negated { assert!(is_always_false(raised.predicate_expr()), "{expr}"); @@ -4064,7 +4499,7 @@ mod tests { std::iter::repeat_n(lit(ScalarValue::Utf8(None)), 21).collect::>(); for negated in [false, true] { let expr = col("c1").in_list(all_null.clone(), negated); - let predicate = large_string_pruning_predicate( + let predicate = large_in_list_pruning_predicate( logical2physical(&expr, &schema), Arc::clone(&schema), )?; @@ -4077,11 +4512,17 @@ mod tests { Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, true)])); let mut integer_values = (0..21).map(lit).collect::>(); integer_values.push(lit(ScalarValue::Int32(None))); - let integer_predicate = large_string_pruning_predicate( + let integer_predicate = large_in_list_pruning_predicate( logical2physical(&col("c1").in_list(integer_values, false), &integer_schema), integer_schema, )?; - assert!(integer_predicate.can_be_inverted_for_full_match()); + assert!( + integer_predicate + .predicate_expr() + .to_string() + .contains("IN_SET_INTERSECTS") + ); + assert!(!integer_predicate.can_be_inverted_for_full_match()); // The compact false predicate has the same filter result as the original // NOT IN expression, which returns UNKNOWN for values outside the list. @@ -4094,6 +4535,66 @@ mod tests { Ok(()) } + #[test] + fn compact_binary_and_primitive_null_lists_are_always_false() -> Result<()> { + let cases = [ + ( + DataType::Binary, + (0..21) + .map(|value| ScalarValue::Binary(Some(vec![value]))) + .collect::>(), + ScalarValue::Binary(None), + ), + ( + DataType::Int32, + (0..21) + .map(|value| ScalarValue::Int32(Some(value))) + .collect(), + ScalarValue::Int32(None), + ), + ]; + + for (data_type, values, null) in cases { + let schema = + Arc::new(Schema::new(vec![Field::new("c1", data_type.clone(), true)])); + let column = Arc::new(phys_expr::Column::new("c1", 0)); + + let mut with_null = values; + with_null.push(null.clone()); + let not_in = phys_expr::in_list( + Arc::clone(&column) as PhysicalExprRef, + with_null + .into_iter() + .map(|value| { + Arc::new(phys_expr::Literal::new(value)) as PhysicalExprRef + }) + .collect(), + &true, + &schema, + )?; + let predicate = large_in_list_pruning_predicate(not_in, Arc::clone(&schema))?; + assert!(is_always_false(predicate.predicate_expr())); + assert!(!predicate.can_be_inverted_for_full_match()); + assert!(predicate.required_columns.columns.is_empty()); + + let all_null = phys_expr::in_list( + column, + std::iter::repeat_n(null, 21) + .map(|value| { + Arc::new(phys_expr::Literal::new(value)) as PhysicalExprRef + }) + .collect(), + &false, + &schema, + )?; + let predicate = large_in_list_pruning_predicate(all_null, schema)?; + assert!(is_always_false(predicate.predicate_expr())); + assert!(!predicate.can_be_inverted_for_full_match()); + assert!(predicate.required_columns.columns.is_empty()); + } + Ok(()) + } + /// Statistics shared by the compact `NOT IN` tests, one row per case. fn not_in_container_stats() -> TestStatistics { TestStatistics::new().with( @@ -4164,7 +4665,7 @@ mod tests { &true, &schema, )?; - let predicate = large_string_pruning_predicate(expr, schema)?; + let predicate = large_in_list_pruning_predicate(expr, schema)?; assert_eq!( predicate.predicate_expr().to_string(), "c1_null_count@3 != row_count@2 AND NOT_IN_SET_MAY_MATCH(c1_min@0, c1_max@1, 21 values)", @@ -4214,7 +4715,7 @@ mod tests { let not_in = col("c1").in_list((0..21).map(|i| lit(format!("a{i:03}"))).collect(), true); let other = col("c2").in_list((0..21).map(|i| lit(i * 10)).collect(), false); - let predicate = large_string_pruning_predicate( + let predicate = large_in_list_pruning_predicate( logical2physical(¬_in.or(other), &schema), Arc::clone(&schema), )?; @@ -4254,13 +4755,13 @@ mod tests { fn large_string_not_in_list_inverts_without_false_full_match() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)])); let values = (0..21).map(|i| lit(format!("a{i:03}"))).collect::>(); - let forward = large_string_pruning_predicate( + let forward = large_in_list_pruning_predicate( logical2physical(&col("c1").in_list(values.clone(), true), &schema), Arc::clone(&schema), )?; // NOT(c1 NOT IN (...)) OR c1 IS NULL, the shape row_group_filter builds // once PhysicalExprSimplifier turns NOT(NOT IN) back into IN. - let inverted = large_string_pruning_predicate( + let inverted = large_in_list_pruning_predicate( logical2physical( &col("c1").in_list(values, false).or(col("c1").is_null()), &schema, @@ -4306,7 +4807,7 @@ mod tests { values.iter().map(|value| lit(value.clone())).collect(), true, ); - let compact = large_string_pruning_predicate( + let compact = large_in_list_pruning_predicate( logical2physical(¬_in, &schema), Arc::clone(&schema), )?; @@ -4318,7 +4819,7 @@ mod tests { .map(|value| col("c1").not_eq(lit(value.clone()))) .reduce(Expr::and) .unwrap(); - let per_value = large_string_pruning_predicate( + let per_value = large_in_list_pruning_predicate( logical2physical(&chain, &schema), Arc::clone(&schema), )?; diff --git a/datafusion/pruning/src/string_in_list.rs b/datafusion/pruning/src/string_in_list.rs index b9d9f68b30511..da1793f5d03a5 100644 --- a/datafusion/pruning/src/string_in_list.rs +++ b/datafusion/pruning/src/string_in_list.rs @@ -15,7 +15,6 @@ // specific language governing permissions and limitations // under the License. -use std::cmp::Ordering; use std::fmt::{self, Display, Formatter}; use std::hash::{Hash, Hasher}; use std::sync::Arc; @@ -28,115 +27,7 @@ use datafusion_common::{Result, assert_eq_or_internal_err}; use datafusion_physical_expr::{PhysicalExpr, PhysicalExprRef}; use datafusion_physical_plan::ColumnarValue; -/// Which `IN` form a sorted string domain is pruning for. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub(crate) enum SetMembership { - /// `col IN (...)`. A row matches only where the domain intersects the - /// interval, so a disjoint interval excludes every row. - In, - /// `col NOT IN (...)`. Overlap proves nothing here: values outside the - /// domain still satisfy the predicate. An interval excludes every row only - /// when it holds a single value that the domain contains. - NotIn, -} - -impl SetMembership { - fn compare_values(self, left: &[u8], right: &[u8]) -> Ordering { - match self { - // IN uses this order for interval searches. - Self::In => left.cmp(right), - // NOT IN only needs exact membership. Reject impossible lengths - // before comparing bytes that may have a long common prefix. - Self::NotIn => left.len().cmp(&right.len()).then_with(|| left.cmp(right)), - } - } -} - -/// Tests an inclusive statistics interval against a sorted string domain. -/// -/// [`PhysicalExpr::evaluate`] returns one nullable Boolean per min/max interval: -/// * `true`: matching rows may exist, so the container must be read. -/// * `false`: the available bounds prove no row can match. -/// * `NULL`: incomplete, invalid, or unusable bounds prevent a safe decision. -/// -/// [`SetMembership`] selects the test. For [`SetMembership::In`] a single known -/// bound can still prove disjointness. For [`SetMembership::NotIn`], one known -/// bound outside the domain proves the container may match, while two equal -/// bounds in the domain prove it cannot. This is the same reach as the per-value -/// `min != v OR v != max` chain it replaces. Otherwise, unknown results keep the -/// container eligible for reading. -/// -/// This expression is used only for pruning; the original IN remains the row filter. -#[derive(Debug, Eq)] -pub(crate) struct StringInListPruningExpr { - membership: SetMembership, - min: PhysicalExprRef, - max: PhysicalExprRef, - values: Arc<[String]>, -} - -impl StringInListPruningExpr { - pub(crate) fn new( - membership: SetMembership, - min: PhysicalExprRef, - max: PhysicalExprRef, - mut values: Vec, - ) -> Self { - values.sort_unstable_by(|left, right| { - membership.compare_values(left.as_bytes(), right.as_bytes()) - }); - values.dedup(); - Self { - membership, - min, - max, - values: values.into(), - } - } - - /// Does the sorted, deduplicated domain hold `value`? - fn contains(&self, value: &[u8]) -> bool { - self.values - .binary_search_by(|candidate| { - self.membership.compare_values(candidate.as_bytes(), value) - }) - .is_ok() - } -} - -impl PartialEq for StringInListPruningExpr { - fn eq(&self, other: &Self) -> bool { - self.membership == other.membership - && self.min.eq(&other.min) - && self.max.eq(&other.max) - && self.values == other.values - } -} - -impl Hash for StringInListPruningExpr { - fn hash(&self, state: &mut H) { - self.membership.hash(state); - self.min.hash(state); - self.max.hash(state); - self.values.hash(state); - } -} - -impl Display for StringInListPruningExpr { - fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { - let name = match self.membership { - SetMembership::In => "IN_SET_INTERSECTS", - SetMembership::NotIn => "NOT_IN_SET_MAY_MATCH", - }; - write!( - f, - "{name}({}, {}, {} values)", - self.min, - self.max, - self.values.len() - ) - } -} +use crate::in_list::{SetMembership, not_in_may_match}; fn has_oversized_string_buffer(array: &dyn Array, limit: usize) -> bool { match array.data_type() { @@ -150,166 +41,317 @@ fn has_oversized_string_buffer(array: &dyn Array, limit: usize) -> bool { } } -impl PhysicalExpr for StringInListPruningExpr { - fn data_type(&self, _input_schema: &Schema) -> Result { - Ok(DataType::Boolean) +fn has_oversized_binary_buffer(array: &dyn Array, limit: usize) -> bool { + match array.data_type() { + DataType::Binary => array.as_binary::().values().len() >= limit, + DataType::LargeBinary => array.as_binary::().values().len() >= limit, + DataType::Dictionary(_, _) => has_oversized_binary_buffer( + array.as_any_dictionary().values().as_ref(), + limit, + ), + _ => false, } +} - fn nullable(&self, _input_schema: &Schema) -> Result { - Ok(true) - } +macro_rules! string_bytes { + ($value:expr) => { + ($value).as_bytes() + }; +} + +macro_rules! binary_bytes { + ($value:expr) => { + &($value)[..] + }; +} - fn evaluate(&self, batch: &RecordBatch) -> Result { - // Normalize Utf8, LargeUtf8, Utf8View, and dictionary-encoded statistics. - let min = self.min.evaluate(batch)?.into_array(batch.num_rows())?; - let max = self.max.evaluate(batch)?.into_array(batch.num_rows())?; - // A short string slice can retain a buffer too large for Utf8View's - // u32 offsets. Avoid a panic in the cast and keep pruning conservative. - if has_oversized_string_buffer(min.as_ref(), u32::MAX as usize) - || has_oversized_string_buffer(max.as_ref(), u32::MAX as usize) - { - return Ok(ColumnarValue::Array(Arc::new(BooleanArray::new_null( - batch.num_rows(), - )))); +/// Defines a byte-domain pruning expression for one variable-length value type. +macro_rules! define_byte_in_list_expr { + ( + $name:ident, + $value_type:ty, + $view_type:expr, + $as_view:ident, + $bytes:ident, + $has_oversized_buffer:path + ) => { + /// Tests an inclusive statistics interval against a sorted byte domain. + /// + /// [`PhysicalExpr::evaluate`] returns one nullable Boolean per min/max + /// interval: `true` means matching rows may exist, `false` means the + /// available bounds prove no row can match, and `NULL` means the bounds + /// cannot support a safe decision. + /// + /// [`SetMembership::In`] can use one known bound to prove disjointness. + /// For [`SetMembership::NotIn`], one known bound outside the domain + /// proves the container may match, while equal bounds in the domain + /// prove it cannot. Otherwise, the result keeps the container eligible. + /// The original `IN` expression remains the row filter. + #[derive(Debug, Eq)] + pub(crate) struct $name { + membership: SetMembership, + min: PhysicalExprRef, + max: PhysicalExprRef, + values: Arc<[$value_type]>, } - // Dictionary values can be NULL behind valid keys. Preserve their - // validity even if the view cast only carries the key nulls. - // TODO: Revisit this workaround once the Arrow dependency includes - // https://github.com/apache/arrow-rs/pull/10510. - let min_nulls = min.logical_nulls(); - let max_nulls = max.logical_nulls(); - let min = cast(&min, &DataType::Utf8View)?; - let max = cast(&max, &DataType::Utf8View)?; - let min = min.as_string_view(); - let max = max.as_string_view(); - let matches: BooleanArray = (0..batch.num_rows()) - .map(|i| { - let min = (min.is_valid(i) - && min_nulls.as_ref().is_none_or(|nulls| nulls.is_valid(i))) - .then(|| min.value(i).as_bytes()); - let max = (max.is_valid(i) - && max_nulls.as_ref().is_none_or(|nulls| nulls.is_valid(i))) - .then(|| max.value(i).as_bytes()); - if self.membership == SetMembership::NotIn { - return match (min, max) { - // Check membership before comparing the bounds. NOT IN - // ordering rejects different byte lengths first, avoiding - // scans of long common prefixes. - (Some(min), Some(max)) if self.contains(min) => { - // A wider interval can hold a value outside the domain, - // which satisfies NOT IN. Only an interval pinned to one - // domain value rules out every row. Truncated Parquet - // bounds cannot fake that: min truncates downward and max - // upward, so equal bounds mean the true values were equal. - Some(min != max) - } - (Some(_), Some(_)) => Some(true), - // One known bound outside the domain is enough to preserve - // the true result that lets an enclosing OR short-circuit. - (Some(bound), None) | (None, Some(bound)) - if !self.contains(bound) => - { - Some(true) - } - _ => None, - }; + + impl $name { + pub(crate) fn new( + membership: SetMembership, + min: PhysicalExprRef, + max: PhysicalExprRef, + mut values: Vec<$value_type>, + ) -> Self { + values.sort_unstable_by(|left, right| { + membership.compare_bytes($bytes!(left), $bytes!(right)) + }); + values.dedup(); + Self { + membership, + min, + max, + values: values.into(), } - match (min, max) { - (Some(min), Some(max)) => { - if min > max { - return None; - } - // Rust string ordering and these byte comparisons both use - // unsigned lexicographic UTF-8 order, as required by the - // PruningStatistics min/max contract. Parquet adapters mask - // bounds with unusable ordering; PartitionPruningStatistics - // uses actual Arrow partition values. PrunableStatistics - // trusts file providers' bounds: there is no ordering gate - // for arbitrary statistics providers here. - let index = self.values.partition_point(|v| v.as_bytes() < min); - Some(self.values.get(index).is_some_and(|v| v.as_bytes() <= max)) - } - // A missing bound makes that end of the interval unbounded. - // Exclude only when the whole domain lies beyond the known bound; - // gaps within the domain and equality cannot prove disjointness. - (Some(min), None) - if self.values.last().is_some_and(|v| v.as_bytes() < min) => - { - Some(false) - } - (None, Some(max)) - if self.values.first().is_some_and(|v| v.as_bytes() > max) => - { - Some(false) - } - _ => None, + } + + /// Does the sorted, deduplicated domain hold `value`? + fn contains(&self, value: &[u8]) -> bool { + self.values + .binary_search_by(|candidate| { + self.membership.compare_bytes($bytes!(candidate), value) + }) + .is_ok() + } + } + + impl PartialEq for $name { + fn eq(&self, other: &Self) -> bool { + self.membership == other.membership + && self.min.eq(&other.min) + && self.max.eq(&other.max) + && self.values == other.values + } + } + + impl Hash for $name { + fn hash(&self, state: &mut H) { + self.membership.hash(state); + self.min.hash(state); + self.max.hash(state); + self.values.hash(state); + } + } + + impl Display for $name { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + let name = self.membership.display_name(); + write!( + f, + "{name}({}, {}, {} values)", + self.min, + self.max, + self.values.len() + ) + } + } + + impl PhysicalExpr for $name { + fn data_type(&self, _input_schema: &Schema) -> Result { + Ok(DataType::Boolean) + } + + fn nullable(&self, _input_schema: &Schema) -> Result { + Ok(true) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + // Normalize byte arrays and dictionary-encoded statistics to a view. + let min = self.min.evaluate(batch)?.into_array(batch.num_rows())?; + let max = self.max.evaluate(batch)?.into_array(batch.num_rows())?; + // A short slice can retain a buffer too large for view arrays' u32 + // offsets. Avoid a panic in the cast and keep pruning conservative. + if $has_oversized_buffer(min.as_ref(), u32::MAX as usize) + || $has_oversized_buffer(max.as_ref(), u32::MAX as usize) + { + return Ok(ColumnarValue::Array(Arc::new(BooleanArray::new_null( + batch.num_rows(), + )))); } - }) - .collect(); - Ok(ColumnarValue::Array(Arc::new(matches))) - } + // Unlike primitive casts, view casts can drop NULLs stored behind + // valid dictionary keys. Preserve that logical validity explicitly. + // TODO: Revisit this workaround once the Arrow dependency includes + // https://github.com/apache/arrow-rs/pull/10510. + let min_nulls = min.logical_nulls(); + let max_nulls = max.logical_nulls(); + let min = cast(&min, &$view_type)?; + let max = cast(&max, &$view_type)?; + let min = min.$as_view(); + let max = max.$as_view(); + let matches: BooleanArray = (0..batch.num_rows()) + .map(|i| { + let min: Option<&[u8]> = (min.is_valid(i) + && min_nulls.as_ref().is_none_or(|nulls| nulls.is_valid(i))) + .then(|| $bytes!(min.value(i))); + let max: Option<&[u8]> = (max.is_valid(i) + && max_nulls.as_ref().is_none_or(|nulls| nulls.is_valid(i))) + .then(|| $bytes!(max.value(i))); + if self.membership == SetMembership::NotIn { + // Membership is checked before bound equality. NOT IN ordering + // rejects different byte lengths first, avoiding scans of long + // common prefixes. + return not_in_may_match(min, max, |value| { + self.contains(value) + }); + } + match (min, max) { + (Some(min), Some(max)) => { + if min > max { + return None; + } + // String and binary values use unsigned lexicographic byte + // order, as required by the PruningStatistics min/max + // contract. Parquet adapters mask bounds with unusable + // ordering; PartitionPruningStatistics uses actual Arrow + // partition values. PrunableStatistics trusts file + // providers' bounds: there is no ordering gate for arbitrary + // statistics providers here. + let index = + self.values.partition_point(|v| $bytes!(v) < min); + Some( + self.values + .get(index) + .is_some_and(|v| $bytes!(v) <= max), + ) + } + // A missing bound makes that end of the interval unbounded. + // Exclude only when the whole domain lies beyond the known bound; + // gaps within the domain and equality cannot prove disjointness. + (Some(min), None) + if self + .values + .last() + .is_some_and(|v| $bytes!(v) < min) => + { + Some(false) + } + (None, Some(max)) + if self + .values + .first() + .is_some_and(|v| $bytes!(v) > max) => + { + Some(false) + } + _ => None, + } + }) + .collect(); + Ok(ColumnarValue::Array(Arc::new(matches))) + } - fn children(&self) -> Vec<&PhysicalExprRef> { - vec![&self.min, &self.max] - } + fn children(&self) -> Vec<&PhysicalExprRef> { + vec![&self.min, &self.max] + } - fn with_new_children( - self: Arc, - children: Vec, - ) -> Result { - assert_eq_or_internal_err!(children.len(), 2); - Ok(Arc::new(Self { - membership: self.membership, - min: Arc::clone(&children[0]), - max: Arc::clone(&children[1]), - values: Arc::clone(&self.values), - })) - } + fn with_new_children( + self: Arc, + children: Vec, + ) -> Result { + assert_eq_or_internal_err!(children.len(), 2); + Ok(Arc::new(Self { + membership: self.membership, + min: Arc::clone(&children[0]), + max: Arc::clone(&children[1]), + values: Arc::clone(&self.values), + })) + } - fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result { - write!(f, "{self}") - } + fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "{self}") + } + } + }; } +define_byte_in_list_expr!( + StringInListPruningExpr, + String, + DataType::Utf8View, + as_string_view, + string_bytes, + has_oversized_string_buffer +); + +define_byte_in_list_expr!( + BinaryInListPruningExpr, + Box<[u8]>, + DataType::BinaryView, + as_binary_view, + binary_bytes, + has_oversized_binary_buffer +); + #[cfg(test)] mod tests { use super::*; - use arrow::array::{ArrayRef, StringArray}; + use arrow::array::{ArrayRef, BinaryArray, StringArray}; - #[test] - fn oversized_buffers_check_retained_data_not_visible_offsets() -> Result<()> { - // Exercise the size boundary without allocating a 4 GiB buffer. - let limit = 32; - let padding = "p".repeat(limit - 1); - let array: ArrayRef = Arc::new(StringArray::from(vec!["a", padding.as_str()])); - - for value_type in [DataType::Utf8, DataType::LargeUtf8] { + fn assert_oversized_buffers( + array: &ArrayRef, + value_types: impl IntoIterator, + view_type: DataType, + limit: usize, + has_oversized_buffer: fn(&dyn Array, usize) -> bool, + ) -> Result<()> { + for value_type in value_types { for data_type in [ value_type.clone(), DataType::Dictionary( Box::new(DataType::Int32), Box::new(value_type.clone()), ), - DataType::Dictionary( - Box::new(DataType::UInt64), - Box::new(value_type.clone()), - ), + DataType::Dictionary(Box::new(DataType::UInt64), Box::new(value_type)), ] { - let slice = cast(&array, &data_type)?.slice(0, 1); - assert!(has_oversized_string_buffer(slice.as_ref(), limit - 1)); - assert!(has_oversized_string_buffer(slice.as_ref(), limit)); - assert!(!has_oversized_string_buffer(slice.as_ref(), limit + 1)); + let slice = cast(array, &data_type)?.slice(0, 1); + assert!(has_oversized_buffer(slice.as_ref(), limit - 1)); + assert!(has_oversized_buffer(slice.as_ref(), limit)); + assert!(!has_oversized_buffer(slice.as_ref(), limit + 1)); } } - // Already-normalized views do not have the byte-array cast limitation. for data_type in [ - DataType::Utf8View, - DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8View)), + view_type.clone(), + DataType::Dictionary(Box::new(DataType::Int32), Box::new(view_type)), ] { - let slice = cast(&array, &data_type)?.slice(0, 1); - assert!(!has_oversized_string_buffer(slice.as_ref(), limit)); + let slice = cast(array, &data_type)?.slice(0, 1); + assert!(!has_oversized_buffer(slice.as_ref(), limit)); } Ok(()) } + + #[test] + fn oversized_buffers_check_retained_data_not_visible_offsets() -> Result<()> { + // Exercise the size boundary without allocating a 4 GiB buffer. + let limit = 32; + let padding = "p".repeat(limit - 1); + let array: ArrayRef = Arc::new(StringArray::from(vec!["a", padding.as_str()])); + assert_oversized_buffers( + &array, + [DataType::Utf8, DataType::LargeUtf8], + DataType::Utf8View, + limit, + has_oversized_string_buffer, + )?; + + let padding = vec![0; limit - 1]; + let array: ArrayRef = Arc::new(BinaryArray::from(vec![&[1][..], &padding])); + assert_oversized_buffers( + &array, + [DataType::Binary, DataType::LargeBinary], + DataType::BinaryView, + limit, + has_oversized_binary_buffer, + )?; + Ok(()) + } } diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 1acdb2058d122..05305b73c637e 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -413,7 +413,7 @@ datafusion.execution.parquet.dictionary_page_size_limit 1048576 (writing) Sets b datafusion.execution.parquet.enable_page_index true (reading) If true, reads the Parquet data page level metadata (the Page Index), if present, to reduce the I/O and number of rows decoded. datafusion.execution.parquet.encoding NULL (writing) Sets default encoding for any column. Valid values are: plain, plain_dictionary, rle, bit_packed, delta_binary_packed, delta_length_byte_array, delta_byte_array, rle_dictionary, and byte_stream_split. These values are not case sensitive. If NULL, uses default parquet writer setting datafusion.execution.parquet.force_filter_selections false (reading) Force the use of RowSelections for filter results, when pushdown_filters is enabled. If false, the reader will automatically choose between a RowSelection and a Bitmap based on the number and pattern of selected rows. -datafusion.execution.parquet.max_in_list_size 20 Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger literal string lists on a string column use a compact representation, for both `IN` and `NOT IN`, including lists with NULL members. `NOT IN` with NULL and all-NULL `IN` lists cannot match any rows. Other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. +datafusion.execution.parquet.max_in_list_size 20 Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger literal lists use a compact representation when the column type is string, variable-length binary, integer, decimal, date, time, timestamp, or duration. This applies to both `IN` and `NOT IN`, including lists with NULL members. `NOT IN` with NULL and all-NULL `IN` lists cannot match any rows. Compact lists containing NULL do not use the fully-matched-row-group optimization. Floating-point and other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. datafusion.execution.parquet.max_predicate_cache_size NULL (reading) The maximum predicate cache size, in bytes. When `pushdown_filters` is enabled, sets the maximum memory used to cache the results of predicate evaluation between filter evaluation and output generation. Decreasing this value will reduce memory usage, but may increase IO and CPU usage. None means use the default parquet reader setting. 0 means no caching. datafusion.execution.parquet.max_row_group_bytes NULL (writing) Target maximum size of each row group in bytes. When set, the writer flushes whenever either this limit or `max_row_group_size` is reached, whichever comes first. Useful for bounding writer memory on wide schemas where a row-count limit can map to very different byte sizes. Matches the behavior of `parquet.block.size` in parquet-mr. If `None` (the default), only the row-count limit applies. Currently only honored when `allow_single_file_parallelism` is `false`; by default the parallel file writer ignores this limit. datafusion.execution.parquet.max_row_group_size 1048576 (writing) Target maximum number of rows in each row group (defaults to 1M rows). Writing larger row groups requires more memory to write, but can get better compression and be faster to read. When `max_row_group_bytes` is also set, the writer flushes a row group when either limit is reached, whichever comes first. diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index bffc69172a201..549cca413487f 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -93,7 +93,7 @@ The following configuration settings are available: | datafusion.execution.parquet.coerce_int96_tz | NULL | (reading) Optional timezone applied to INT96 columns when `coerce_int96` is set. When `Some`, INT96 columns coerce to `Timestamp(, Some())` instead of the default `Timestamp(, None)`. Spark and other systems write INT96 values as UTC-adjusted instants, so callers that need the resulting Arrow type to be timezone-aware (e.g. for Spark `TimestampType` semantics) should set this to `"UTC"`. No effect when `coerce_int96` is `None`. | | datafusion.execution.parquet.bloom_filter_on_read | true | (reading) Use any available bloom filters when reading parquet files | | datafusion.execution.parquet.max_predicate_cache_size | NULL | (reading) The maximum predicate cache size, in bytes. When `pushdown_filters` is enabled, sets the maximum memory used to cache the results of predicate evaluation between filter evaluation and output generation. Decreasing this value will reduce memory usage, but may increase IO and CPU usage. None means use the default parquet reader setting. 0 means no caching. | -| datafusion.execution.parquet.max_in_list_size | 20 | Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger literal string lists on a string column use a compact representation, for both `IN` and `NOT IN`, including lists with NULL members. `NOT IN` with NULL and all-NULL `IN` lists cannot match any rows. Other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. | +| datafusion.execution.parquet.max_in_list_size | 20 | Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger literal lists use a compact representation when the column type is string, variable-length binary, integer, decimal, date, time, timestamp, or duration. This applies to both `IN` and `NOT IN`, including lists with NULL members. `NOT IN` with NULL and all-NULL `IN` lists cannot match any rows. Compact lists containing NULL do not use the fully-matched-row-group optimization. Floating-point and other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. | | datafusion.execution.parquet.data_pagesize_limit | 1048576 | (writing) Sets best effort maximum size of data page in bytes | | datafusion.execution.parquet.write_batch_size | 1024 | (writing) Sets write_batch_size in rows | | datafusion.execution.parquet.writer_version | 1.0 | (writing) Sets parquet writer version valid values are "1.0" and "2.0" | From 390e174d1b1bbd42749a603dee6f2702985d71a1 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Tue, 29 Sep 2026 12:29:30 +0000 Subject: [PATCH 13/19] datafusion-physical-plan: handles nulls correctly in aggregate dynamic filters (#25579) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The fix prevents one aggregate’s filter from discarding rows that another aggregate still needs in dynamic filters. ### Example: Query: `MIN(a), MAX(a), MIN(b), MAX(b)` ``` First batch: (1, NULL), (8, NULL) | v a has bounds [1, 8] b has no bounds yet | +---------+----------+ | | Before After | | a < 1 OR a > 8 filter = true | | v v Next batch: (2, 5), (3, 7) | | Both rows rejected Both rows accepted | | v v Result: Now b has bounds [5, 7] 1 8 NULL NULL | v Filtering can start: a < 1 OR a > 8 OR b < 5 OR b > 7 | v Correct result: 1 8 5 7 ``` The latter is correct 1. According to [Postgres](https://www.postgresql.org/docs/current/functions-aggregate.html): > - MIN: “Computes the minimum of the non-null input values.” > - MAX: “Computes the maximum of the non-null input values.” 2. According to [DataFusion](https://github.com/apache/datafusion/blob/6ff484a96d7eca4b1bf4a263d7190ada5af2343c/datafusion/sqllogictest/test_files/aggregate.slt#L5490-L5500) Executing `1, NULL, 3, 4, 5` expects `MIN = 1, MAX = 5` --------- Co-authored-by: Gabriel <45515538+gabotechs@users.noreply.github.com> Backport provenance: - Upstream PR: https://github.com/apache/datafusion/pull/25579 - Upstream commit: 0af68188912349e9c865548458b37b0e66b5c31d - Conflict resolution only supplies the missing test module/import context; the upstream logic, native test body and SQL assertions are preserved. - Original upstream tests are retained; fork-only adaptation, if any, follows separately. Signed-off-by: discord9 --- .../src/aggregates/aggregate_stream.rs | 87 +++++++++++++++++-- .../push_down_filter_regression.slt | 49 +++++++++++ 2 files changed, 130 insertions(+), 6 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs index e1e869105d8db..946cb70cdb1f9 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs @@ -77,7 +77,6 @@ struct AggregateStreamInner { } impl AggregateStreamInner { - // TODO: check if we get Null handling correct /// # Examples /// - Example 1 /// Accumulators: min(c1) @@ -105,11 +104,12 @@ impl AggregateStreamInner { Vec::with_capacity(filter_state.accumulator_dyn_filter_info.len()); for acc_info in &filter_state.accumulator_dyn_filter_info { - // Skip if we don't yet have a meaningful bound + // Every aggregate needs a bound before we can prune rows. Otherwise, + // another aggregate's predicate could discard its first non-NULL value. let bound = { let guard = acc_info.shared_bound.lock(); if (*guard).is_null() { - continue; + return Ok(lit(true)); } guard.clone() }; @@ -255,9 +255,9 @@ fn scalar_cmp_null_short_circuit( v1: &ScalarValue, v2: &ScalarValue, ) -> Option { - match (v1, v2) { - (ScalarValue::Null, ScalarValue::Null) => Some(ScalarValue::Null), - (ScalarValue::Null, other) | (other, ScalarValue::Null) => Some(other.clone()), + match (v1.is_null(), v2.is_null()) { + (true, _) => Some(v2.clone()), + (_, true) => Some(v1.clone()), _ => None, } } @@ -476,3 +476,78 @@ fn aggregate_batch( Ok(allocated) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::ExecutionPlan; + use crate::aggregates::PhysicalGroupBy; + use crate::common::collect as collect_stream; + use crate::expressions::col; + use crate::test::TestMemoryExec; + use arrow::array::Int32Array; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_functions_aggregate::min_max::{max_udaf, min_udaf}; + use datafusion_physical_expr::aggregate::AggregateExprBuilder; + + #[tokio::test] + async fn aggregate_dynamic_filter_recovers_from_null_bounds() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, true), + ])); + let first = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 8])), + Arc::new(Int32Array::from(vec![None, None])), + ], + )?; + let second = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![2, 3])), + Arc::new(Int32Array::from(vec![5, 7])), + ], + )?; + let input = TestMemoryExec::try_new_exec( + &[vec![first], vec![second]], + Arc::clone(&schema), + None, + )?; + let mut aggregates = Vec::new(); + for column in ["a", "b"] { + for function in [min_udaf(), max_udaf()] { + let alias = format!("{}({column})", function.name()); + aggregates.push(Arc::new( + AggregateExprBuilder::new(function, vec![col(column, &schema)?]) + .schema(Arc::clone(&schema)) + .alias(alias) + .build()?, + )); + } + } + let aggregate = AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::default(), + aggregates, + vec![None; 4], + input, + schema, + )?; + let filter = &aggregate.dynamic_filter.as_ref().unwrap().filter; + let context = Arc::new(TaskContext::default()); + + // Execute partitions sequentially to make filter publication deterministic. + collect_stream(aggregate.execute(0, Arc::clone(&context))?).await?; + assert_eq!(filter.current()?.to_string(), "true"); + + // Typed NULL bounds must recover so filtering resumes for every aggregate. + collect_stream(aggregate.execute(1, context)?).await?; + assert_eq!( + filter.current()?.to_string(), + "a@0 < 1 OR a@0 > 8 OR b@1 < 5 OR b@1 > 7" + ); + Ok(()) + } +} diff --git a/datafusion/sqllogictest/test_files/push_down_filter_regression.slt b/datafusion/sqllogictest/test_files/push_down_filter_regression.slt index 5a038427d0c65..525999fa411db 100644 --- a/datafusion/sqllogictest/test_files/push_down_filter_regression.slt +++ b/datafusion/sqllogictest/test_files/push_down_filter_regression.slt @@ -443,6 +443,55 @@ drop table agg_dyn_mixed; statement ok reset datafusion.execution.batch_size; +# --- one aggregate has a bound before another sees its first non-NULL value --- +# Both files establish MIN(a)=1 and MAX(a)=8 in their first row group. Until b +# has a non-NULL bound, filtering only on a would prune every remaining row and +# incorrectly leave MIN(b) and MAX(b) NULL, regardless of partition order. + +statement ok +set datafusion.execution.batch_size = 2; + +statement ok +COPY ( + SELECT * FROM (VALUES (1, NULL), (8, NULL), (2, 5), (3, 7)) AS v(a, b) +) TO 'test_files/scratch/push_down_filter_regression/agg_dyn_late_bound/file_0.parquet' +STORED AS PARQUET +OPTIONS ('format.max_row_group_size' '2'); + +statement ok +COPY ( + SELECT * FROM (VALUES (1, NULL), (8, NULL), (2, 5), (3, 7)) AS v(a, b) +) TO 'test_files/scratch/push_down_filter_regression/agg_dyn_late_bound/file_1.parquet' +STORED AS PARQUET +OPTIONS ('format.max_row_group_size' '2'); + +statement ok +CREATE EXTERNAL TABLE agg_dyn_late_bound (a INT, b INT) +STORED AS PARQUET +LOCATION 'test_files/scratch/push_down_filter_regression/agg_dyn_late_bound/'; + +statement ok +set datafusion.optimizer.enable_aggregate_dynamic_filter_pushdown = false; + +query IIII +SELECT MIN(a), MAX(a), MIN(b), MAX(b) FROM agg_dyn_late_bound; +---- +1 8 5 7 + +statement ok +set datafusion.optimizer.enable_aggregate_dynamic_filter_pushdown = true; + +query IIII +SELECT MIN(a), MAX(a), MIN(b), MAX(b) FROM agg_dyn_late_bound; +---- +1 8 5 7 + +statement ok +drop table agg_dyn_late_bound; + +statement ok +reset datafusion.execution.batch_size; + # --- all-NULLs input: filter should stay `true` (no meaningful bound) --- statement ok From 38ea24dc8fd6b6128a615d8fcb42b7824604dd27 Mon Sep 17 00:00:00 2001 From: Huaijin Date: Thu, 24 Sep 2026 12:05:17 +0000 Subject: [PATCH 14/19] feat: recognize more lossless casts for statistics and ordering (#25500) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Which issue does this PR close? Follow-up to #25227 and [its discussion](https://github.com/apache/datafusion/pull/25227#discussion_r3999343833). ## Rationale for this change Some lossless casts unnecessarily discard column statistics or fail to preserve strict ordering information. For example, UInt32 → Int64 currently requires both bounds to be exact to propagate statistics, even though every UInt32 value fits in Int64. Recognizing this conversion as globally safe preserves available statistics even when one bound is absent or bounds are inexact. ## What changes are included in this PR? - Recognize unsigned-to-signed widening casts: UInt8 → Int16 / Int32 / Int64, UInt16 → Int32 / Int64, and UInt32 → Int64. - Recognize Utf8 → Utf8View and Binary → LargeBinary / BinaryView. - Extend the existing `check_bigger_cast` / `is_bigger_cast` predicates to recognize these lossless conversions. ## What is the testing strategy for this PR? - Test values, nulls, integer extrema, and strict ordering properties. - Test empty values, Unicode, long strings, and non-UTF8 binary values. - Test statistics propagation with exact, inexact, and missing bounds. - Verify unsafe reverse conversions remain excluded. Validation on the latest revision: 25 relevant unit tests passed; one existing test is ignored. Formatting, all-target/all-feature Clippy, and the standard `rust_lint.sh` suite passed. ## Are there any user-facing changes? More lossless casts retain statistics and strict ordering information, enabling existing optimizer rules. Backport provenance: - Upstream PR: https://github.com/apache/datafusion/pull/25500 - Upstream commit: b172f73f054d9dd6d51663ff713a9747d4f245a3 - Retain the upstream matcher and its complete self-tests; only the absent fork example-doc context is adjusted. - Original upstream tests are retained; fork-only adaptation, if any, follows separately. Signed-off-by: discord9 --- .../physical-expr/src/expressions/cast.rs | 189 +++++++++++++++++- datafusion/physical-expr/src/projection.rs | 75 +++++++ 2 files changed, 257 insertions(+), 7 deletions(-) diff --git a/datafusion/physical-expr/src/expressions/cast.rs b/datafusion/physical-expr/src/expressions/cast.rs index d9ee0798433c1..7bb8ddd85408b 100644 --- a/datafusion/physical-expr/src/expressions/cast.rs +++ b/datafusion/physical-expr/src/expressions/cast.rs @@ -250,7 +250,7 @@ impl CastExpr { } /// Check if casting from the specified source type to the target type is a - /// widening cast (e.g. from `Int8` to `Int16`). + /// lossless and strictly order-preserving cast (e.g. from `UInt8` to `Int16`). pub fn check_bigger_cast(cast_type: &DataType, src: &DataType) -> bool { if cast_type.eq(src) { return true; @@ -260,12 +260,13 @@ impl CastExpr { (Int8, Int16 | Int32 | Int64) | (Int16, Int32 | Int64) | (Int32, Int64) - | (UInt8, UInt16 | UInt32 | UInt64) - | (UInt16, UInt32 | UInt64) - | (UInt32, UInt64) + | (UInt8, UInt16 | UInt32 | UInt64 | Int16 | Int32 | Int64) + | (UInt16, UInt32 | UInt64 | Int32 | Int64) + | (UInt32, UInt64 | Int64) | (Int8 | Int16 | UInt8 | UInt16, Float32) | (Int8 | Int16 | Int32 | UInt8 | UInt16 | UInt32, Float64) - | (Utf8, LargeUtf8) + | (Utf8, LargeUtf8 | Utf8View) + | (Binary, LargeBinary | BinaryView) ) } @@ -1480,10 +1481,184 @@ mod tests { assert!(!CastExpr::check_bigger_cast(&Float64, &Int64)); assert!(!CastExpr::check_bigger_cast(&Float64, &UInt64)); - // Signed <-> Unsigned conversions should return false (not order-preserving due to negative values) + // Signed-to-unsigned and unsigned-to-signed casts whose target cannot + // represent the entire source range are not lossless for all values. assert!(!CastExpr::check_bigger_cast(&UInt16, &Int8)); assert!(!CastExpr::check_bigger_cast(&UInt32, &Int16)); - assert!(!CastExpr::check_bigger_cast(&Int16, &UInt8)); + assert!(!CastExpr::check_bigger_cast(&Int8, &UInt8)); + assert!(!CastExpr::check_bigger_cast(&Int16, &UInt16)); + assert!(!CastExpr::check_bigger_cast(&Int32, &UInt32)); + assert!(!CastExpr::check_bigger_cast(&Int64, &UInt64)); + assert!(!CastExpr::check_bigger_cast(&Int8, &UInt16)); + } + #[test] + fn test_byte_representation_cast_preserves_values_and_ordering() -> Result<()> { + use arrow::array::{ + BinaryArray, BinaryViewArray, LargeBinaryArray, StringViewArray, + }; + use arrow::compute::SortOptions; + use datafusion_expr_common::sort_properties::SortProperties; + + // Cover nulls, empty values, inline/long views, Unicode, and non-UTF8 bytes. + let strings = vec![ + None, + Some(""), + Some("a"), + Some("a longer shared string"), + Some("a longer shared string"), + Some("🦀"), + ]; + let bytes: Vec> = vec![ + None, + Some(b""), + Some(b"\0"), + Some(b"a longer shared byte string"), + Some(b"a longer shared byte string"), + Some(b"\xff"), + ]; + let binary: ArrayRef = Arc::new(BinaryArray::from(bytes.clone())); + let cases: [(ArrayRef, ArrayRef); 3] = [ + ( + Arc::new(StringArray::from(strings.clone())), + Arc::new(StringViewArray::from(strings)), + ), + ( + Arc::clone(&binary), + Arc::new(LargeBinaryArray::from(bytes.clone())), + ), + (binary, Arc::new(BinaryViewArray::from(bytes))), + ]; + for (input, expected) in cases { + let schema = Arc::new(Schema::new(vec![Field::new( + "a", + input.data_type().clone(), + true, + )])); + let expr = + CastExpr::new(col("a", &schema)?, expected.data_type().clone(), None); + assert!(expr.is_bigger_cast(input.data_type())); + for descending in [false, true] { + for nulls_first in [false, true] { + let child = ExprProperties::new_unknown() + .with_range(Interval::make_unbounded(input.data_type())?) + .with_order(SortProperties::Ordered(SortOptions { + descending, + nulls_first, + })) + .with_strictly_order_preserving(true); + let properties = expr.get_properties(std::slice::from_ref(&child))?; + assert_eq!(properties.sort_properties, child.sort_properties); + assert!(properties.strictly_order_preserving); + assert_eq!(properties.range.data_type(), *expected.data_type()); + } + } + let batch = RecordBatch::try_new(schema, vec![input])?; + let actual = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + assert_eq!(actual.as_ref(), expected.as_ref()); + } + for (source, target) in [ + (Utf8View, Utf8), + (LargeBinary, Binary), + (BinaryView, Binary), + ] { + assert!(!CastExpr::check_bigger_cast(&target, &source)); + } + Ok(()) + } + + #[test] + fn test_unsigned_to_signed_cast_preserves_values_and_ordering() { + use arrow::array::{UInt8Array, UInt16Array, UInt32Array}; + use arrow::compute::SortOptions; + use datafusion_expr_common::sort_properties::SortProperties; + + let inputs: [(ArrayRef, Vec, i64); 3] = [ + ( + Arc::new(UInt8Array::from(vec![ + None, + Some(0), + Some(1), + Some(u8::MAX), + ])), + vec![Int16, Int32, Int64], + i64::from(u8::MAX), + ), + ( + Arc::new(UInt16Array::from(vec![ + None, + Some(0), + Some(1), + Some(u16::MAX), + ])), + vec![Int32, Int64], + i64::from(u16::MAX), + ), + ( + Arc::new(UInt32Array::from(vec![ + None, + Some(0), + Some(1), + Some(u32::MAX), + ])), + vec![Int64], + i64::from(u32::MAX), + ), + ]; + for (input, target_types, max) in inputs { + let schema = Arc::new(Schema::new(vec![Field::new( + "a", + input.data_type().clone(), + true, + )])); + let batch = + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::clone(&input)]) + .expect("valid input batch"); + for target_type in target_types { + let expr = + CastExpr::new(col("a", &schema).unwrap(), target_type.clone(), None); + let actual = expr + .evaluate(&batch) + .unwrap() + .into_array(batch.num_rows()) + .unwrap(); + for (index, value) in + [None, Some(0), Some(1), Some(max)].into_iter().enumerate() + { + let expected = match target_type { + Int16 => { + ScalarValue::Int16(value.map(|v| i16::try_from(v).unwrap())) + } + Int32 => { + ScalarValue::Int32(value.map(|v| i32::try_from(v).unwrap())) + } + Int64 => ScalarValue::Int64(value), + _ => unreachable!(), + }; + assert_eq!( + ScalarValue::try_from_array(&actual, index).unwrap(), + expected + ); + } + for descending in [false, true] { + for nulls_first in [false, true] { + let child = ExprProperties::new_unknown() + .with_range( + Interval::make_unbounded(input.data_type()).unwrap(), + ) + .with_order(SortProperties::Ordered(SortOptions { + descending, + nulls_first, + })) + .with_strictly_order_preserving(true); + let properties = + expr.get_properties(std::slice::from_ref(&child)).unwrap(); + assert_eq!(properties.sort_properties, child.sort_properties); + assert!(properties.strictly_order_preserving); + assert_eq!(properties.range.data_type(), target_type); + } + } + } + } } } diff --git a/datafusion/physical-expr/src/projection.rs b/datafusion/physical-expr/src/projection.rs index 0e8876f017379..412161143ae15 100644 --- a/datafusion/physical-expr/src/projection.rs +++ b/datafusion/physical-expr/src/projection.rs @@ -2490,6 +2490,81 @@ pub(crate) mod tests { Ok(()) } + #[test] + fn test_project_statistics_lossless_cast() { + use Precision::{Absent, Exact, Inexact}; + + for (lower, upper, targets) in [ + ( + ScalarValue::UInt8(Some(0)), + ScalarValue::UInt8(Some(u8::MAX)), + vec![DataType::Int16, DataType::Int32, DataType::Int64], + ), + ( + ScalarValue::UInt16(Some(0)), + ScalarValue::UInt16(Some(u16::MAX)), + vec![DataType::Int32, DataType::Int64], + ), + ( + ScalarValue::UInt32(Some(0)), + ScalarValue::UInt32(Some(u32::MAX)), + vec![DataType::Int64], + ), + ( + ScalarValue::Utf8(Some(String::new())), + ScalarValue::Utf8(Some("🦀".to_string())), + vec![DataType::Utf8View], + ), + ( + ScalarValue::Binary(Some(vec![])), + ScalarValue::Binary(Some(vec![0xff])), + vec![DataType::LargeBinary, DataType::BinaryView], + ), + ] { + for target in targets { + // A globally safe cast does not need two exact extrema. Preserve + // whichever bounds are available without upgrading their precision. + for (min_value, max_value) in [ + (Exact(lower.clone()), Exact(upper.clone())), + (Exact(lower.clone()), Absent), + (Absent, Exact(upper.clone())), + (Inexact(lower.clone()), Inexact(upper.clone())), + (Inexact(lower.clone()), Absent), + (Absent, Inexact(upper.clone())), + ] { + let input = ColumnStatistics { + min_value: min_value.clone(), + max_value: max_value.clone(), + null_count: Exact(1), + distinct_count: Inexact(3), + sum_value: Absent, + byte_size: Absent, + }; + let expr = CastExpr::new( + Arc::new(Column::new("a", 0)), + target.clone(), + None, + ); + let output = + project_column_statistics_through_expr(&expr, &[input], None); + assert_eq!( + output, + ColumnStatistics { + min_value: min_value.cast_to(&target).unwrap(), + max_value: max_value.cast_to(&target).unwrap(), + null_count: Exact(1), + distinct_count: Inexact(3), + sum_value: Absent, + byte_size: Absent, + }, + "{} -> {target}, {min_value:?}..{max_value:?}", + lower.data_type() + ); + } + } + } + } + #[test] fn test_column_indices_duplicates() -> Result<()> { // Test that duplicate column indices appear only once From 5d5cec02ac9fb85636417518f4d1ccd04bef3c7b Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 6 Nov 2026 21:23:14 +0800 Subject: [PATCH 15/19] test: adapt #25500 projection statistics test helper Use the older fork's two-argument project_column_statistics_through_expr helper in the upstream regression test. This commit contains only the test helper adaptation; upstream production logic and tests are in the previous commit. Upstream PR: https://github.com/apache/datafusion/pull/25500 Upstream commit: b172f73f054d9dd6d51663ff713a9747d4f245a3 Signed-off-by: discord9 --- datafusion/physical-expr/src/projection.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/datafusion/physical-expr/src/projection.rs b/datafusion/physical-expr/src/projection.rs index 412161143ae15..21e42903ada38 100644 --- a/datafusion/physical-expr/src/projection.rs +++ b/datafusion/physical-expr/src/projection.rs @@ -2545,8 +2545,7 @@ pub(crate) mod tests { target.clone(), None, ); - let output = - project_column_statistics_through_expr(&expr, &[input], None); + let output = project_column_statistics_through_expr(&expr, &[input]); assert_eq!( output, ColumnStatistics { From d787bc3f90ecff351820307747424704af8dba8c Mon Sep 17 00:00:00 2001 From: Marko Grujic Date: Fri, 2 Oct 2026 13:08:41 +0000 Subject: [PATCH 16/19] Prune hash join probe side via a build-side key range bitmap (#25602) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Which issue does this PR close? - Closes https://github.com/apache/datafusion/issues/25291. - Alternative to https://github.com/apache/datafusion/pull/25292. ## Rationale for this change Avoid scanning redundant files/row-groups/pages in the probe side of hash joins, based on the values dictated by the build side. ## What changes are included in this PR? - wiring up PushdownStrategy::Map/HashTableLookupExpr to carry a representation of the build side values from a hash join, through - `KeyRangeBitmap` implementation, which maps the build-side values array into a finite-sized bitmap bucket, and can answer probing questions for certain ranges - extend `build_predicate_expression` to build the associated pruning expression from `HashTableLookupExpr` using the new `KeyRangeBitmapPruningExpr`, which implements `PhysicalExpr` on top of `KeyRangeBitmap` - also extend `build_predicate_expression` so that it now pushes down pruning for `CaseExprs`, since that also unlocks the partitioned hash-join scenario this pr targets ## What is the testing strategy for this PR? A number of unit tests added, and one SLT added. Also verified on the query shape that motivated the original issue ```sql > copy (select i as k, random() as v from generate_series(0, 1999999) t(i)) to '/tmp/fact.parquet' stored as parquet options ('format.max_row_group_size' '1000'); +---------+ | count | +---------+ | 2000000 | +---------+ 1 row(s) fetched. Elapsed 0.093 seconds. > create external table fact stored as parquet location '/tmp/fact.parquet'; 0 row(s) fetched. Elapsed 0.008 seconds. > create table dim as select i as k from generate_series(0, 1999999) t(i) where i % 10000 < 200; 0 row(s) fetched. Elapsed 0.014 seconds. > select count(*), sum(v) from fact join dim on fact.k = dim.k; select count(*), sum(v) from fact join dim on fact.k = dim.k; select count(*), sum(v) from fact join dim on fact.k = dim.k; select count(*), sum(v) from fact join dim on fact.k = dim.k; +----------+--------------------+ | count(*) | sum(fact.v) | +----------+--------------------+ | 40000 | 20121.308621485692 | +----------+--------------------+ 1 row(s) fetched. Elapsed 0.031 seconds. +----------+--------------------+ | count(*) | sum(fact.v) | +----------+--------------------+ | 40000 | 20121.308621485692 | +----------+--------------------+ 1 row(s) fetched. Elapsed 0.017 seconds. +----------+-------------------+ | count(*) | sum(fact.v) | +----------+-------------------+ | 40000 | 20121.30862148569 | +----------+-------------------+ 1 row(s) fetched. Elapsed 0.011 seconds. +----------+-------------------+ | count(*) | sum(fact.v) | +----------+-------------------+ | 40000 | 20121.30862148569 | +----------+-------------------+ 1 row(s) fetched. Elapsed 0.011 seconds. > explain analyze select count(*), sum(v) from fact join dim on fact.k = dim.k; ... | | DataSourceExec: file_groups={12 groups: [[tmp/fact.parquet:0..1927333], [tmp/fact.parquet:1927333..3854666], [tmp/fact.parquet:3854666..5781999], [tmp/fact.parquet:5781999..7709332], [tmp/fact.parquet:7709332..9636665], ...]}, projection=[k, v], output_ordering=[k@0 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ k@0 >= 0 AND k@0 <= 1990199 AND hash_lookup ], dynamic_rg_pruning=eligible, pruning_predicate=k_null_count@1 != row_count@2 AND k_max@0 >= 0 AND k_null_count@1 != row_count@2 AND k_min@3 <= 1990199 AND k_null_count@1 != row_count@2 AND KEY_RANGE_BITMAP(k_min@3, k_max@0, 20000/1048576 buckets), required_guarantees=[], metrics=[output_rows=200.0 K, elapsed_compute=9.65ms, output_bytes=25.0 MB, output_batches=200, files_ranges_pruned_statistics=12 total → 12 matched, row_groups_pruned_statistics=2.00 K total → 200 matched, row_groups_pruned_bloom_filter=200 total → 200 matched, page_index_pages_pruned=200 total → 200 matched, page_index_rows_pruned=200.0 K total → 200.0 K matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, bytes_processed=22.1 MB, bytes_scanned=2.1 MB, file_open_errors=0, file_scan_errors=0, files_opened=12, files_processed=12, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=0, pushdown_rows_pruned=0, row_groups_pruned_dynamic_filter=0, predicate_cache_inner_records=0, predicate_cache_records=0, bloom_filter_eval_time=64.23µs, metadata_load_time=345.27µs, page_index_eval_time=1.33ms, row_pushdown_eval_time=36ns, statistics_eval_time=896.06µs, time_elapsed_opening=6.99ms, time_elapsed_processing=24.90ms, time_elapsed_scanning_total=224.99ms, time_elapsed_scanning_until_data=25.01ms, output_rows_skew=1.19%, scan_efficiency_ratio=9.7% (2.24 M/23.13 M)] | ... ``` ## Are there any user-facing changes? --------- Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Jay Zhan Backport provenance: - Upstream PR: https://github.com/apache/datafusion/pull/25602 - Upstream commit: 4bb9f8d3192f6c945f2bf933231380e60341af66 - Upstream additions/removals and self-tests retained; fork codec/import context is resolved separately. Signed-off-by: discord9 --- Cargo.lock | 1 + .../physical-plan/src/joins/hash_join/exec.rs | 115 ++++++- .../joins/hash_join/exec/prepared/tests.rs | 15 +- .../joins/hash_join/partitioned_hash_eval.rs | 65 +++- .../src/joins/hash_join/shared_bounds.rs | 21 +- .../src/joins/key_range_bitmap.rs | 238 +++++++++++++++ datafusion/physical-plan/src/joins/mod.rs | 1 + datafusion/proto/tests/cases/plans/exprs.rs | 1 + datafusion/pruning/Cargo.toml | 1 + .../pruning/src/key_range_bitmap_expr.rs | 103 +++++++ datafusion/pruning/src/lib.rs | 1 + datafusion/pruning/src/pruning_predicate.rs | 283 ++++++++++++++++-- .../test_files/dynamic_row_group_pruning.slt | 32 +- .../test_files/parquet_filter_pushdown.slt | 47 +++ 14 files changed, 876 insertions(+), 48 deletions(-) create mode 100644 datafusion/physical-plan/src/joins/key_range_bitmap.rs create mode 100644 datafusion/pruning/src/key_range_bitmap_expr.rs diff --git a/Cargo.lock b/Cargo.lock index 8d0821fb8a0e0..9f0fa41743af7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2600,6 +2600,7 @@ dependencies = [ "insta", "itertools 0.15.0", "log", + "rstest", ] [[package]] diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 2f5d8abf91269..738724ae903da 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -40,6 +40,7 @@ use crate::joins::hash_join::stream::{ BuildSide, BuildSideInitialState, HashJoinStream, HashJoinStreamState, }; use crate::joins::join_hash_map::{JoinHashMapU32, JoinHashMapU64}; +use crate::joins::key_range_bitmap::KeyRangeBitmap; use crate::joins::utils::{ OnceAsync, OnceFut, asymmetric_join_output_partitioning, reorder_output_after_swap, swap_join_projection, update_hash, @@ -2615,7 +2616,10 @@ async fn collect_left_input( let map = Arc::new(join_hash_map); - let membership = if num_rows == 0 { + // For an ordinary build, nothing reads the strategy unless this join's own + // dynamic filter accumulator exists. A prepared build is shared with joins + // decided later, so it must always be ready to serve one. + let membership = if num_rows == 0 || (!prepared && !should_compute_dynamic_filters) { PushdownStrategy::Empty } else { // If the build side is small enough we can use IN list pushdown. @@ -2625,19 +2629,45 @@ async fn collect_left_input( .iter() .map(|arr| arr.get_array_memory_size()) .sum::(); - if left_values.is_empty() - || left_values[0].is_empty() - || estimated_size > config.optimizer.hash_join_inlist_pushdown_max_size - || map.num_of_distinct_key() - > config + + let pushdown_inlist = !left_values.is_empty() + && !left_values[0].is_empty() + && estimated_size <= config.optimizer.hash_join_inlist_pushdown_max_size + && map.num_of_distinct_key() + <= config .optimizer - .hash_join_inlist_pushdown_max_distinct_values + .hash_join_inlist_pushdown_max_distinct_values; + + if pushdown_inlist + && let Some(in_list_values) = build_struct_inlist_values(&left_values)? { - PushdownStrategy::Map(Arc::clone(&map)) - } else if let Some(in_list_values) = build_struct_inlist_values(&left_values)? { PushdownStrategy::InList(in_list_values) } else { - PushdownStrategy::Map(Arc::clone(&map)) + // Past the InList threshold use a bucket bitmap for container pruning. + let pruning_bitmap = match (left_values.as_slice(), bounds.as_ref()) { + ([keys], Some(bounds)) if !keys.is_empty() => bounds + .get_column_bounds(0) + .and_then(|b| { + KeyRangeBitmap::try_new( + keys, + &b.min, + &b.max, + map.num_of_distinct_key(), + ) + }) + .map(Arc::new), + _ => None, + }; + // Held for the join's lifetime, so charge it like the maps; it is + // optional, so skip it rather than fail when the pool is full. + let pruning_bitmap = pruning_bitmap.filter(|bitmap| { + let ok = reservation.try_grow(bitmap.size()).is_ok(); + if ok { + metrics.build_mem_used.add(bitmap.size()); + } + ok + }); + PushdownStrategy::Map(Arc::clone(&map), pruning_bitmap) } }; @@ -2785,6 +2815,71 @@ mod tests { Ok(()) } + /// Runs a join whose 200 keys spread over 2M would size a pruning bitmap + /// at the 128 KiB cap, and reports the bytes the build side reserved. + /// + /// Forces the IN-list threshold off: under `force_hash_collisions` every + /// key hashes to the same bucket, so `num_of_distinct_key()` no longer + /// reflects the real 200 distinct keys and would otherwise take the + /// IN-list branch instead of the one under test. + async fn build_mem_used(limit: usize, dynamic_filters: bool) -> Result { + let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int64, false)])); + let batch = |keys: Vec| { + let column = Arc::new(Int64Array::from(keys)) as ArrayRef; + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![column])?; + TestMemoryExec::try_new_exec(&[vec![batch]], Arc::clone(&schema), None) + }; + let on = vec![( + Arc::new(Column::new_with_schema("k", &schema)?) as _, + Arc::new(Column::new_with_schema("k", &schema)?) as _, + )]; + let build = batch((0..200).map(|i| i * 10_000).collect())?; + let probe = batch(vec![0, 500_000, 1_500_000])?; + let join = if dynamic_filters { + hash_join_with_dynamic_filter(build, probe, on, JoinType::Inner)?.0 + } else { + join( + build, + probe, + on, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + )? + }; + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(limit, 1.0) + .build_arc()?; + let mut config = SessionConfig::new(); + config + .options_mut() + .optimizer + .hash_join_inlist_pushdown_max_size = 0; + let task_ctx = Arc::new( + TaskContext::default() + .with_runtime(runtime) + .with_session_config(config), + ); + let batches = common::collect(join.execute(0, task_ctx)?).await?; + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 3); + Ok(join + .metrics() + .unwrap() + .sum_by_name("build_mem_used") + .unwrap() + .as_usize()) + } + + /// The bitmap is pruning-only: never built when no dynamic filter will read + /// it, and dropped rather than fatal when the pool cannot fit it. + #[tokio::test] + async fn pruning_bitmap_is_optional() -> Result<()> { + assert!(build_mem_used(1_000_000, false).await? < 100_000); + assert!(build_mem_used(1_000_000, true).await? > 100_000); + build_mem_used(100_000, true).await?; + Ok(()) + } + #[derive(Debug)] struct PartitionedTestExec { cache: Arc, diff --git a/datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs b/datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs index e83f953a4c9f0..48281e5f481b1 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs @@ -785,7 +785,9 @@ async fn prepared_inlist_shares_accounted_build_buffers() -> Result<()> { assert_eq!(counter.count_batch(&membership_batch), 0); drop(membership_batch); - // Disabling IN-list membership changes no retained key buffers. + // Disabling IN-list membership shares the same key buffers; a + // pruning bitmap, when the sparse two-key range earns one, adds + // its own bytes on top. let mut config = ConfigOptions::default(); config.optimizer.hash_join_inlist_pushdown_max_size = 0; let map = base @@ -795,9 +797,12 @@ async fn prepared_inlist_shares_accounted_build_buffers() -> Result<()> { Arc::new(config), ) .await?; - assert!(matches!(map.build.membership, PushdownStrategy::Map(_))); - assert_eq!(map.reserved_bytes(), bytes); - assert_eq!(pool.reserved(), 2 * bytes); + let PushdownStrategy::Map(_, bitmap) = &map.build.membership else { + panic!("expected hash-map membership"); + }; + let bitmap_bytes = bitmap.as_ref().map_or(0, |bitmap| bitmap.size()); + assert_eq!(map.reserved_bytes(), bytes + bitmap_bytes); + assert_eq!(pool.reserved(), 2 * bytes + bitmap_bytes); drop(map); let lease = Arc::clone(&prepared); drop(prepared); @@ -888,7 +893,7 @@ async fn prepared_byte_keys_use_hash_membership() -> Result<()> { assert!(prepared.build.bounds.is_none()); assert!(matches!( &prepared.build.membership, - PushdownStrategy::Map(_) + PushdownStrategy::Map(_, _) )); let task = with_probe_filter(base)? .builder() diff --git a/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs b/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs index fdd66d0e54ae2..fca91ccfd975b 100644 --- a/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs +++ b/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs @@ -35,6 +35,7 @@ use datafusion_physical_expr_common::physical_expr::{ }; use crate::joins::Map; +use crate::joins::key_range_bitmap::KeyRangeBitmap; /// RandomState wrapper that preserves the seed used to create it. /// @@ -265,6 +266,8 @@ pub struct HashTableLookupExpr { map: Arc, /// Description for display description: String, + /// Shared pruning-only view of the build side's key range + pruning_bitmap: Option>, } impl HashTableLookupExpr { /// Create a new HashTableLookupExpr @@ -274,6 +277,7 @@ impl HashTableLookupExpr { /// * `random_state` - SeededRandomState for hashing /// * `map` - Map to check membership (hash table or array map) /// * `description` - Description for debugging + /// * `pruning_bitmap` - key-range summary for pruning only, or `None` /// # Note /// This is public for internal testing purposes only and is not /// guaranteed to be stable across versions. @@ -282,14 +286,21 @@ impl HashTableLookupExpr { random_state: SeededRandomState, map: Arc, description: String, + pruning_bitmap: Option>, ) -> Self { Self { on_columns, random_state, map, description, + pruning_bitmap, } } + + /// Which parts of this build side's key range hold keys, for container pruning. + pub fn pruning_bitmap(&self) -> Option<&Arc> { + self.pruning_bitmap.as_ref() + } } impl std::fmt::Debug for HashTableLookupExpr { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { @@ -342,12 +353,13 @@ impl PhysicalExpr for HashTableLookupExpr { self: Arc, children: Vec>, ) -> Result> { - Ok(Arc::new(HashTableLookupExpr::new( - children, - self.random_state.clone(), - Arc::clone(&self.map), - self.description.clone(), - ))) + Ok(Arc::new(HashTableLookupExpr { + on_columns: children, + random_state: self.random_state.clone(), + map: Arc::clone(&self.map), + description: self.description.clone(), + pruning_bitmap: self.pruning_bitmap.clone(), + })) } fn data_type(&self, _input_schema: &Schema) -> Result { @@ -427,6 +439,7 @@ fn evaluate_columns( mod tests { use super::*; use crate::joins::join_hash_map::JoinHashMapU32; + use datafusion_common::ScalarValue; use datafusion_physical_expr::expressions::Column; use std::collections::hash_map::DefaultHasher; use std::hash::Hasher; @@ -437,6 +450,36 @@ mod tests { hasher.finish() } + /// The pruning bitmap must survive `with_new_children`, which a scan calls + /// once per file. + #[test] + fn test_pruning_bitmap_shared_with_derived_children() { + let keys: ArrayRef = Arc::new(arrow::array::Int64Array::from(vec![0, 1_000_000])); + let bitmap = Arc::new( + KeyRangeBitmap::try_new( + &keys, + &ScalarValue::Int64(Some(0)), + &ScalarValue::Int64(Some(1_000_000)), + 2, + ) + .expect("two keys a million apart leave prunable gaps"), + ); + let expr = Arc::new(HashTableLookupExpr::new( + vec![Arc::new(Column::new("a", 0))], + SeededRandomState::with_seed(1), + Arc::new(Map::HashMap(Box::new(JoinHashMapU32::with_capacity(10)))), + "hash_lookup".to_string(), + Some(Arc::clone(&bitmap)), + )); + + let derived = expr + .with_new_children(vec![Arc::new(Column::new("a", 7))]) + .unwrap(); + let derived = derived.downcast_ref::().unwrap(); + assert_eq!(derived.children()[0].to_string(), "a@7"); + assert!(Arc::ptr_eq(&bitmap, derived.pruning_bitmap().unwrap())); + } + #[test] fn test_hash_expr_eq_same() { let col_a: PhysicalExprRef = Arc::new(Column::new("a", 0)); @@ -715,6 +758,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup".to_string(), + None, ); let expr2 = HashTableLookupExpr::new( @@ -722,6 +766,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup".to_string(), + None, ); assert_eq!(expr1, expr2); @@ -740,6 +785,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup".to_string(), + None, ); let expr2 = HashTableLookupExpr::new( @@ -747,6 +793,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup".to_string(), + None, ); assert_ne!(expr1, expr2); @@ -763,6 +810,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup_one".to_string(), + None, ); let expr2 = HashTableLookupExpr::new( @@ -770,6 +818,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup_two".to_string(), + None, ); assert_ne!(expr1, expr2); @@ -789,6 +838,7 @@ mod tests { SeededRandomState::with_seed(1), hash_map1, "lookup".to_string(), + None, ); let expr2 = HashTableLookupExpr::new( @@ -796,6 +846,7 @@ mod tests { SeededRandomState::with_seed(1), hash_map2, "lookup".to_string(), + None, ); // Different Arc pointers means not equal (uses Arc::ptr_eq) @@ -813,6 +864,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup".to_string(), + None, ); let expr2 = HashTableLookupExpr::new( @@ -820,6 +872,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup".to_string(), + None, ); // Equal expressions should have equal hashes diff --git a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs index 94ec4565a4cef..9182fc4173dea 100644 --- a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs +++ b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs @@ -31,6 +31,7 @@ use crate::joins::hash_join::inlist_builder::build_struct_fields; use crate::joins::hash_join::partitioned_hash_eval::{ HashExpr, HashTableLookupExpr, SeededRandomState, }; +use crate::joins::key_range_bitmap::KeyRangeBitmap; use crate::repartition::RangeExpr; use arrow::array::ArrayRef; use arrow::datatypes::{DataType, Field, Schema}; @@ -136,12 +137,15 @@ fn create_membership_predicate( )?))) } // Use hash table lookup for large build sides - PushdownStrategy::Map(hash_map) => Ok(Some(Arc::new(HashTableLookupExpr::new( - on_right.to_vec(), - random_state.clone(), - hash_map, - "hash_lookup".to_string(), - )) as Arc)), + PushdownStrategy::Map(hash_map, pruning_bitmap) => { + Ok(Some(Arc::new(HashTableLookupExpr::new( + on_right.to_vec(), + random_state.clone(), + hash_map, + "hash_lookup".to_string(), + pruning_bitmap, + )) as Arc)) + } // Empty partition - should not create a filter for this PushdownStrategy::Empty => Ok(None), } @@ -277,8 +281,9 @@ pub(crate) struct SharedBuildAccumulator { pub(crate) enum PushdownStrategy { /// Use InList for small build sides (< 128MB) InList(ArrayRef), - /// Use map lookup for large build sides - Map(Arc), + /// Use map lookup for large build sides. If `Some` the second field + /// represents the build side's keys using a bucket bitmap. + Map(Arc, Option>), /// There was no data in this partition, do not build a dynamic filter for it Empty, } diff --git a/datafusion/physical-plan/src/joins/key_range_bitmap.rs b/datafusion/physical-plan/src/joins/key_range_bitmap.rs new file mode 100644 index 0000000000000..efdb5467c897e --- /dev/null +++ b/datafusion/physical-plan/src/joins/key_range_bitmap.rs @@ -0,0 +1,238 @@ +// 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. + +//! Bitmap over a hash join build side's key range, for container pruning. + +use arrow::array::{Array, ArrayRef, BooleanArray, downcast_integer_array}; +use arrow::buffer::{BooleanBuffer, MutableBuffer}; +use arrow::util::bit_util; +use datafusion_common::ScalarValue; +use num_traits::AsPrimitive; + +use crate::joins::array_map::ArrayMap; + +/// Largest bitmap, in bits (128 KiB), which bounds what one build side can hold. +const MAX_BUCKETS: u64 = 1 << 20; + +/// Which parts of a build side's key range hold at least one key. +/// +/// Keys map order-preservingly onto buckets spanning `[min, max]`, one bit each, +/// so "can a container holding `[lo, hi]` match?" is a scan of the bits spanning +/// `[lo, hi]`: all clear proves no build key lies in it. +/// +/// Below, keys 3 to 97 over eight buckets: `offset` is 3, `span` is 97 - 3, and +/// eight buckets need `shift` = 4, so each covers `1 << 4` keys starting at +/// `offset`. Only buckets 0 to `span >> shift` can hold a key. Real bitmaps run +/// to `MAX_BUCKETS`; eight is what fits on a line. +/// +/// ```text +/// keys 3 12 42 97 +/// v v v v +/// bucket | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | +/// covers 3-18 19-34 35-50 51-66 67-82 83-98 unused unused +/// bit 1 0 1 0 0 1 0 0 +/// ``` +/// +/// A container is kept when every bucket its `[min, max]` touches has a bit set. +#[derive(Debug)] +pub struct KeyRangeBitmap { + /// Smallest key, in the `u64` ordering the join's array map uses. + offset: u64, + /// `max - min`: a key is in range iff `key.wrapping_sub(offset) <= span`. + span: u64, + /// Each bucket covers `1 << shift` keys. + shift: u32, + /// One bit per bucket. + bits: BooleanBuffer, +} + +impl KeyRangeBitmap { + /// Maps `array`'s non-null values, which must lie within `[min, max]`. + /// + /// `None` when the key type has no `u64` ordering, or when the bitmap could + /// not exclude anything the `[min, max]` bounds predicate does not already. + pub fn try_new( + array: &ArrayRef, + min: &ScalarValue, + max: &ScalarValue, + distinct_keys: usize, + ) -> Option { + if !ArrayMap::is_supported_type(array.data_type()) { + return None; + } + let offset = ArrayMap::key_to_u64(min)?; + let span = ArrayMap::key_to_u64(max)?.wrapping_sub(offset); + + // One bucket per key value where the range allows it, so the bitmap is + // exact; otherwise as many as the cap permits. + let buckets = span.saturating_add(1).min(MAX_BUCKETS).next_power_of_two(); + + // Use a heuristic to bail-out for the common shape which renders each bucket + // set (contiguous surrogate keys, a date range) before actually computing it. + if (distinct_keys as u64).saturating_mul(2) > span.min(buckets) { + return None; + } + + let shift = (0u32..63).find(|s| (span >> s) < buckets).unwrap_or(63); + + let bits = bucket_bits(array, offset, shift, buckets); + + // Every bucket set means nothing can be excluded. + let reachable = ((span >> shift) + 1) as usize; + bits.slice(0, reachable).has_false().then_some(Self { + offset, + span, + shift, + bits, + }) + } + + /// One Boolean per `[min[i], max[i]]` container interval: `false` proves no + /// build key lies in that container, `None` where a bound is absent or the + /// statistics are not integer keys. + pub fn may_contain_ranges(&self, min: &dyn Array, max: &dyn Array) -> BooleanArray { + // Both bounds must be the same integer type for the tuple pattern to match. + downcast_integer_array!( + (min, max) => { + min.iter() + .zip(max.iter()) + .map(|bounds| match bounds { + (Some(lo), Some(hi)) => { + Some(self.may_contain_range(lo.as_(), hi.as_())) + } + _ => None, + }) + .collect() + } + _ => BooleanArray::new_null(min.len()), + ) + } + + /// Might a container spanning `[lo, hi]` hold a build key? `false` proves it + /// cannot. Bounds use the same `u64` key ordering as the join's array map. + pub fn may_contain_range(&self, lo: u64, hi: u64) -> bool { + // Offsets from the minimum; anything outside `[min, max]` wraps to + // something huge, which the clamping below turns into "past the end". + let lo_off = lo.wrapping_sub(self.offset); + let hi_off = hi.wrapping_sub(self.offset); + if lo_off > self.span && hi_off > self.span && lo_off <= hi_off { + return false; + } + let lo_b = if lo_off > self.span { + 0 + } else { + lo_off >> self.shift + }; + let hi_b = (hi_off.min(self.span) >> self.shift).max(lo_b); + self.bits + .slice(lo_b as usize, (hi_b - lo_b + 1) as usize) + .has_true() + } + + /// Bytes held, for the build-side memory reservation. + pub fn size(&self) -> usize { + self.bits.inner().capacity() + } + + /// Buckets set and buckets total, for `EXPLAIN`. + pub fn fill_stats(&self) -> (usize, usize) { + (self.bits.count_set_bits(), self.bits.len()) + } +} + +/// One bit per bucket, set for each bucket holding a non-null key of `keys`. +fn bucket_bits(keys: &dyn Array, offset: u64, shift: u32, buckets: u64) -> BooleanBuffer { + let mut words = MutableBuffer::new_null(buckets as usize); + downcast_integer_array!( + keys => { + let bytes = words.as_slice_mut(); + let mut set = |key: u64| { + let bucket = (key.wrapping_sub(offset) >> shift).min(buckets - 1); + bit_util::set_bit(bytes, bucket as usize); + }; + if keys.null_count() == 0 { + keys.values().iter().for_each(|key| set((*key).as_())); + } else { + keys.iter().flatten().for_each(|key| set(key.as_())); + } + } + _ => unreachable!("guarded by ArrayMap::is_supported_type"), + ); + BooleanBuffer::new(words.into(), 0, buckets as usize) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Int64Array; + use std::sync::Arc; + + fn bitmap(keys: &[i64]) -> Option { + let a: ArrayRef = Arc::new(Int64Array::from(keys.to_vec())); + KeyRangeBitmap::try_new( + &a, + &ScalarValue::Int64(Some(*keys.iter().min().unwrap())), + &ScalarValue::Int64(Some(*keys.iter().max().unwrap())), + keys.iter().collect::>().len(), + ) + } + fn may(b: &KeyRangeBitmap, lo: i64, hi: i64) -> bool { + b.may_contain_range(lo as u64, hi as u64) + } + + #[test] + fn prunes_gaps_and_never_drops_a_match() { + let keys: Vec = (0..2_000_000).filter(|k| k % 10000 < 200).collect(); + let s = bitmap(&keys).expect("not saturated"); + let mut kept = 0; + for i in 0..2000i64 { + let (lo, hi) = (1000 * i, 1000 * i + 999); + let truth = keys + .binary_search(&lo) + .map_or_else(|p| keys.get(p).is_some_and(|k| *k <= hi), |_| true); + let got = may(&s, lo, hi); + assert!( + !truth || got, + "dropped a container holding a key: {lo}..{hi}" + ); + kept += got as usize; + } + assert!(kept <= 220, "expected ~200 kept, got {kept}"); + } + + #[test] + fn contiguous_keys_no_bitmap() { + assert!(bitmap(&(0..5000i64).collect::>()).is_none()); + } + + #[test] + fn test_bitmap_range_probing() { + let s = bitmap(&[0, 1_000_000]).expect("not saturated"); + assert!(!may(&s, -500, -10)); + assert!(!may(&s, 2_000_000, 3_000_000)); + assert!(may(&s, -500, 10)); + + let s = bitmap(&[-1_000_000, -5, 0, 5, 1_000_000]).expect("not saturated"); + assert!(may(&s, -10, 10)); + assert!(!may(&s, -900_000, -800_000)); + + // The widest possible span must not overflow the bucket sizing. + let s = bitmap(&[i64::MIN, 0, i64::MAX]).expect("not saturated"); + assert!(may(&s, -1, 1)); + assert!(!may(&s, 1 << 50, 1 << 51)); + } +} diff --git a/datafusion/physical-plan/src/joins/mod.rs b/datafusion/physical-plan/src/joins/mod.rs index 39b80cd3d0cc0..f8e9788bc3356 100644 --- a/datafusion/physical-plan/src/joins/mod.rs +++ b/datafusion/physical-plan/src/joins/mod.rs @@ -49,6 +49,7 @@ mod join_filter; /// Note: This module is public for internal testing purposes only /// and is not guaranteed to be stable across versions. pub mod join_hash_map; +pub mod key_range_bitmap; use array_map::ArrayMap; use utils::JoinHashMapType; diff --git a/datafusion/proto/tests/cases/plans/exprs.rs b/datafusion/proto/tests/cases/plans/exprs.rs index 518b4a62ce072..cc96f2aeb3f77 100644 --- a/datafusion/proto/tests/cases/plans/exprs.rs +++ b/datafusion/proto/tests/cases/plans/exprs.rs @@ -117,6 +117,7 @@ fn roundtrip_hash_table_lookup_expr_to_lit() -> Result<()> { datafusion::physical_plan::joins::SeededRandomState::with_seed(0), hash_map, "test_lookup".to_string(), + None, )); // Create a filter with the lookup expression diff --git a/datafusion/pruning/Cargo.toml b/datafusion/pruning/Cargo.toml index a914a0a079bcd..48b1ed009b28c 100644 --- a/datafusion/pruning/Cargo.toml +++ b/datafusion/pruning/Cargo.toml @@ -31,6 +31,7 @@ datafusion-expr = { workspace = true } datafusion-functions-nested = { workspace = true } insta = { workspace = true } itertools = { workspace = true } +rstest = { workspace = true } [[bench]] harness = false diff --git a/datafusion/pruning/src/key_range_bitmap_expr.rs b/datafusion/pruning/src/key_range_bitmap_expr.rs new file mode 100644 index 0000000000000..ba30a052cac10 --- /dev/null +++ b/datafusion/pruning/src/key_range_bitmap_expr.rs @@ -0,0 +1,103 @@ +// 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 std::fmt::{self, Display, Formatter}; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +use arrow::datatypes::{DataType, Schema}; +use arrow::record_batch::RecordBatch; +use datafusion_common::{Result, assert_eq_or_internal_err}; +use datafusion_physical_expr::{PhysicalExpr, PhysicalExprRef}; +use datafusion_physical_plan::ColumnarValue; +use datafusion_physical_plan::joins::key_range_bitmap::KeyRangeBitmap; + +/// Tests containers' `[min, max]` statistics against a build side's +/// [`KeyRangeBitmap`]: one nullable Boolean per container, where `false` proves +/// no build key lies in it and `NULL` means the bounds could not decide. +#[derive(Debug)] +pub(crate) struct KeyRangeBitmapPruningExpr { + pub(crate) min: PhysicalExprRef, + pub(crate) max: PhysicalExprRef, + pub(crate) bitmap: Arc, +} + +impl PartialEq for KeyRangeBitmapPruningExpr { + fn eq(&self, other: &Self) -> bool { + self.min.eq(&other.min) + && self.max.eq(&other.max) + && Arc::ptr_eq(&self.bitmap, &other.bitmap) + } +} +impl Eq for KeyRangeBitmapPruningExpr {} + +impl Hash for KeyRangeBitmapPruningExpr { + fn hash(&self, state: &mut H) { + self.min.hash(state); + self.max.hash(state); + Arc::as_ptr(&self.bitmap).hash(state); + } +} + +impl Display for KeyRangeBitmapPruningExpr { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + let (set, total) = self.bitmap.fill_stats(); + write!( + f, + "KEY_RANGE_BITMAP({}, {}, {set}/{total} buckets)", + self.min, self.max + ) + } +} + +impl PhysicalExpr for KeyRangeBitmapPruningExpr { + fn data_type(&self, _input_schema: &Schema) -> Result { + Ok(DataType::Boolean) + } + + fn nullable(&self, _input_schema: &Schema) -> Result { + Ok(true) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + let rows = batch.num_rows(); + let min = self.min.evaluate(batch)?.into_array(rows)?; + let max = self.max.evaluate(batch)?.into_array(rows)?; + let matches = self.bitmap.may_contain_ranges(&min, &max); + Ok(ColumnarValue::Array(Arc::new(matches))) + } + + fn children(&self) -> Vec<&PhysicalExprRef> { + vec![&self.min, &self.max] + } + + fn with_new_children( + self: Arc, + children: Vec, + ) -> Result { + assert_eq_or_internal_err!(children.len(), 2); + Ok(Arc::new(Self { + min: Arc::clone(&children[0]), + max: Arc::clone(&children[1]), + bitmap: Arc::clone(&self.bitmap), + })) + } + + fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "{self}") + } +} diff --git a/datafusion/pruning/src/lib.rs b/datafusion/pruning/src/lib.rs index 88a7fdda5e733..3e7c2ad7aa7ab 100644 --- a/datafusion/pruning/src/lib.rs +++ b/datafusion/pruning/src/lib.rs @@ -19,6 +19,7 @@ mod file_pruner; mod in_list; +mod key_range_bitmap_expr; mod primitive_in_list; mod pruning_predicate; mod string_in_list; diff --git a/datafusion/pruning/src/pruning_predicate.rs b/datafusion/pruning/src/pruning_predicate.rs index 2118892d089d2..65fb737bfbe14 100644 --- a/datafusion/pruning/src/pruning_predicate.rs +++ b/datafusion/pruning/src/pruning_predicate.rs @@ -23,6 +23,7 @@ use std::collections::HashSet; use std::sync::Arc; use crate::in_list::{SetMembership, unwrap_scalar}; +use crate::key_range_bitmap_expr::KeyRangeBitmapPruningExpr; use crate::primitive_in_list::PrimitiveInListDomain; use crate::string_in_list::{BinaryInListPruningExpr, StringInListPruningExpr}; @@ -52,6 +53,7 @@ use datafusion_expr_common::operator::Operator; use datafusion_physical_expr::utils::{Guarantee, LiteralGuarantee}; use datafusion_physical_expr::{PhysicalExprRef, expressions as phys_expr}; use datafusion_physical_expr_common::physical_expr::snapshot_physical_expr_opt; +use datafusion_physical_plan::joins::HashTableLookupExpr; use datafusion_physical_plan::{ColumnarValue, PhysicalExpr}; /// Used to prove that arbitrary predicates (boolean expression) can not @@ -1525,6 +1527,78 @@ impl CompactInListDomain { } } +/// Statistics expressions for `column`: its min, its max, and "not every row is +/// NULL". Appends to `required_columns`, rolling the appends back if any part +/// cannot be built, since [`RequiredColumns::stat_column_expr`] only appends. +fn min_max_non_null_exprs( + column: &phys_expr::Column, + column_expr: &Arc, + field: &Field, + schema: &Schema, + required_columns: &mut RequiredColumns, +) -> Option<(PhysicalExprRef, PhysicalExprRef, PhysicalExprRef)> { + let appended = required_columns.columns.len(); + let statistics = (|| { + let min = required_columns + .min_column_expr(column, column_expr, field) + .ok()?; + let max = required_columns + .max_column_expr(column, column_expr, field) + .ok()?; + // A cast on `column_expr` (widening it towards min/max's target type) + // does not affect nullability, so check it on the bare column beneath. + let bare_expr = match column_expr.downcast_ref::() { + Some(cast) => cast.expr(), + None => column_expr, + }; + let non_null = + build_is_null_column_expr(bare_expr, schema, required_columns, true)?; + Some((min, max, non_null)) + })(); + if statistics.is_none() { + required_columns.columns.truncate(appended); + } + statistics +} + +/// Builds a "may match" expression for a [`HashTableLookupExpr`] (a large join +/// build side pushed down as an opaque hash-table lookup): tests container +/// min/max stats against the build side's key-range bitmap, which is fixed size +/// and built once with the join's hash table. +fn build_hash_lookup_pruning_expr( + lookup: &HashTableLookupExpr, + schema: &Schema, + required_columns: &mut RequiredColumns, +) -> Option> { + let bitmap = Arc::clone(lookup.pruning_bitmap()?); + let on_columns = lookup.children(); + let [column_expr] = on_columns[..] else { + return None; + }; + // The schema adapter often wraps the column in a same-type or widening cast + // (a nullable table column over a REQUIRED file column, or a narrower file + // type). `stat_column_expr` rewrites only the inner `Column`, keeping it. + let column = match column_expr.downcast_ref::() { + Some(cast) => { + let column = cast.expr().downcast_ref::()?; + let from = schema.fields().get(column.index())?.data_type(); + cast.is_bigger_cast(from).then_some(column)? + } + None => column_expr.downcast_ref::()?, + }; + let field = schema.fields().get(column.index())?; + if field.name() != column.name() { + return None; + } + let (min, max, non_null) = + min_max_non_null_exprs(column, column_expr, field, schema, required_columns)?; + Some(Arc::new(phys_expr::BinaryExpr::new( + non_null, + Operator::And, + Arc::new(KeyRangeBitmapPruningExpr { min, max, bitmap }), + ))) +} + /// Keep large literal lists of supported ordered types compact instead of /// building a per-value tree: an OR tree for `IN`, an AND chain for `NOT IN`. /// @@ -1601,24 +1675,8 @@ fn build_compact_in_list_expr( Some(false), )))); } - // Roll back appended statistics columns if the compact rewrite cannot be - // completed. `RequiredColumns::stat_column_expr` only appends entries. - let required_columns_len = required_columns.columns.len(); - let statistics = (|| { - let min = required_columns - .min_column_expr(column, in_list.expr(), field) - .ok()?; - let max = required_columns - .max_column_expr(column, in_list.expr(), field) - .ok()?; - let non_null = - build_is_null_column_expr(in_list.expr(), schema, required_columns, true)?; - Some((min, max, non_null)) - })(); - let Some((min, max, non_null)) = statistics else { - required_columns.columns.truncate(required_columns_len); - return None; - }; + let (min, max, non_null) = + min_max_non_null_exprs(column, in_list.expr(), field, schema, required_columns)?; let may_match = match domain { CompactInListDomain::String(values) => { Arc::new(StringInListPruningExpr::new(membership, min, max, values)) @@ -1821,6 +1879,45 @@ fn build_predicate_expression( return unhandled_hook.handle(expr); } } + if let Some(lookup) = expr.downcast_ref::() { + return build_hash_lookup_pruning_expr(lookup, schema, required_columns) + .unwrap_or_else(|| unhandled_hook.handle(expr)); + } + // A partitioned hash join hides its per-partition filters under a `CASE` on the + // repartition hash. A row takes exactly one branch, so a container may match + // only if some branch may: the branches' disjunction is a sound relaxation, and + // the `WHEN`s (a hash, which no statistics describe) can be dropped. + if let Some(case) = expr.downcast_ref::() { + // Only a Boolean `CASE` is a predicate; anything else is a value for + // whatever compares it to handle. + if !matches!(case.data_type(schema), Ok(DataType::Boolean)) { + return unhandled_hook.handle(expr); + } + // Without an `ELSE`, unmatched rows are UNKNOWN: not a match, but not + // FALSE either, which full-match inference must be able to tell apart. + if case.else_expr().is_none() { + return unhandled_hook.handle(expr); + } + return case + .when_then_expr() + .iter() + .map(|(_, then)| then) + .chain(case.else_expr()) + .map(|branch| { + build_predicate_expression( + branch, + schema, + required_columns, + unhandled_hook, + max_in_list_size, + properties, + ) + }) + .reduce(|acc, branch| { + Arc::new(phys_expr::BinaryExpr::new(acc, Operator::Or, branch)) as _ + }) + .unwrap_or_else(|| unhandled_hook.handle(expr)); + } let (left, op, right) = { if let Some(bin_expr) = expr.downcast_ref::() { @@ -2412,7 +2509,11 @@ mod tests { self as phys_expr, DynamicFilterPhysicalExpr, }; use datafusion_physical_expr::planner::logical2physical; + use datafusion_physical_plan::joins::join_hash_map::JoinHashMapU32; + use datafusion_physical_plan::joins::key_range_bitmap::KeyRangeBitmap; + use datafusion_physical_plan::joins::{Map, SeededRandomState}; use itertools::Itertools; + use rstest::rstest; #[derive(Debug, Default)] /// Mock statistic provider for tests @@ -7344,4 +7445,150 @@ mod tests { "c1_null_count@2 != row_count@3 AND c1_min@0 <= a AND a <= c1_max@1"; assert_eq!(res.to_string(), expected); } + + /// A hash-table lookup over `column` carrying a bitmap over `keys`. + fn hash_lookup( + column: &Arc, + keys: &[i64], + ) -> Arc { + let array: ArrayRef = Arc::new(Int64Array::from(keys.to_vec())); + let bitmap = KeyRangeBitmap::try_new( + &array, + &ScalarValue::Int64(Some(*keys.iter().min().unwrap())), + &ScalarValue::Int64(Some(*keys.iter().max().unwrap())), + keys.len(), + ); + Arc::new(HashTableLookupExpr::new( + vec![Arc::clone(column)], + SeededRandomState::with_seed(1), + Arc::new(Map::HashMap(Box::new(JoinHashMapU32::with_capacity(1)))), + "hash_lookup".to_string(), + bitmap.map(Arc::new), + )) + } + + #[test] + fn test_hash_lookup_pruning_via_min_max() { + let schema = Arc::new(Schema::new(vec![Field::new("b", DataType::Int64, true)])); + let column: Arc = Arc::new(phys_expr::Column::new("b", 0)); + // Gapped keys: 10,20,30 then a far-away 100000, so the bitmap is sparse. + let lookup = hash_lookup(&column, &[10, 20, 30, 100_000]); + + let predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(lookup) + .unwrap(); + + // An opaque lookup yields no literal guarantees, so `contained()` is never + // consulted and every exclusion below comes from the min/max rewrite. + assert!(predicate.literal_guarantees().is_empty()); + + let statistics = TestStatistics::new().with( + "b", + ContainerStats::new_i64( + vec![Some(5000), Some(15), Some(50000)], + vec![Some(8000), Some(25), Some(60000)], + ), + ); + + let result = predicate.prune(&statistics).unwrap(); + // Containers 0 and 2 sit in the gap between 30 and 100000; container 1 + // holds 20. The bitmap excludes the first and third. + assert_eq!(result, vec![false, true, false]); + } + + #[rstest] + #[case::widening_i32_to_i64(DataType::Int32, DataType::Int64, true)] + #[case::widening_u32_to_i64(DataType::UInt32, DataType::Int64, true)] + #[case::widening_u16_to_i64(DataType::UInt16, DataType::Int64, true)] + #[case::same_type_i64_to_i64(DataType::Int64, DataType::Int64, true)] + #[case::narrowing_i64_to_i32(DataType::Int64, DataType::Int32, false)] + #[case::narrowing_u32_to_i16(DataType::UInt32, DataType::Int16, false)] + fn test_hash_lookup_pruning_through_cast( + #[case] file_type: DataType, + #[case] cast_type: DataType, + #[case] pruned: bool, + ) { + let schema = + Arc::new(Schema::new(vec![Field::new("b", file_type.clone(), true)])); + let column: Arc = Arc::new(phys_expr::Column::new("b", 0)); + let cast_column: Arc = + Arc::new(phys_expr::CastExpr::new(column, cast_type, None)); + let lookup = hash_lookup(&cast_column, &[10, 20, 30, 100_000]); + + let predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(lookup) + .unwrap(); + + let to_file_type = |values: Vec>| -> ArrayRef { + let values: ArrayRef = Arc::new(values.into_iter().collect::()); + arrow::compute::cast(&values, &file_type).unwrap() + }; + let statistics = TestStatistics::new().with( + "b", + ContainerStats::new() + .with_min(to_file_type(vec![Some(5000), Some(15), Some(50000)])) + .with_max(to_file_type(vec![Some(8000), Some(25), Some(60000)])), + ); + + let result = predicate.prune(&statistics).unwrap(); + // Containers 0 and 2 sit in the gap between 30 and 100000; container 1 + // holds 20. A working cast excludes the first and third; a declined one + // falls back to keeping every container. + let expected = if pruned { + vec![false, true, false] + } else { + vec![true, true, true] + }; + assert_eq!(result, expected); + } + + #[test] + fn test_partition_routed_hash_lookup_pruning() { + let schema = Arc::new(Schema::new(vec![Field::new("b", DataType::Int32, true)])); + let column: Arc = Arc::new(phys_expr::Column::new("b", 0)); + let literal = |value: ScalarValue| -> Arc { + Arc::new(phys_expr::Literal::new(value)) + }; + // One branch per build partition, each holding its own slice of the build + // side, as `build_partitioned_filter` produces them. + let case: Arc = Arc::new( + phys_expr::CaseExpr::try_new( + Some(literal(ScalarValue::UInt64(Some(0)))), + vec![ + ( + literal(ScalarValue::UInt64(Some(0))), + hash_lookup(&column, &[10, 20, 30]), + ), + ( + literal(ScalarValue::UInt64(Some(1))), + hash_lookup(&column, &[40, 50, 60]), + ), + ], + // Partitions with no build rows reject everything routed to them. + Some(literal(ScalarValue::Boolean(Some(false)))), + ) + .unwrap(), + ); + + let predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(case) + .unwrap(); + + let statistics = TestStatistics::new().with( + "b", + ContainerStats::new_i32( + vec![Some(5), Some(15), Some(45), Some(100)], + vec![Some(8), Some(25), Some(55), Some(200)], + ), + ); + + let result = predicate.prune(&statistics).unwrap(); + // Container 1 ([15,25]) holds 20 from the first branch and container 2 + // ([45,55]) holds 50 from the second, so a branch may match in each. The + // other two intersect neither branch, nor the `ELSE`. + assert_eq!(result, vec![false, true, true, false]); + } } diff --git a/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt b/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt index c6700ebf0b97c..7c2e0f6baf4ca 100644 --- a/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt +++ b/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt @@ -227,7 +227,7 @@ SELECT WHEN i < 200 THEN 2000 + (i - 100) WHEN i < 300 THEN 3000 + (i - 200) ELSE (i - 300) - END AS BIGINT) AS b + END AS INT) AS b FROM generate_series(0, 399) AS t(i); statement ok @@ -243,6 +243,9 @@ OPTIONS ( statement ok drop table rgsel_src; +# `b` is stored as INT32 but declared BIGINT, so the schema adapter wraps every +# reference in a widening `CAST(b AS BIGINT)`. The hash-lookup pruning expr must +# see through that cast rather than silently losing the bitmap below. statement ok CREATE EXTERNAL TABLE rgsel (a BIGINT NOT NULL, b BIGINT NOT NULL) STORED AS PARQUET @@ -276,6 +279,33 @@ SELECT b FROM rgsel WHERE a >= 50 ORDER BY b ASC LIMIT 5; statement ok set datafusion.execution.parquet.pushdown_filters = true; +# Regression test: a hash join's build side must get the same compact, +# sorted-domain min/max pruning an ordinary large `IN (...)` list already +# gets - and do better than the plain min/max *bounds* check every join +# already gets for free. `dim`'s overall envelope [1000, 3050] spans RG 1 +# (b=2000..2099) entirely, so bounds alone cannot exclude it; only the +# discrete check can prove none of {1000, 1050, 3050} falls inside it. +# Force the hash-table-lookup path (rather than `InList`) regardless of size. +statement ok +CREATE TABLE dim AS VALUES (1000), (1050), (3050); + +statement ok +set datafusion.optimizer.hash_join_inlist_pushdown_max_distinct_values = 0; + +query TT +explain analyze select rgsel.b from dim join rgsel on dim.column1 = rgsel.b; +---- +Plan with Metrics +01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(column1@0, b@0)], projection=[b@1], metrics=[output_rows=3, elapsed_compute=, output_bytes=, output_batches=1, build_mem_used=, array_map_created_count=0, build_input_batches=1, build_input_rows=3, input_batches=2, input_rows=3, build_time=, join_time=, avg_fanout=100% (3/3), probe_hit_rate=100% (3/3)] +02)--DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] +03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_row_group_pruning/rgsel.parquet]]}, projection=[b], file_type=parquet, predicate=DynamicFilter [ b@1 >= 1000 AND b@1 <= 3050 AND hash_lookup ], dynamic_rg_pruning=eligible, pruning_predicate=b_null_count@1 != row_count@2 AND b_max@0 >= 1000 AND b_null_count@1 != row_count@2 AND b_min@3 <= 3050 AND b_null_count@1 != row_count@2 AND KEY_RANGE_BITMAP(b_min@3, b_max@0, 3/4096 buckets), required_guarantees=[], metrics=[output_rows=3, elapsed_compute=, output_bytes=, output_batches=2, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=4 total → 2 matched, row_groups_pruned_bloom_filter=2 total → 2 matched, page_index_pages_pruned=20 total → 3 matched, page_index_rows_pruned=200 total → 30 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, bytes_processed=, bytes_scanned=, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=3, pushdown_rows_pruned=27, row_groups_pruned_dynamic_filter=0, predicate_cache_inner_records=200, predicate_cache_records=13, bloom_filter_eval_time=, metadata_load_time=, page_index_eval_time=, row_pushdown_eval_time=, statistics_eval_time=, time_elapsed_opening=, time_elapsed_processing=, time_elapsed_scanning_total=, time_elapsed_scanning_until_data=, output_rows_skew=, scan_efficiency_ratio=] + +statement ok +RESET datafusion.optimizer.hash_join_inlist_pushdown_max_distinct_values; + +statement ok +drop table dim; + statement ok drop table rgsel; diff --git a/datafusion/sqllogictest/test_files/parquet_filter_pushdown.slt b/datafusion/sqllogictest/test_files/parquet_filter_pushdown.slt index f7d23001fcf50..83dae8c520de5 100644 --- a/datafusion/sqllogictest/test_files/parquet_filter_pushdown.slt +++ b/datafusion/sqllogictest/test_files/parquet_filter_pushdown.slt @@ -1053,3 +1053,50 @@ set datafusion.execution.parquet.reorder_filters = false; statement ok DROP TABLE dict_filter_bug; + +# Regression: a `CASE` with no `ELSE` leaves unmatched rows UNKNOWN, not FALSE. +# Treating them as FALSE lets the inverse predicate call a row group fully +# matched and skip the row filter. Required columns admit no `IS NULL` rescue. +statement ok +set datafusion.execution.parquet.pushdown_filters = true; + +statement ok +COPY (SELECT column1 AS a, column2 AS b FROM (VALUES (2, 1), (1, 2))) +TO 'test_files/scratch/parquet_filter_pushdown/case_without_else.parquet' +STORED AS PARQUET +OPTIONS ('format.max_row_group_size' '1'); + +statement ok +CREATE EXTERNAL TABLE case_without_else (a BIGINT NOT NULL, b BIGINT NOT NULL) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_filter_pushdown/case_without_else.parquet'; + +query II +SELECT * FROM case_without_else +WHERE NOT (CASE WHEN a = 1 THEN false END) AND b > 0 ORDER BY a; +---- +1 2 + +query II +SELECT * FROM case_without_else +WHERE NOT (CASE WHEN a = 1 THEN false END) AND b > 0 LIMIT 1; +---- +1 2 + +query II +SELECT * FROM case_without_else +WHERE NOT (CASE WHEN a = 1 THEN false ELSE NULL END) AND b > 0 ORDER BY a; +---- +1 2 + +query II +SELECT * FROM case_without_else +WHERE NOT (CASE WHEN a = 1 THEN NULL ELSE false END) AND b > 0 ORDER BY a; +---- +2 1 + +statement ok +set datafusion.execution.parquet.pushdown_filters = false; + +statement ok +DROP TABLE case_without_else; From 563832c4f004e5058ec0287196257b87133ce446 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 6 Nov 2026 21:25:03 +0800 Subject: [PATCH 17/19] fix: complete join pruning NULL and producer-lifetime safety Local fork correctness changes on top of the twelve upstream backports: - treat computed NULL/TRY_CAST/CASE results conservatively for full-match predicates, and keep the inverse-row-group guard; - retain dynamic-filter producer reservations through snapshots and partitioned remapping; - make visited-bitmap charges execution-local and carry the accepted key-range-bitmap reservation with the derived lookup. These are local changes, not upstream PR content. Local regression tests are kept in the following commit. Signed-off-by: discord9 --- .../src/row_group_filter.rs | 6 ++ .../physical-expr/src/expressions/in_list.rs | 21 ++++++- .../physical-plan/src/joins/hash_join/exec.rs | 58 ++++++++++-------- .../src/joins/hash_join/exec/prepared.rs | 7 ++- .../joins/hash_join/partitioned_hash_eval.rs | 12 ++++ .../src/joins/hash_join/shared_bounds.rs | 43 ++++++++++---- .../src/joins/hash_join/stream.rs | 2 + .../src/joins/key_range_bitmap.rs | 16 +++++ datafusion/pruning/src/pruning_predicate.rs | 59 ++++++++++++++++++- 9 files changed, 182 insertions(+), 42 deletions(-) diff --git a/datafusion/datasource-parquet/src/row_group_filter.rs b/datafusion/datasource-parquet/src/row_group_filter.rs index 2b487ec80f944..67d139f8d1ee2 100644 --- a/datafusion/datasource-parquet/src/row_group_filter.rs +++ b/datafusion/datasource-parquet/src/row_group_filter.rs @@ -435,6 +435,12 @@ impl RowGroupAccessPlanFilter { else { return; }; + // The inverse can introduce filter-only rewrites that were hidden + // inside an unhandled expression in the forward predicate. In that + // case pruning UNKNOWN cannot establish that every row matches. + if !inverted_predicate.can_be_inverted_for_full_match() { + return; + } let inverted_pruning_stats = RowGroupPruningStatistics { parquet_schema: pruning_stats.parquet_schema, diff --git a/datafusion/physical-expr/src/expressions/in_list.rs b/datafusion/physical-expr/src/expressions/in_list.rs index 874e149b58328..3c60f5d59a32f 100644 --- a/datafusion/physical-expr/src/expressions/in_list.rs +++ b/datafusion/physical-expr/src/expressions/in_list.rs @@ -52,6 +52,9 @@ pub struct InListExpr { list: Vec>, negated: bool, static_filter: Option>, + // Opaque lease for producer-owned buffers referenced by join-created lists. + // It does not account for this expression's own allocations. + producer_lease: Option>, } impl Debug for InListExpr { @@ -155,9 +158,21 @@ impl InListExpr { list, negated, static_filter, + producer_lease: None, } } + /// Retain an existing producer's reservation while this expression can + /// reference its buffers. This does not account for expression allocations. + #[doc(hidden)] + pub fn with_producer_lease( + mut self, + lease: Arc, + ) -> Self { + self.producer_lease = Some(lease); + self + } + /// Input expression pub fn expr(&self) -> &Arc { &self.expr @@ -447,12 +462,14 @@ impl PhysicalExpr for InListExpr { children: Vec>, ) -> Result> { // assume the static_filter will not change during the rewrite process - Ok(Arc::new(InListExpr::new( + let mut remapped = InListExpr::new( Arc::clone(&children[0]), children[1..].to_vec(), self.negated, self.static_filter.as_ref().map(Arc::clone), - ))) + ); + remapped.producer_lease.clone_from(&self.producer_lease); + Ok(Arc::new(remapped)) } fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 738724ae903da..96464c6691ec2 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -199,19 +199,17 @@ fn array_map_key_range( } /// Immutable build buffers and their durable reservation. -struct JoinBuildData { +pub(super) struct JoinBuildData { /// Hash table with row indices into `batch`, also shared with dynamic filters. - map: Arc, + pub(super) map: Arc, /// The input rows for the build side. - batch: RecordBatch, + pub(super) batch: RecordBatch, /// Evaluated build-side key expressions. - values: Vec, + pub(super) values: Vec, /// Bounds computed from the build side; absent for an empty partition. - bounds: Option, - /// IN-list values or a hash-table reference used for filter pushdown. - membership: PushdownStrategy, + pub(super) bounds: Option, // Declared last so accounted-for build allocations are released first. - reservation: MemoryReservation, + pub(super) reservation: MemoryReservation, } /// A build lease and mutable bookkeeping for one join execution. @@ -224,7 +222,10 @@ pub(super) struct JoinLeftData { pub(super) probe_side_non_empty: AtomicBool, /// Whether any probe partition saw a NULL key (for null-aware anti joins). pub(super) probe_side_has_null: AtomicBool, - // Drop the reservation after the mutable state it also accounts for. + /// IN-list values or a hash-table reference used for filter pushdown. + membership: PushdownStrategy, + /// The mutable visited bitmap belongs to this execution, not the build owner. + _visited_reservation: MemoryReservation, build: Arc, } @@ -268,7 +269,11 @@ impl JoinLeftData { /// returns a reference to the InList values for filter pushdown pub(super) fn membership(&self) -> &PushdownStrategy { - &self.build.membership + &self.membership + } + + pub(super) fn build_lease(&self) -> Arc { + Arc::clone(&self.build) } /// Borrow bounds without sharing a consuming join's dynamic-filter state. @@ -2602,16 +2607,17 @@ async fn collect_left_input( metrics.build_mem_used.add(keys_size); // Reserve additional memory for visited indices bitmap and create shared builder - let visited_indices_bitmap = if with_visited_indices_bitmap { + let (visited_indices_bitmap, visited_reservation) = if with_visited_indices_bitmap { let bitmap_size = bit_util::ceil(batch.num_rows(), 8); reservation.try_grow(bitmap_size)?; metrics.build_mem_used.add(bitmap_size); + let visited_reservation = reservation.split(bitmap_size); let mut bitmap_buffer = BooleanBufferBuilder::new(batch.num_rows()); bitmap_buffer.append_n(num_rows, false); - bitmap_buffer + (bitmap_buffer, visited_reservation) } else { - BooleanBufferBuilder::new(0) + (BooleanBufferBuilder::new(0), reservation.new_empty()) }; let map = Arc::new(join_hash_map); @@ -2645,9 +2651,8 @@ async fn collect_left_input( } else { // Past the InList threshold use a bucket bitmap for container pruning. let pruning_bitmap = match (left_values.as_slice(), bounds.as_ref()) { - ([keys], Some(bounds)) if !keys.is_empty() => bounds - .get_column_bounds(0) - .and_then(|b| { + ([keys], Some(bounds)) if !keys.is_empty() => { + bounds.get_column_bounds(0).and_then(|b| { KeyRangeBitmap::try_new( keys, &b.min, @@ -2655,17 +2660,19 @@ async fn collect_left_input( map.num_of_distinct_key(), ) }) - .map(Arc::new), + } _ => None, }; - // Held for the join's lifetime, so charge it like the maps; it is - // optional, so skip it rather than fail when the pool is full. - let pruning_bitmap = pruning_bitmap.filter(|bitmap| { - let ok = reservation.try_grow(bitmap.size()).is_ok(); - if ok { - metrics.build_mem_used.add(bitmap.size()); + // A derived pruning expression may outlive the build. Transfer the + // accepted charge to the bitmap; admission failure remains optional. + let pruning_bitmap = pruning_bitmap.and_then(|bitmap| { + let bytes = bitmap.size(); + if reservation.try_grow(bytes).is_ok() { + metrics.build_mem_used.add(bytes); + Some(Arc::new(bitmap.with_reservation(reservation.split(bytes)))) + } else { + None } - ok }); PushdownStrategy::Map(Arc::clone(&map), pruning_bitmap) } @@ -2692,12 +2699,13 @@ async fn collect_left_input( probe_threads_counter: AtomicUsize::new(probe_threads_count), probe_side_non_empty: AtomicBool::new(false), probe_side_has_null: AtomicBool::new(false), + membership, + _visited_reservation: visited_reservation, build: Arc::new(JoinBuildData { map, batch, values: left_values, bounds, - membership, reservation, }), }; diff --git a/datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs b/datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs index a19ad5fdbcc1a..9ec405e919bab 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs @@ -38,6 +38,8 @@ use datafusion_execution::memory_pool::MemoryPool; /// nested build columns remain unsupported. UTF-8 and fixed-size binary keys /// use hash-table membership filters instead of copying range or IN-list values. pub struct PreparedHashJoinBuild { + // Membership may retain build buffers; drop it before the build reservation. + membership: PushdownStrategy, build: Arc, keys: Vec, null_equality: NullEquality, @@ -58,7 +60,7 @@ impl fmt::Debug for PreparedHashJoinBuild { impl PreparedHashJoinBuild { /// Return the retained build reservation, excluding all per-probe state. pub fn reserved_bytes(&self) -> usize { - self.build.reservation.size() + self.build.reservation.size() + self.membership.bitmap_reserved_bytes() } /// Return the complete build row count, including duplicate and null keys. @@ -74,6 +76,8 @@ impl PreparedHashJoinBuild { probe_threads_counter: AtomicUsize::new(probe_threads), probe_side_non_empty: AtomicBool::new(false), probe_side_has_null: AtomicBool::new(false), + membership: self.membership.clone(), + _visited_reservation: self.build.reservation.new_empty(), } } @@ -199,6 +203,7 @@ impl HashJoinExec { .await?; Ok(Arc::new(PreparedHashJoinBuild { build: data.build, + membership: data.membership, keys, null_equality: self.null_equality, })) diff --git a/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs b/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs index fca91ccfd975b..f249a6873896a 100644 --- a/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs +++ b/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs @@ -264,6 +264,8 @@ pub struct HashTableLookupExpr { random_state: SeededRandomState, /// Map to check against (hash table or array map) map: Arc, + /// Retain the build reservation as long as the published map is retained. + build_lease: Option>, /// Description for display description: String, /// Shared pruning-only view of the build side's key range @@ -292,11 +294,20 @@ impl HashTableLookupExpr { on_columns, random_state, map, + build_lease: None, description, pruning_bitmap, } } + pub(super) fn with_build_lease( + mut self, + lease: Option>, + ) -> Self { + self.build_lease = lease; + self + } + /// Which parts of this build side's key range hold keys, for container pruning. pub fn pruning_bitmap(&self) -> Option<&Arc> { self.pruning_bitmap.as_ref() @@ -357,6 +368,7 @@ impl PhysicalExpr for HashTableLookupExpr { on_columns: children, random_state: self.random_state.clone(), map: Arc::clone(&self.map), + build_lease: self.build_lease.clone(), description: self.description.clone(), pruning_bitmap: self.pruning_bitmap.clone(), })) diff --git a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs index 9182fc4173dea..02eb96af97dfd 100644 --- a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs +++ b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs @@ -26,7 +26,7 @@ use crate::ExecutionPlanProperties; use crate::Partitioning; use crate::joins::Map; use crate::joins::PartitionMode; -use crate::joins::hash_join::exec::HASH_JOIN_SEED; +use crate::joins::hash_join::exec::{HASH_JOIN_SEED, JoinBuildData}; use crate::joins::hash_join::inlist_builder::build_struct_fields; use crate::joins::hash_join::partitioned_hash_eval::{ HashExpr, HashTableLookupExpr, SeededRandomState, @@ -96,6 +96,7 @@ impl PartitionBounds { fn create_membership_predicate( on_right: &[PhysicalExprRef], pushdown: PushdownStrategy, + build_lease: Option>, random_state: &SeededRandomState, schema: &Schema, ) -> Result>> { @@ -129,23 +130,25 @@ fn create_membership_predicate( }; // Use InListExpr::try_new_from_array() to build an InList with static_filter optimization (hash-based lookup) - Ok(Some(Arc::new(InListExpr::try_new_from_array( - expr, - in_list_array, - false, - schema, - )?))) + let mut in_list = + InListExpr::try_new_from_array(expr, in_list_array, false, schema)?; + if let Some(build_lease) = build_lease { + in_list = in_list.with_producer_lease(build_lease); + } + Ok(Some(Arc::new(in_list))) } // Use hash table lookup for large build sides - PushdownStrategy::Map(hash_map, pruning_bitmap) => { - Ok(Some(Arc::new(HashTableLookupExpr::new( + PushdownStrategy::Map(hash_map, pruning_bitmap) => Ok(Some(Arc::new( + HashTableLookupExpr::new( on_right.to_vec(), random_state.clone(), hash_map, "hash_lookup".to_string(), pruning_bitmap, - )) as Arc)) - } + ) + .with_build_lease(build_lease), + ) + as Arc)), // Empty partition - should not create a filter for this PushdownStrategy::Empty => Ok(None), } @@ -288,16 +291,27 @@ pub(crate) enum PushdownStrategy { Empty, } +impl PushdownStrategy { + pub(super) fn bitmap_reserved_bytes(&self) -> usize { + match self { + Self::Map(_, Some(bitmap)) => bitmap.reserved_bytes(), + _ => 0, + } + } +} + /// Build-side data reported by a single partition pub(crate) enum PartitionBuildData { Partitioned { partition_id: usize, pushdown: PushdownStrategy, + build_lease: Option>, bounds: PartitionBounds, keys_have_null: bool, }, CollectLeft { pushdown: PushdownStrategy, + build_lease: Option>, bounds: PartitionBounds, keys_have_null: bool, }, @@ -308,6 +322,7 @@ pub(crate) enum PartitionBuildData { struct PartitionData { bounds: PartitionBounds, pushdown: PushdownStrategy, + build_lease: Option>, /// Whether any build key of this partition is NULL. Decides whether the pushed /// filter must keep probe-side NULL rows for a null-equal join to match them. keys_have_null: bool, @@ -494,6 +509,7 @@ impl SharedBuildAccumulator { PartitionBuildData::Partitioned { partition_id, pushdown, + build_lease, bounds, keys_have_null, }, @@ -507,6 +523,7 @@ impl SharedBuildAccumulator { } partitions[partition_id] = PartitionStatus::Reported(PartitionData { pushdown, + build_lease, bounds, keys_have_null, }); @@ -514,6 +531,7 @@ impl SharedBuildAccumulator { ( PartitionBuildData::CollectLeft { pushdown, + build_lease, bounds, keys_have_null, }, @@ -526,6 +544,7 @@ impl SharedBuildAccumulator { if matches!(data, PartitionStatus::Pending) { *data = PartitionStatus::Reported(PartitionData { pushdown, + build_lease, bounds, keys_have_null, }); @@ -621,6 +640,7 @@ impl SharedBuildAccumulator { let membership_expr = create_membership_predicate( &self.on_right, partition_data.pushdown.clone(), + partition_data.build_lease.clone(), &HASH_JOIN_SEED, self.probe_schema.as_ref(), )?; @@ -669,6 +689,7 @@ impl SharedBuildAccumulator { let membership_expr = create_membership_predicate( &self.on_right, partition.pushdown.clone(), + partition.build_lease.clone(), &HASH_JOIN_SEED, self.probe_schema.as_ref(), )?; diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs index fc2bcd82e775d..302ffb77dc528 100644 --- a/datafusion/physical-plan/src/joins/hash_join/stream.rs +++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs @@ -569,11 +569,13 @@ impl HashJoinStream { PartitionMode::Partitioned => PartitionBuildData::Partitioned { partition_id: self.partition, pushdown, + build_lease: Some(left_data.build_lease()), bounds, keys_have_null, }, PartitionMode::CollectLeft => PartitionBuildData::CollectLeft { pushdown, + build_lease: Some(left_data.build_lease()), bounds, keys_have_null, }, diff --git a/datafusion/physical-plan/src/joins/key_range_bitmap.rs b/datafusion/physical-plan/src/joins/key_range_bitmap.rs index efdb5467c897e..9a8767d858132 100644 --- a/datafusion/physical-plan/src/joins/key_range_bitmap.rs +++ b/datafusion/physical-plan/src/joins/key_range_bitmap.rs @@ -21,6 +21,7 @@ use arrow::array::{Array, ArrayRef, BooleanArray, downcast_integer_array}; use arrow::buffer::{BooleanBuffer, MutableBuffer}; use arrow::util::bit_util; use datafusion_common::ScalarValue; +use datafusion_execution::memory_pool::MemoryReservation; use num_traits::AsPrimitive; use crate::joins::array_map::ArrayMap; @@ -58,6 +59,8 @@ pub struct KeyRangeBitmap { shift: u32, /// One bit per bucket. bits: BooleanBuffer, + /// Reservation follows this bitmap even when a pruning expression retains it. + reservation: Option, } impl KeyRangeBitmap { @@ -98,9 +101,22 @@ impl KeyRangeBitmap { span, shift, bits, + reservation: None, }) } + /// Attach the existing build-side charge before publishing this bitmap. + pub(crate) fn with_reservation(mut self, reservation: MemoryReservation) -> Self { + debug_assert_eq!(reservation.size(), self.size()); + self.reservation = Some(reservation); + self + } + + /// Bytes currently charged to this bitmap. + pub(crate) fn reserved_bytes(&self) -> usize { + self.reservation.as_ref().map_or(0, MemoryReservation::size) + } + /// One Boolean per `[min[i], max[i]]` container interval: `false` proves no /// build key lies in that container, `None` where a bound is absent or the /// statistics are not integer keys. diff --git a/datafusion/pruning/src/pruning_predicate.rs b/datafusion/pruning/src/pruning_predicate.rs index 65fb737bfbe14..158e74a54aaed 100644 --- a/datafusion/pruning/src/pruning_predicate.rs +++ b/datafusion/pruning/src/pruning_predicate.rs @@ -388,12 +388,16 @@ pub struct PruningPredicate { max_in_list_size: usize, /// Whether its logical inverse can safely prove every row matches. can_be_inverted_for_full_match: bool, + /// A simplified snapshot may retain producer buffers through its literal + /// children while dropping the input expression's producer lease. This is + /// dropped after the expressions and literal guarantees that may share them. + _snapshot_lease: Option>, } #[derive(Default)] struct PruningExpressionProperties { - /// The rewrite preserves which rows can be TRUE, but not whether rejected - /// rows are FALSE or UNKNOWN. + /// A rewrite or computed NULL preserves which rows can be TRUE, but does + /// not prove whether rejected rows are FALSE rather than UNKNOWN. has_filter_semantics_only: bool, } @@ -531,6 +535,7 @@ impl<'a> PruningPredicateBuilder<'a> { // In particular this unravels any `DynamicFilterPhysicalExpr`s by snapshotting them // so that PruningPredicate can work with a static expression. let tf = snapshot_physical_expr_opt(predicate)?; + let mut snapshot_lease = None; if tf.transformed { // If we had an expression such as Dynamic(part_col < 5 and col < 10) // (this could come from something like `select * from t order by part_col, col, limit 10`) @@ -539,7 +544,11 @@ impl<'a> PruningPredicateBuilder<'a> { // the expression we have now is `8 < 5 and col < 10`. // Thus we need as simplifier pass to get `false and col < 10` => `false` here. let simplifier = PhysicalExprSimplifier::new(&file_schema); - predicate = simplifier.simplify(tf.data)?; + let snapshot = tf.data; + predicate = simplifier.simplify(Arc::clone(&snapshot))?; + if !Arc::ptr_eq(&predicate, &snapshot) { + snapshot_lease = Some(snapshot); + } } else { predicate = tf.data; } @@ -548,6 +557,37 @@ impl<'a> PruningPredicateBuilder<'a> { // build predicate expression once let mut required_columns = RequiredColumns::new(); let mut properties = PruningExpressionProperties::default(); + // A failing TRY_CAST or safe CAST can manufacture NULL even when its + // input column is known non-null. Statistics over that column cannot + // then prove the inverse predicate is FALSE for every row. Inspect the + // original expression before statistics rewriting loses cast options. + if predicate + .apply(|node| { + let cast = if let Some(cast) = + node.downcast_ref::() + { + Some((cast.expr(), cast.cast_type())) + } else if let Some(cast) = node.downcast_ref::() { + cast.cast_options() + .safe + .then_some((cast.expr(), cast.cast_type())) + } else { + None + }; + if let Some((expr, target)) = cast + && !expr.data_type(&file_schema).is_ok_and(|source| { + phys_expr::CastExpr::check_bigger_cast(target, &source) + }) + { + properties.has_filter_semantics_only = true; + } + // Continue through total outer casts to inspect nested casts too. + Ok(TreeNodeRecursion::Continue) + }) + .is_err() + { + properties.has_filter_semantics_only = true; + } let predicate_expr = build_predicate_expression( &predicate, &file_schema, @@ -567,6 +607,7 @@ impl<'a> PruningPredicateBuilder<'a> { predicate_expr, required_columns, orig_expr: predicate, + _snapshot_lease: snapshot_lease, literal_guarantees, max_in_list_size: self.max_in_list_size, can_be_inverted_for_full_match: !properties.has_filter_semantics_only, @@ -1898,6 +1939,18 @@ fn build_predicate_expression( if case.else_expr().is_none() { return unhandled_hook.handle(expr); } + // OR-ing results is a valid may-be-TRUE relaxation, but it cannot + // prove a full match when a selected result can be UNKNOWN. + for branch in case + .when_then_expr() + .iter() + .map(|(_, then)| then) + .chain(case.else_expr()) + { + if !matches!(branch.nullable(schema), Ok(false)) { + properties.has_filter_semantics_only = true; + } + } return case .when_then_expr() .iter() From b05a450f82fdebdb89578efdc41804afb76db4a1 Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 6 Nov 2026 21:25:03 +0800 Subject: [PATCH 18/19] test: add local join pruning and producer-lifetime regressions Add the local Arrow 59 all-valid, Parquet inverse-order/full-match, aggregate producer, prepared-build lease and execution-local bitmap regression coverage. These tests are intentionally last, after the upstream backports and local correctness changes. Signed-off-by: discord9 --- datafusion/common/src/utils/memory.rs | 116 ++- .../core/tests/parquet/file_statistics.rs | 156 +++ .../src/row_group_filter.rs | 954 ++++++++++++++++++ .../src/statistics_order_tests.rs | 25 + .../physical-plan/src/aggregates/mod.rs | 216 +++- .../physical-plan/src/joins/hash_join/exec.rs | 325 ++++++ .../joins/hash_join/exec/prepared/tests.rs | 33 +- .../src/joins/hash_join/shared_bounds.rs | 77 ++ .../src/joins/hash_join/stream.rs | 1 + datafusion/pruning/src/pruning_predicate.rs | 308 ++++++ 10 files changed, 2202 insertions(+), 9 deletions(-) diff --git a/datafusion/common/src/utils/memory.rs b/datafusion/common/src/utils/memory.rs index e235a578587cb..4522ef35c95a5 100644 --- a/datafusion/common/src/utils/memory.rs +++ b/datafusion/common/src/utils/memory.rs @@ -451,8 +451,9 @@ mod record_batch_tests { use arrow::array::{ ArrayData, ArrayRef, BinaryViewArray, Float64Array, Int16Array, Int32Array, Int64Array, LargeListViewArray, ListArray, ListViewArray, RunArray, StringArray, - StringViewArray, new_null_array, + StringViewArray, StructArray, new_null_array, }; + use arrow::buffer::{NullBuffer, OffsetBuffer}; use arrow::datatypes::{ DataType, Field, Int16Type, Int32Type, Int64Type, Schema, UnionFields, UnionMode, }; @@ -611,6 +612,119 @@ mod record_batch_tests { assert_eq!(counter.memory_usage(), get_record_batch_memory_size(&batch)); } + #[test] + fn test_record_batch_memory_counter_retained_all_valid_buffer() { + // ArrayDataBuilder drops all-valid null buffers even though the original + // array still retains their allocation. + let array = + Int32Array::new(vec![1, 2, 3, 4].into(), Some(NullBuffer::new_valid(4))); + let validity = array.nulls().unwrap().buffer(); + let expected = array.values().inner().capacity() + validity.capacity(); + assert!(validity.capacity() > 0); + assert!(array.to_data().nulls().is_none()); + assert_eq!( + array_data_memory_size(&array), + array.values().inner().capacity() + ); + + let schema = Arc::new(Schema::new(vec![ + Field::new("first", DataType::Int32, true), + Field::new("second", DataType::Int32, true), + ])); + let array = Arc::new(array); + let batch = + RecordBatch::try_new(schema, vec![Arc::clone(&array) as ArrayRef, array]) + .unwrap(); + let mut counter = RecordBatchMemoryCounter::new(); + assert_eq!(counter.count_batch(&batch), expected); + assert_eq!(counter.count_batch(&batch), 0); + assert_eq!(counter.memory_usage(), expected); + } + + #[test] + fn test_record_batch_memory_counter_all_valid_slice_of_nullable_array() { + let array = Int32Array::from(vec![None, Some(1), Some(2), Some(3), None]); + let slice = array.slice(1, 3); + let nulls = slice.nulls().unwrap(); + assert_eq!(nulls.null_count(), 0); + assert_eq!( + nulls.buffer().data_ptr(), + array.nulls().unwrap().buffer().data_ptr() + ); + assert!(slice.to_data().nulls().is_none()); + let expected = array.values().inner().capacity() + nulls.buffer().capacity(); + assert!(nulls.buffer().capacity() > 0); + assert_eq!( + array_data_memory_size(&slice), + array.values().inner().capacity() + ); + + let schema = + Arc::new(Schema::new(vec![Field::new("ints", DataType::Int32, true)])); + let slice_batch = + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(slice)]).unwrap(); + let original_batch = RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap(); + let mut counter = RecordBatchMemoryCounter::new(); + assert_eq!(counter.count_batch(&slice_batch), expected); + assert_eq!(counter.count_batch(&original_batch), 0); + assert_eq!(counter.memory_usage(), expected); + } + + #[test] + fn test_record_batch_memory_counter_nested_retained_all_valid_buffer() { + let child = Arc::new(Int32Array::new( + vec![1, 2, 3, 4].into(), + Some(NullBuffer::new_valid(4)), + )); + let child_validity = child.nulls().unwrap().buffer().capacity(); + let child_values = child.values().inner().capacity(); + assert!(child_validity > 0); + let field = Arc::new(Field::new("value", DataType::Int32, true)); + let struct_array = StructArray::new( + vec![Arc::clone(&field)].into(), + vec![Arc::clone(&child) as ArrayRef], + None, + ); + let list_array = ListArray::new( + field, + OffsetBuffer::new(vec![0_i32, 1, 2, 3, 4].into()), + child, + None, + ); + let offset_size = list_array.offsets().inner().inner().capacity(); + assert!(struct_array.to_data().child_data()[0].nulls().is_none()); + assert!(list_array.to_data().child_data()[0].nulls().is_none()); + assert_eq!(array_data_memory_size(&struct_array), child_values); + assert_eq!( + array_data_memory_size(&list_array), + child_values + offset_size + ); + let mut struct_counter = RecordBatchMemoryCounter::new(); + struct_counter.count_array_memory_size(&struct_array); + assert_eq!(struct_counter.memory_usage(), child_values + child_validity); + let mut list_counter = RecordBatchMemoryCounter::new(); + list_counter.count_array_memory_size(&list_array); + assert_eq!( + list_counter.memory_usage(), + child_values + child_validity + offset_size + ); + + let schema = Arc::new(Schema::new(vec![ + Field::new("struct", struct_array.data_type().clone(), false), + Field::new("list", list_array.data_type().clone(), false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![Arc::new(struct_array), Arc::new(list_array)], + ) + .unwrap(); + let expected = child_validity + child_values + offset_size; + let mut counter = RecordBatchMemoryCounter::new(); + assert_eq!(counter.count_batch(&batch), expected); + assert_eq!(counter.count_batch(&batch), 0); + assert_eq!(counter.memory_usage(), expected); + } + #[test] fn test_record_batch_memory_counter_promotes_buffer_set() { let fields = (0..=INLINE_BUFFER_IDS) diff --git a/datafusion/core/tests/parquet/file_statistics.rs b/datafusion/core/tests/parquet/file_statistics.rs index f6d733ec69720..cc1d69f31e76a 100644 --- a/datafusion/core/tests/parquet/file_statistics.rs +++ b/datafusion/core/tests/parquet/file_statistics.rs @@ -412,3 +412,159 @@ fn get_list_file_cache_size(state1: &SessionState) -> usize { .unwrap() .len() } + +/// Regression for SQL CASE simplification hiding a computed NULL from +/// row-group full-match inference and scan LIMIT pruning. +#[tokio::test] +async fn sql_computed_null_cast_limit_preserves_later_row() +-> datafusion_common::Result<()> { + use crate::parquet::utils::MetricsFinder; + use arrow::array::{BooleanArray, Int8Array, Int32Array}; + use arrow::record_batch::RecordBatch; + use datafusion_physical_plan::{collect, displayable, metrics::MetricValue}; + use parquet::arrow::ArrowWriter; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use parquet::file::properties::WriterProperties; + + const PREDICATE: &str = + "NOT(CASE WHEN a=1 THEN TRY_CAST(x AS TINYINT)>0 ELSE false END) AND b>0"; + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("x", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let file = tempfile::Builder::new().suffix(".parquet").tempfile()?; + let properties = WriterProperties::builder() + .set_max_row_group_row_count(Some(2)) + .set_bloom_filter_enabled(false) + .build(); + let mut writer = + ArrowWriter::try_new(file.reopen()?, Arc::clone(&schema), Some(properties))?; + for values in [vec![-200, 0], vec![0]] { + let rows = values.len(); + writer.write(&RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1; rows])), + Arc::new(Int32Array::from(values)), + Arc::new(Int32Array::from(vec![1; rows])), + ], + )?)?; + writer.flush()?; + } + writer.close()?; + let reader = ParquetRecordBatchReaderBuilder::try_new(file.reopen()?)?; + assert_eq!(reader.metadata().num_row_groups(), 2); + assert_eq!(reader.metadata().row_group(0).num_rows(), 2); + assert_eq!(reader.metadata().row_group(1).num_rows(), 1); + drop(reader); + + for pushdown in [false, true] { + let mut config = SessionConfig::new() + .with_target_partitions(1) + .with_batch_size(8192); + config.options_mut().execution.parquet.pushdown_filters = pushdown; + config.options_mut().execution.parquet.reorder_filters = false; + config.options_mut().execution.parquet.enable_page_index = false; + config.options_mut().execution.parquet.bloom_filter_on_read = false; + config.options_mut().execution.parquet.pruning = true; + let ctx = SessionContext::new_with_config(config); + ctx.register_parquet( + "t", + file.path().to_str().unwrap(), + ParquetReadOptions::default(), + ) + .await?; + let oracle = ctx + .sql(&format!( + "SELECT TRY_CAST(x AS TINYINT), {PREDICATE} FROM t" + )) + .await? + .collect() + .await?; + let casts: Vec<_> = oracle + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + }) + .collect(); + let predicate: Vec<_> = oracle + .iter() + .flat_map(|batch| { + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + }) + .collect(); + assert_eq!(casts, [None, Some(0), Some(0)]); + assert_eq!(predicate, [None, Some(true), Some(true)]); + + for limit in ["", " LIMIT 2"] { + let frame = ctx + .sql(&format!("SELECT x FROM t WHERE {PREDICATE}{limit}")) + .await?; + let plan = frame.create_physical_plan().await?; + let plan_text = displayable(plan.as_ref()).indent(true).to_string(); + assert!( + plan_text.contains("TRY_CAST") && !plan_text.contains("CASE"), + "{plan_text}" + ); + if pushdown && !limit.is_empty() { + assert!( + plan_text.contains("limit=2"), + "scan limit not exercised: {plan_text}" + ); + } + let batches = collect(Arc::clone(&plan), ctx.task_ctx()).await?; + let values: Vec<_> = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .map(Option::unwrap) + }) + .collect(); + assert_eq!(values, [0, 0], "pushdown={pushdown}, limit={limit:?}"); + let metrics = + MetricsFinder::find_metrics(plan.as_ref()).expect("Parquet scan metrics"); + for metric in ["row_groups_pruned_statistics", "limit_pruned_row_groups"] { + let value = metrics + .sum(|m| m.value().name() == metric) + .expect("pruning metric"); + let MetricValue::PruningMetrics { + pruning_metrics, .. + } = value + else { + panic!("expected pruning metric {metric}"); + }; + if metric == "row_groups_pruned_statistics" { + assert_eq!(pruning_metrics.fully_matched(), 0); + } else { + assert_eq!(pruning_metrics.pruned(), 0); + } + } + if pushdown { + assert_eq!( + metrics + .sum_by_name("pushdown_rows_pruned") + .unwrap() + .as_usize(), + 1 + ); + } + } + } + Ok(()) +} diff --git a/datafusion/datasource-parquet/src/row_group_filter.rs b/datafusion/datasource-parquet/src/row_group_filter.rs index 67d139f8d1ee2..56b8ce08b531f 100644 --- a/datafusion/datasource-parquet/src/row_group_filter.rs +++ b/datafusion/datasource-parquet/src/row_group_filter.rs @@ -698,6 +698,235 @@ mod tests { } } + /// Exercise a real join producer through the pruning snapshot boundary, + /// and the independently owned bitmap used by a detached pruning expr. + #[tokio::test] + async fn pruning_pipeline_retains_join_producer_and_detached_bitmap() -> Result<()> { + use arrow::array::Int64Array; + use arrow::record_batch::RecordBatch; + use datafusion_common::pruning::PrunableStatistics; + use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; + use datafusion_common::{JoinType, NullEquality, Statistics}; + use datafusion_datasource::memory::MemorySourceConfig; + use datafusion_execution::TaskContext; + use datafusion_execution::config::SessionConfig; + use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryPool}; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_physical_expr::PhysicalExprRef; + use datafusion_physical_expr::expressions::{ + Column as PhysicalColumn, DynamicFilterPhysicalExpr, InListExpr, Literal, + lit as physical_lit, + }; + use datafusion_physical_plan::ExecutionPlan; + use datafusion_physical_plan::common::collect; + use datafusion_physical_plan::filter::FilterExecBuilder; + use datafusion_physical_plan::joins::{ + HashJoinExec, HashTableLookupExpr, PartitionMode, + }; + + for struct_in_list in [true, false] { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int64, false), + Field::new("v", DataType::Int64, false), + ])); + let batch = |keys: Vec, values: Vec| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from(keys)), + Arc::new(Int64Array::from(values)), + ], + ) + }; + let build = batch(vec![0, 1_000_000], vec![10, 20])?; + let build_buffers = build + .columns() + .iter() + .map(|column| column.to_data().buffers()[0].as_ptr() as usize) + .collect::>(); + let probe = batch(vec![0, 1_000_000, 99], vec![10, 20, 10])?; + let left = MemorySourceConfig::try_new_exec( + &[vec![build]], + Arc::clone(&schema), + None, + )?; + let right = MemorySourceConfig::try_new_exec( + &[vec![probe.clone()]], + Arc::clone(&schema), + None, + )?; + let keys = (0..if struct_in_list { 2 } else { 1 }) + .map(|index| { + Arc::new(PhysicalColumn::new(schema.field(index).name(), index)) + as PhysicalExprRef + }) + .collect::>(); + let on = keys + .iter() + .map(|key| (Arc::clone(key), Arc::clone(key))) + .collect(); + let dynamic = Arc::new(DynamicFilterPhysicalExpr::new( + keys.clone(), + physical_lit(true), + )); + let consumer: PhysicalExprRef = dynamic.clone(); + let right = Arc::new(FilterExecBuilder::new(consumer, right).build()?); + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::Inner, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + false, + )? + .with_dynamic_filter_expr(Arc::clone(&dynamic))?; + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 20)); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .build_arc()?; + let mut config = SessionConfig::new(); + if !struct_in_list { + config + .options_mut() + .optimizer + .hash_join_inlist_pushdown_max_size = 0; + } + let task = Arc::new( + TaskContext::default() + .with_runtime(runtime) + .with_session_config(config), + ); + let batches = collect(join.execute(0, task)?).await?; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); + let current = dynamic.current()?; + let mut membership = None; + current.apply(|node| { + if node.downcast_ref::().is_some() + || node.downcast_ref::().is_some() + { + membership = Some(Arc::clone(node)); + return Ok(TreeNodeRecursion::Stop); + } + Ok(TreeNodeRecursion::Continue) + })?; + let membership = membership.expect("join-created membership expression"); + assert!(pool.reserved() > 0); + + if struct_in_list { + let list = membership + .downcast_ref::() + .expect("join IN-list"); + let literal = list.list()[0].downcast_ref::().unwrap(); + let ScalarValue::Struct(values) = literal.value() else { + panic!("join must publish Struct IN literals"); + }; + let buffers = values + .columns() + .iter() + .map(|column| column.to_data().buffers()[0].as_ptr() as usize) + .collect::>(); + assert_eq!( + buffers, build_buffers, + "literal buffers are producer buffers" + ); + let negated: PhysicalExprRef = + Arc::new(NotExpr::new(Arc::clone(&membership))); + let snapshot: PhysicalExprRef = + Arc::new(DynamicFilterPhysicalExpr::new(keys, negated)); + let pruning = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(Arc::clone(&snapshot))?; + assert!( + pruning + .orig_expr() + .downcast_ref::() + .unwrap() + .negated(), + "exercise actual snapshot NOT(IN) simplification" + ); + let producer_charge = pool.reserved(); + assert!(producer_charge > 0); + drop(snapshot); + drop(membership); + drop(current); + drop(batches); + drop(dynamic); + drop(join); + assert_eq!( + pool.reserved(), + producer_charge, + "pruner must retain the entire original producer charge" + ); + let actual = pruning + .orig_expr() + .evaluate(&probe)? + .into_array(probe.num_rows())?; + assert_eq!( + actual.as_ref(), + &BooleanArray::from(vec![false, false, true]) + ); + let unknown = PrunableStatistics::new( + vec![Arc::new(Statistics::new_unknown(&schema))], + Arc::clone(&schema), + ); + assert_eq!(pruning.prune(&unknown)?, vec![true]); + drop(actual); + drop(pruning); + assert_eq!(pool.reserved(), 0, "pruner releases producer charge"); + } else { + let bitmap_bytes = membership + .downcast_ref::() + .unwrap() + .pruning_bitmap() + .expect("sparse bitmap") + .size(); + let dynamic_expr: PhysicalExprRef = dynamic.clone(); + let pruning = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(dynamic_expr)?; + let mut bitmap_expr = None; + pruning.predicate_expr().apply(|node| { + if node.to_string().starts_with("KEY_RANGE_BITMAP(") { + bitmap_expr = Some(Arc::clone(node)); + return Ok(TreeNodeRecursion::Stop); + } + Ok(TreeNodeRecursion::Continue) + })?; + let bitmap_expr = bitmap_expr.expect("actual bitmap pruning expression"); + let remapped = bitmap_expr + .with_new_children(vec![physical_lit(0_i64), physical_lit(0_i64)])?; + drop(pruning); + drop(membership); + drop(current); + drop(batches); + drop(dynamic); + drop(join); + assert_eq!( + pool.reserved(), + bitmap_bytes, + "detached pruning expression retains only the bitmap charge" + ); + let actual = remapped.evaluate(&probe)?.into_array(probe.num_rows())?; + assert_eq!(actual.as_ref(), &BooleanArray::from(vec![true, true, true])); + let absent = remapped.with_new_children(vec![ + physical_lit(500_000_i64), + physical_lit(500_000_i64), + ])?; + let actual = absent.evaluate(&probe)?.into_array(probe.num_rows())?; + assert_eq!( + actual.as_ref(), + &BooleanArray::from(vec![false, false, false]) + ); + drop(absent); + assert_eq!(pool.reserved(), 0, "detached bitmap charge releases last"); + } + } + Ok(()) + } + #[test] fn remaining_row_group_count_reports_non_skipped_groups() { let mut filter = RowGroupAccessPlanFilter::new(ParquetAccessPlan::new_all(4)); @@ -814,6 +1043,731 @@ mod tests { assert_eq!(row_groups.is_fully_matched(), &vec![false, true, false]); } + #[test] + fn null_in_list_under_not_does_not_prove_full_match() { + use arrow::array::Int32Array; + use arrow::record_batch::RecordBatch; + use bytes::Bytes; + use datafusion_physical_expr::expressions::{ + col as physical_col, in_list, lit as physical_lit, + }; + use parquet::arrow::ArrowWriter; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use parquet::file::properties::WriterProperties; + + let schema = + Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![100, 1000]))], + ) + .unwrap(); + let mut bytes = Vec::new(); + let properties = WriterProperties::builder() + .set_max_row_group_row_count(Some(1)) + .set_bloom_filter_enabled(false) + .build(); + let mut writer = + ArrowWriter::try_new(&mut bytes, Arc::clone(&schema), Some(properties)) + .unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let bytes = Bytes::from(bytes); + let metadata = Arc::clone( + ParquetRecordBatchReaderBuilder::try_new(bytes.clone()) + .unwrap() + .metadata(), + ); + let column = physical_col("c1", &schema).unwrap(); + let mut list = (0..21) + .map(|value| physical_lit(ScalarValue::Int32(Some(value)))) + .collect::>(); + list.push(physical_lit(ScalarValue::Int32(None))); + let in_list = in_list(Arc::clone(&column), list, &false, &schema).unwrap(); + let not_in: Arc = Arc::new(NotExpr::new(in_list)); + let at_least_100: Arc = Arc::new(BinaryExpr::new( + Arc::clone(&column), + Operator::GtEq, + physical_lit(ScalarValue::Int32(Some(100))), + )); + let thousand: Arc = Arc::new(BinaryExpr::new( + column, + Operator::Eq, + physical_lit(ScalarValue::Int32(Some(1000))), + )); + // At c1=100, the first arm is UNKNOWN and the second is false. + // At c1=1000, the second arm is true, so LIMIT must not discard it. + let predicate: Arc = Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new(not_in, Operator::And, at_least_100)), + Operator::Or, + thousand, + )); + let pruning_predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .with_max_in_list_size(32) + .try_build(Arc::clone(&predicate)) + .unwrap(); + assert!(pruning_predicate.can_be_inverted_for_full_match()); + let metrics = parquet_file_metrics(); + let mut row_groups = RowGroupAccessPlanFilter::new(ParquetAccessPlan::new_all(2)); + row_groups.prune_by_statistics_with_metadata( + &schema, + &metadata, + &pruning_predicate, + &metrics, + ); + let full_match_flags = row_groups.is_fully_matched().clone(); + row_groups.prune_by_limit(1, metadata.row_groups(), &metrics); + let selected = row_groups.row_group_indexes().collect::>(); + let reader = ParquetRecordBatchReaderBuilder::try_new(bytes) + .unwrap() + .with_row_groups(selected.clone()) + .build() + .unwrap(); + let matching_values = reader + .flat_map(|batch| { + let batch = batch.unwrap(); + let filter = predicate + .evaluate(&batch) + .unwrap() + .into_array(batch.num_rows()) + .unwrap(); + let filter = filter.as_any().downcast_ref::().unwrap(); + let filtered = + arrow::compute::filter_record_batch(&batch, filter).unwrap(); + filtered + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect::>(); + assert_eq!( + (full_match_flags, selected, matching_values), + (vec![false, false], vec![0, 1], vec![1000]), + "UNKNOWN from NOT IN with a NULL member must not prove a full match" + ); + } + + /// A CASE arm can manufacture NULL even when every input column is non-null. + /// The first row is UNKNOWN and the second is TRUE; LIMIT must keep the latter. + #[tokio::test] + async fn computed_null_case_does_not_prove_full_match() -> Result<()> { + use crate::row_filter::build_row_filter; + use crate::source::ParquetSource; + use arrow::array::Int32Array; + use arrow::record_batch::RecordBatch; + use bytes::Bytes; + use datafusion_datasource::PartitionedFile; + use datafusion_datasource::file_scan_config::FileScanConfigBuilder; + use datafusion_datasource::source::DataSourceExec; + use datafusion_execution::TaskContext; + use datafusion_execution::object_store::ObjectStoreUrl; + use datafusion_physical_expr::expressions::{ + CaseExpr, TryCastExpr, col as physical_col, lit as physical_lit, + }; + use datafusion_physical_plan::filter::FilterExec; + use datafusion_physical_plan::limit::GlobalLimitExec; + use datafusion_physical_plan::{ExecutionPlan, collect, displayable}; + use parquet::arrow::ArrowWriter; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use parquet::file::properties::WriterProperties; + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("x", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1])), + Arc::new(Int32Array::from(vec![-200, 0])), + Arc::new(Int32Array::from(vec![1, 1])), + ], + )?; + let column = |name| physical_col(name, &schema).unwrap(); + let lit_i32 = |v| physical_lit(ScalarValue::Int32(Some(v))); + let when = Arc::new(BinaryExpr::new(column("a"), Operator::Eq, lit_i32(1))); + let cast = Arc::new(TryCastExpr::new(column("x"), DataType::Int8)); + let then = Arc::new(BinaryExpr::new( + cast, + Operator::Gt, + physical_lit(ScalarValue::Int8(Some(0))), + )); + let case = Arc::new(CaseExpr::try_new( + None, + vec![(when, then)], + Some(physical_lit(ScalarValue::Boolean(Some(false)))), + )?); + let predicate: Arc = Arc::new(BinaryExpr::new( + Arc::new(NotExpr::new(case)), + Operator::And, + Arc::new(BinaryExpr::new(column("b"), Operator::Gt, lit_i32(0))), + )); + let evaluated = predicate.evaluate(&batch)?.into_array(batch.num_rows())?; + assert_eq!( + evaluated.as_ref(), + &BooleanArray::from(vec![None, Some(true)]) + ); + + let properties = WriterProperties::builder() + .set_max_row_group_row_count(Some(1)) + .set_bloom_filter_enabled(false) + .build(); + let mut bytes = Vec::new(); + let mut writer = + ArrowWriter::try_new(&mut bytes, Arc::clone(&schema), Some(properties))?; + writer.write(&batch)?; + writer.close()?; + let bytes = Bytes::from(bytes); + let builder = ParquetRecordBatchReaderBuilder::try_new(bytes.clone())?; + let metadata = Arc::clone(builder.metadata()); + let pruning_predicate = + build_test_pruning_predicate(Arc::clone(&predicate), Arc::clone(&schema)); + let metrics = parquet_file_metrics(); + let mut groups = RowGroupAccessPlanFilter::new(ParquetAccessPlan::new_all(2)); + groups.prune_by_statistics_with_metadata( + &schema, + &metadata, + &pruning_predicate, + &metrics, + ); + let flags = groups.is_fully_matched().clone(); + groups.prune_by_limit(1, metadata.row_groups(), &metrics); + let selected = groups.row_group_indexes().collect::>(); + let row_filter = + build_row_filter(&predicate, &schema, &metadata, false, &metrics)? + .expect("real parquet reader filter"); + let reader = builder + .with_row_groups(selected.clone()) + .with_row_filter(row_filter) + .build()?; + let read_x = reader + .flat_map(|batch| { + let batch = batch.unwrap(); + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect::>(); + + // The same two rows can share a row group; max(x)=0 then masks the + // NULL produced by TRY_CAST(-200 AS Int8) in the inverse rewrite. + let mut mixed_bytes = Vec::new(); + let mut mixed_writer = + ArrowWriter::try_new(&mut mixed_bytes, Arc::clone(&schema), None)?; + mixed_writer.write(&batch)?; + mixed_writer.close()?; + let mixed_bytes = Bytes::from(mixed_bytes); + let mixed_builder = + ParquetRecordBatchReaderBuilder::try_new(mixed_bytes.clone())?; + let mixed_metadata = Arc::clone(mixed_builder.metadata()); + assert_eq!(mixed_metadata.num_row_groups(), 1); + let mut mixed_groups = + RowGroupAccessPlanFilter::new(ParquetAccessPlan::new_all(1)); + mixed_groups.prune_by_statistics_with_metadata( + &schema, + &mixed_metadata, + &pruning_predicate, + &metrics, + ); + let mixed_flags = mixed_groups.is_fully_matched().clone(); + mixed_groups.prune_by_limit(1, mixed_metadata.row_groups(), &metrics); + let mixed_selected = mixed_groups.row_group_indexes().collect::>(); + let mixed_filter = + build_row_filter(&predicate, &schema, &mixed_metadata, false, &metrics)? + .expect("mixed reader filter"); + let mixed_reader = mixed_builder + .with_row_groups(mixed_selected.clone()) + .with_row_filter(mixed_filter) + .build()?; + let mixed_read_x = mixed_reader + .flat_map(|batch| { + let batch = batch.unwrap(); + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect::>(); + eprintln!( + "mixed CASE flags={mixed_flags:?}, selected={mixed_selected:?}, reader={mixed_read_x:?}" + ); + // Additionally place a later matching row in a second row group. A + // LIMIT of two cannot be satisfied by the first row group alone. + let batch3 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 1])), + Arc::new(Int32Array::from(vec![-200, 0, 0])), + Arc::new(Int32Array::from(vec![1, 1, 1])), + ], + )?; + let properties3 = WriterProperties::builder() + .set_max_row_group_row_count(Some(2)) + .set_bloom_filter_enabled(false) + .build(); + let mut bytes3 = Vec::new(); + let mut writer3 = + ArrowWriter::try_new(&mut bytes3, Arc::clone(&schema), Some(properties3))?; + writer3.write(&batch3)?; + writer3.close()?; + let bytes3 = Bytes::from(bytes3); + let builder3 = ParquetRecordBatchReaderBuilder::try_new(bytes3.clone())?; + let metadata3 = Arc::clone(builder3.metadata()); + assert_eq!(metadata3.num_row_groups(), 2); + let mut groups3 = RowGroupAccessPlanFilter::new(ParquetAccessPlan::new_all(2)); + groups3.prune_by_statistics_with_metadata( + &schema, + &metadata3, + &pruning_predicate, + &metrics, + ); + let flags3 = groups3.is_fully_matched().clone(); + let before_limit3 = groups3.row_group_indexes().collect::>(); + groups3.prune_by_limit(2, metadata3.row_groups(), &metrics); + let selected3 = groups3.row_group_indexes().collect::>(); + let filter3 = build_row_filter(&predicate, &schema, &metadata3, false, &metrics)? + .expect("two-group reader filter"); + let reader3 = builder3 + .with_row_groups(selected3.clone()) + .with_row_filter(filter3) + .build()?; + let read3 = reader3 + .flat_map(|batch| { + let batch = batch.unwrap(); + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect::>(); + eprintln!( + "LIMIT2 before={before_limit3:?}, flags={flags3:?}, selected={selected3:?}, reader={read3:?}" + ); + + // Exercise the real Parquet DataSourceExec with row-filter pushdown both + // enabled and disabled. The external filter gives identical SQL semantics. + let file = tempfile::NamedTempFile::new()?; + std::fs::write(file.path(), &bytes)?; + let path = object_store::path::Path::from_filesystem_path(file.path()).unwrap(); + let mut exec_results = Vec::new(); + for pushdown in [false, true] { + let source = Arc::new( + ParquetSource::new(Arc::clone(&schema)) + .with_predicate(Arc::clone(&predicate)) + .with_pushdown_filters(pushdown) + .with_enable_page_index(false) + .with_bloom_filter_on_read(false), + ); + let config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + .with_file(PartitionedFile::new(path.to_string(), bytes.len() as u64)) + .build(); + let scan: Arc = DataSourceExec::from_data_source(config); + let plan: Arc = + Arc::new(FilterExec::try_new(Arc::clone(&predicate), scan)?); + let batches = collect(plan, Arc::new(TaskContext::default())).await?; + let values = batches + .iter() + .flat_map(|batch| { + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect::>(); + exec_results.push(values); + } + let file3 = tempfile::NamedTempFile::new()?; + std::fs::write(file3.path(), &bytes3)?; + let path3 = object_store::path::Path::from_filesystem_path(file3.path()).unwrap(); + let mut exec3 = Vec::new(); + for pushdown in [false, true] { + for limit in [None, Some(2)] { + let source = Arc::new( + ParquetSource::new(Arc::clone(&schema)) + .with_predicate(Arc::clone(&predicate)) + .with_pushdown_filters(pushdown) + .with_enable_page_index(false) + .with_bloom_filter_on_read(false), + ); + let config = FileScanConfigBuilder::new( + ObjectStoreUrl::local_filesystem(), + source, + ) + .with_file(PartitionedFile::new(path3.to_string(), bytes3.len() as u64)) + // A scan fetch is valid only when its row filter enforces the predicate. + .with_limit(if pushdown { limit } else { None }) + .build(); + let scan: Arc = + DataSourceExec::from_data_source(config); + let filtered: Arc = Arc::new(FilterExec::try_new( + Arc::clone(&predicate), + Arc::clone(&scan), + )?); + let plan: Arc = if limit.is_some() { + Arc::new(GlobalLimitExec::new(filtered, 0, limit)) + } else { + filtered + }; + eprintln!( + "pushdown={pushdown} limit={limit:?} plan:\n{}", + displayable(plan.as_ref()).indent(true) + ); + let batches = collect(plan, Arc::new(TaskContext::default())).await?; + if pushdown && limit.is_none() { + let metrics = scan.metrics().expect("scan metrics"); + let pruned = metrics + .sum_by_name("pushdown_rows_pruned") + .expect("row filter metric") + .as_usize(); + eprintln!("actual reader row-filter pruned={pruned}"); + assert!(pruned > 0, "configured row filter must actually run"); + } + let mut values = batches + .iter() + .flat_map(|batch| { + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect::>(); + values.sort_unstable(); + exec3.push(values); + } + } + eprintln!("LIMIT2 exec pushdown-off/on [no-limit, limit]: {exec3:?}"); + assert_eq!( + exec3, + vec![vec![0, 0]; 4], + "LIMIT must agree with the unbounded results with and without pushdown" + ); + assert_eq!( + (mixed_flags, mixed_selected, mixed_read_x), + (vec![false], vec![0], vec![0]) + ); + assert_eq!( + (before_limit3, flags3, selected3, read3), + (vec![0, 1], vec![false, false], vec![0, 1], vec![0, 0]) + ); + assert_eq!( + (flags, selected, read_x, exec_results), + ( + vec![false, false], + vec![0, 1], + vec![0], + vec![vec![0], vec![0]] + ), + "computed NULL must not prove a full match or discard the later true row" + ); + Ok(()) + } + + #[test] + fn safe_cast_case_computed_null_does_not_prove_full_match() -> Result<()> { + use arrow::array::Int32Array; + use arrow::compute::CastOptions; + use arrow::record_batch::RecordBatch; + use bytes::Bytes; + use datafusion_physical_expr::expressions::{ + CaseExpr, CastExpr, col as physical_col, lit as physical_lit, + }; + use parquet::arrow::ArrowWriter; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("x", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1])), + Arc::new(Int32Array::from(vec![-200, 0])), + Arc::new(Int32Array::from(vec![1, 1])), + ], + )?; + let column = |name| physical_col(name, &schema).unwrap(); + let i32_lit = |value| physical_lit(ScalarValue::Int32(Some(value))); + let cast = Arc::new(CastExpr::new( + column("x"), + DataType::Int8, + Some(CastOptions { + safe: true, + ..Default::default() + }), + )); + assert!( + !cast.nullable(&schema)?, + "safe CAST currently reports non-null for non-null input" + ); + let then: Arc = Arc::new(BinaryExpr::new( + cast, + Operator::Gt, + physical_lit(ScalarValue::Int8(Some(0))), + )); + assert!(!then.nullable(&schema)?); + let case = Arc::new(CaseExpr::try_new( + None, + vec![( + Arc::new(BinaryExpr::new(column("a"), Operator::Eq, i32_lit(1))), + then, + )], + Some(physical_lit(ScalarValue::Boolean(Some(false)))), + )?); + let predicate: Arc = Arc::new(BinaryExpr::new( + Arc::new(NotExpr::new(case)), + Operator::And, + Arc::new(BinaryExpr::new(column("b"), Operator::Gt, i32_lit(0))), + )); + let evaluated = predicate.evaluate(&batch)?.into_array(batch.num_rows())?; + assert_eq!( + evaluated.as_ref(), + &BooleanArray::from(vec![None, Some(true)]) + ); + let mut bytes = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut bytes, Arc::clone(&schema), None)?; + writer.write(&batch)?; + writer.close()?; + let reader = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(bytes))?; + let metadata = Arc::clone(reader.metadata()); + let pruning = build_test_pruning_predicate(predicate, Arc::clone(&schema)); + let mut groups = RowGroupAccessPlanFilter::new(ParquetAccessPlan::new_all(1)); + groups.prune_by_statistics_with_metadata( + &schema, + &metadata, + &pruning, + &parquet_file_metrics(), + ); + assert_eq!( + groups.is_fully_matched(), + &vec![false], + "safe CAST may create NULL despite nullable=false" + ); + Ok(()) + } + + #[test] + fn optimized_computed_null_cast_does_not_prove_full_match() -> Result<()> { + use arrow::array::Int32Array; + use arrow::record_batch::RecordBatch; + use bytes::Bytes; + use datafusion_physical_expr::expressions::{ + TryCastExpr, col as physical_col, lit as physical_lit, + }; + use parquet::arrow::ArrowWriter; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use parquet::file::properties::WriterProperties; + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("x", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 1])), + Arc::new(Int32Array::from(vec![-200, 0, 0])), + Arc::new(Int32Array::from(vec![1, 1, 1])), + ], + )?; + let column = |name| physical_col(name, &schema).unwrap(); + let int32 = |v| physical_lit(ScalarValue::Int32(Some(v))); + // The SQL optimizer eliminates CASE to this non-CASE expression. + let predicate: Arc = Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new(column("a"), Operator::NotEq, int32(1))), + Operator::Or, + Arc::new(BinaryExpr::new( + Arc::new(TryCastExpr::new(column("x"), DataType::Int8)), + Operator::LtEq, + physical_lit(ScalarValue::Int8(Some(0))), + )), + )), + Operator::And, + Arc::new(BinaryExpr::new(column("b"), Operator::Gt, int32(0))), + )); + let values = predicate.evaluate(&batch)?.into_array(batch.num_rows())?; + assert_eq!( + values.as_ref(), + &BooleanArray::from(vec![None, Some(true), Some(true)]) + ); + let properties = WriterProperties::builder() + .set_max_row_group_row_count(Some(2)) + .build(); + let mut bytes = Vec::new(); + let mut writer = + ArrowWriter::try_new(&mut bytes, Arc::clone(&schema), Some(properties))?; + writer.write(&batch)?; + writer.close()?; + let metadata = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(bytes))? + .metadata() + .clone(); + let pruning = build_test_pruning_predicate(predicate, Arc::clone(&schema)); + assert!(!pruning.can_be_inverted_for_full_match()); + let metrics = parquet_file_metrics(); + let mut groups = RowGroupAccessPlanFilter::new(ParquetAccessPlan::new_all(2)); + groups.prune_by_statistics_with_metadata(&schema, &metadata, &pruning, &metrics); + assert_eq!(groups.is_fully_matched(), &vec![false, false]); + groups.prune_by_limit(2, metadata.row_groups(), &metrics); + assert_eq!(groups.row_group_indexes().collect::>(), vec![0, 1]); + Ok(()) + } + + #[test] + fn total_casts_can_prove_full_match_and_satisfy_limit() -> Result<()> { + use arrow::array::Int32Array; + use arrow::compute::CastOptions; + use arrow::record_batch::RecordBatch; + use bytes::Bytes; + use datafusion_physical_expr::expressions::{ + CastExpr, TryCastExpr, col as physical_col, lit as physical_lit, + }; + use parquet::arrow::ArrowWriter; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use parquet::file::properties::WriterProperties; + + let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![1, 2]))], + )?; + let mut bytes = Vec::new(); + let properties = WriterProperties::builder() + .set_max_row_group_row_count(Some(1)) + .build(); + let mut writer = + ArrowWriter::try_new(&mut bytes, Arc::clone(&schema), Some(properties))?; + writer.write(&batch)?; + writer.close()?; + let metadata = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(bytes))? + .metadata() + .clone(); + for safe_cast in [false, true] { + for target in [DataType::Int32, DataType::Int64] { + let column = physical_col("x", &schema)?; + let cast: Arc = if safe_cast { + Arc::new(CastExpr::new( + column, + target.clone(), + Some(CastOptions { + safe: true, + ..Default::default() + }), + )) + } else { + Arc::new(TryCastExpr::new(column, target.clone())) + }; + let zero = match target { + DataType::Int32 => ScalarValue::Int32(Some(0)), + DataType::Int64 => ScalarValue::Int64(Some(0)), + _ => unreachable!(), + }; + let predicate = build_test_pruning_predicate( + Arc::new(BinaryExpr::new(cast, Operator::Gt, physical_lit(zero))), + Arc::clone(&schema), + ); + assert!( + predicate.can_be_inverted_for_full_match(), + "safe_cast={safe_cast} target={target}" + ); + let metrics = parquet_file_metrics(); + let mut groups = + RowGroupAccessPlanFilter::new(ParquetAccessPlan::new_all(2)); + groups.prune_by_statistics_with_metadata( + &schema, &metadata, &predicate, &metrics, + ); + assert_eq!( + groups.is_fully_matched(), + &vec![true, true], + "safe_cast={safe_cast} target={target}" + ); + groups.prune_by_limit(1, metadata.row_groups(), &metrics); + assert_eq!(groups.row_group_indexes().collect::>(), vec![0]); + assert_eq!(metrics.limit_pruned_row_groups.pruned(), 1); + } + } + Ok(()) + } + + #[test] + fn non_nullable_case_can_still_prove_full_match_and_satisfy_limit() -> Result<()> { + use arrow::array::Int32Array; + use arrow::record_batch::RecordBatch; + use bytes::Bytes; + use datafusion_physical_expr::expressions::{ + CaseExpr, col as physical_col, lit as physical_lit, + }; + use parquet::arrow::ArrowWriter; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use parquet::file::properties::WriterProperties; + let schema = Arc::new(Schema::new(vec![Field::new("b", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![1, 2]))], + )?; + let column = physical_col("b", &schema)?; + let zero = physical_lit(ScalarValue::Int32(Some(0))); + let when = Arc::new(BinaryExpr::new( + Arc::clone(&column), + Operator::Gt, + Arc::clone(&zero), + )); + let then = Arc::new(BinaryExpr::new(column, Operator::LtEq, zero)); + let case = Arc::new(CaseExpr::try_new( + None, + vec![(when, then)], + Some(physical_lit(ScalarValue::Boolean(Some(false)))), + )?); + let predicate = build_test_pruning_predicate( + Arc::new(NotExpr::new(case)), + Arc::clone(&schema), + ); + assert!(predicate.can_be_inverted_for_full_match()); + let mut bytes = Vec::new(); + let properties = WriterProperties::builder() + .set_max_row_group_row_count(Some(1)) + .build(); + let mut writer = + ArrowWriter::try_new(&mut bytes, Arc::clone(&schema), Some(properties))?; + writer.write(&batch)?; + writer.close()?; + let metadata = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(bytes))? + .metadata() + .clone(); + let metrics = parquet_file_metrics(); + let mut groups = RowGroupAccessPlanFilter::new(ParquetAccessPlan::new_all(2)); + groups + .prune_by_statistics_with_metadata(&schema, &metadata, &predicate, &metrics); + assert_eq!(groups.is_fully_matched(), &vec![true, true]); + groups.prune_by_limit(1, metadata.row_groups(), &metrics); + assert_eq!(groups.row_group_indexes().collect::>(), vec![0]); + Ok(()) + } + #[test] fn prune_by_limit_preserves_row_selection() { let field = PrimitiveTypeField::new("c1", PhysicalType::INT32); diff --git a/datafusion/datasource-parquet/src/statistics_order_tests.rs b/datafusion/datasource-parquet/src/statistics_order_tests.rs index 5081abf2cb47b..95291d48d384f 100644 --- a/datafusion/datasource-parquet/src/statistics_order_tests.rs +++ b/datafusion/datasource-parquet/src/statistics_order_tests.rs @@ -359,6 +359,31 @@ fn byte_array_order_preserves_matching_rows_at_every_pruning_level() { } } +#[test] +fn byte_array_order_does_not_fully_match_from_untrusted_bounds() { + for order in [ + StatisticsOrder::Modern, + StatisticsOrder::Deprecated, + StatisticsOrder::Missing, + StatisticsOrder::Unknown, + ] { + let file = TestFile::new(order); + // The first row group contains "az", which is below "aé" in + // Arrow's order, but its legacy minimum is "aé". + let (physical, predicate) = file.predicate(&col("s").gt_eq(lit("aé"))); + let plan = file.row_group_plan(&predicate); + assert!(plan.should_scan(0), "order={order:?}"); + assert!(!plan.is_fully_matched(0), "order={order:?}"); + assert!(plan.should_scan(2), "order={order:?}"); + assert_eq!( + plan.is_fully_matched(2), + matches!(order, StatisticsOrder::Modern | StatisticsOrder::Deprecated), + "order={order:?}", + ); + assert_eq!(file.matching_rows(&physical, plan), 5, "order={order:?}"); + } +} + #[test] fn large_string_in_list_preserves_rows_with_untrusted_page_order() { let max_in_list_size = MAX_IN_LIST_SIZE + 2; diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index e0d3c2488b8c3..b9e5cc665922e 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -3152,7 +3152,7 @@ mod tests { use crate::empty::EmptyExec; use crate::execution_plan::Boundedness; use crate::expressions::col; - use crate::filter::FilterExecBuilder; + use crate::filter::{FilterExec, FilterExecBuilder, batch_filter}; use crate::metrics::MetricValue; use crate::statistics::{StatisticsArgs, StatisticsContext}; use crate::test::TestMemoryExec; @@ -3183,7 +3183,7 @@ mod tests { use datafusion_functions_aggregate::count::count_udaf; use datafusion_functions_aggregate::first_last::{first_value_udaf, last_value_udaf}; use datafusion_functions_aggregate::median::median_udaf; - use datafusion_functions_aggregate::min_max::min_udaf; + use datafusion_functions_aggregate::min_max::{max_udaf, min_udaf}; use datafusion_functions_aggregate::sum::sum_udaf; use datafusion_physical_expr::Partitioning; use datafusion_physical_expr::PhysicalSortExpr; @@ -7837,6 +7837,218 @@ mod tests { } } + // The FilterExec uses a batch size of one so that the aggregate updates the + // shared predicate before the filter reads the next input batch. + async fn run_aggregate_dynamic_filter_null_case( + rows: &[(Option, Option, Option)], + max_a: bool, + ) -> Result<( + Vec, + Arc, + Arc, + )> { + let schema = Arc::new(Schema::new(vec![ + Field::new("ts", DataType::Int64, true), + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Int64, true), + ])); + let batches = rows + .iter() + .map(|(ts, a, b)| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from(vec![*ts])), + Arc::new(Int64Array::from(vec![*a])), + Arc::new(Int64Array::from(vec![*b])), + ], + ) + }) + .collect::, _>>()?; + let input = TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema), None)?; + let aggregates = [ + ("ts", min_udaf()), + ("a", if max_a { max_udaf() } else { min_udaf() }), + ("b", min_udaf()), + ] + .into_iter() + .map(|(name, function)| { + Ok(Arc::new( + AggregateExprBuilder::new(function, vec![col(name, &schema)?]) + .schema(Arc::clone(&schema)) + .alias(name) + .build()?, + )) + }) + .collect::>>()?; + let aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::default(), + aggregates, + vec![None; 3], + Arc::new(EmptyExec::new(Arc::clone(&schema))), + Arc::clone(&schema), + )?); + let mut produced = aggregate.dynamic_expressions_produced(); + assert_eq!(produced.len(), 1); + let dynamic_filter = (produced.pop().unwrap() + as Arc) + .downcast::() + .expect("aggregate should produce a dynamic filter"); + let filter = Arc::new( + FilterExecBuilder::new( + Arc::::clone(&dynamic_filter), + input, + ) + .with_batch_size(1) + .build()?, + ); + let aggregate = aggregate.replace_children( + vec![Arc::::clone(&filter)], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + let results = + collect(aggregate.execute(0, Arc::new(TaskContext::default()))?).await?; + Ok((results, filter, dynamic_filter)) + } + + fn assert_aggregate_dynamic_filter_null_result( + batches: &[RecordBatch], + expected: [Option; 3], + ) -> Result<()> { + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 1); + for (index, value) in expected.into_iter().enumerate() { + assert_eq!( + ScalarValue::try_from_array(batches[0].column(index), 0)?, + ScalarValue::Int64(value), + "aggregate {index}" + ); + } + Ok(()) + } + + #[tokio::test] + async fn aggregate_dynamic_filter_null_keeps_unknown_columns() -> Result<()> { + let (results, filter, dynamic_filter) = run_aggregate_dynamic_filter_null_case( + &[(Some(100), None, None), (Some(101), None, Some(2))], + false, + ) + .await?; + assert_aggregate_dynamic_filter_null_result( + &results, + [Some(100), None, Some(2)], + )?; + assert_eq!(filter.metrics().unwrap().output_rows(), Some(2)); + assert_eq!( + format!("{}", dynamic_filter.current()?), + format!("{}", lit(true)) + ); + Ok(()) + } + + #[tokio::test] + async fn aggregate_dynamic_filter_null_to_value_becomes_selective() -> Result<()> { + let (results, filter, dynamic_filter) = run_aggregate_dynamic_filter_null_case( + &[ + (Some(100), None, Some(1)), + (Some(99), Some(2), None), + // No aggregate can improve: the second batch must have tightened + // the filter before this batch is read. + (Some(101), Some(3), Some(2)), + // Only min(a) improves; it must be present in the predicate. + (Some(101), Some(1), Some(2)), + ], + false, + ) + .await?; + assert_aggregate_dynamic_filter_null_result( + &results, + [Some(99), Some(1), Some(1)], + )?; + assert_eq!(filter.metrics().unwrap().output_rows(), Some(3)); + let predicate = dynamic_filter.current()?; + assert_ne!(format!("{predicate}"), format!("{}", lit(true))); + let probe = RecordBatch::try_new( + filter.schema(), + vec![ + Arc::new(Int64Array::from(vec![101, 101, 101, 98])), + Arc::new(Int64Array::from(vec![0, 2, 2, 2])), + Arc::new(Int64Array::from(vec![2, 0, 2, 2])), + ], + )?; + assert_eq!(batch_filter(&probe, &predicate)?.num_rows(), 3); + Ok(()) + } + + #[tokio::test] + async fn aggregate_dynamic_filter_null_mixed_min_max() -> Result<()> { + let (results, filter, dynamic_filter) = run_aggregate_dynamic_filter_null_case( + &[(None, None, None), (None, None, None)], + true, + ) + .await?; + assert_aggregate_dynamic_filter_null_result(&results, [None, None, None])?; + assert_eq!(filter.metrics().unwrap().output_rows(), Some(2)); + assert_eq!( + format!("{}", dynamic_filter.current()?), + format!("{}", lit(true)) + ); + + let (results, filter, dynamic_filter) = run_aggregate_dynamic_filter_null_case( + &[(Some(100), None, None), (Some(101), Some(2), None)], + true, + ) + .await?; + assert_aggregate_dynamic_filter_null_result( + &results, + [Some(100), Some(2), None], + )?; + assert_eq!(filter.metrics().unwrap().output_rows(), Some(2)); + assert_eq!( + format!("{}", dynamic_filter.current()?), + format!("{}", lit(true)) + ); + + let (results, filter, dynamic_filter) = run_aggregate_dynamic_filter_null_case( + &[ + (Some(100), None, None), + (Some(101), Some(2), None), + (Some(101), Some(3), Some(2)), + (Some(101), Some(4), Some(3)), + (Some(101), Some(2), Some(3)), + ], + true, + ) + .await?; + assert_aggregate_dynamic_filter_null_result( + &results, + [Some(100), Some(4), Some(2)], + )?; + assert_eq!(filter.metrics().unwrap().output_rows(), Some(4)); + assert_ne!( + format!("{}", dynamic_filter.current()?), + format!("{}", lit(true)) + ); + Ok(()) + } + + #[tokio::test] + async fn aggregate_dynamic_filter_null_consumer_uses_updated_predicate() -> Result<()> + { + let (results, filter, _) = run_aggregate_dynamic_filter_null_case( + &[(Some(100), Some(2), Some(1)), (Some(101), Some(3), Some(2))], + false, + ) + .await?; + assert_aggregate_dynamic_filter_null_result( + &results, + [Some(100), Some(2), Some(1)], + )?; + assert_eq!(filter.metrics().unwrap().output_rows(), Some(1)); + Ok(()) + } + /// Test that [`AggregateExec::with_dynamic_filter_expr`] overrides the existing dynamic filter #[test] fn test_with_dynamic_filter() -> Result<()> { diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 96464c6691ec2..b1f677a8cb13b 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -2823,6 +2823,157 @@ mod tests { Ok(()) } + /// Published membership expressions must keep the build's reservation alive + /// after the originating stream and plan are gone, including remapped copies. + #[tokio::test] + async fn dynamic_membership_retains_build_reservation() -> Result<()> { + use crate::joins::hash_join::partitioned_hash_eval::HashTableLookupExpr; + use arrow::array::BooleanArray; + use datafusion_common::tree_node::TreeNode; + use datafusion_physical_expr::expressions::InListExpr; + + for mode in [PartitionMode::CollectLeft, PartitionMode::Partitioned] { + for (in_list, key_count) in [(false, 1), (true, 1), (true, 2)] { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int64, false), + Field::new("v", DataType::Int64, false), + ])); + let batch = |keys: Vec, values: Vec| -> Result { + Ok(RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from(keys)), + Arc::new(Int64Array::from(values)), + ], + )?) + }; + let build = batch(vec![0, 1_000_000], vec![10, 20])?; + let original_buffers = build + .columns() + .iter() + .map(|column| column.to_data().buffers()[0].as_ptr() as usize) + .collect::>(); + let probe = batch(vec![0, 1_000_000, 99], vec![10, 20, 10])?; + let left = TestMemoryExec::try_new_exec( + &[vec![build]], + Arc::clone(&schema), + None, + )?; + let right = TestMemoryExec::try_new_exec( + &[vec![probe.clone()]], + Arc::clone(&schema), + None, + )?; + let on = (0..key_count) + .map(|index| { + let key: PhysicalExprRef = + Arc::new(Column::new(schema.field(index).name(), index)); + (Arc::clone(&key), key) + }) + .collect(); + let (join, dynamic) = hash_join_with_dynamic_filter_and_mode( + left, + right, + on, + JoinType::Inner, + mode, + )?; + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 20)); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .build_arc()?; + let mut config = SessionConfig::new(); + if !in_list { + config + .options_mut() + .optimizer + .hash_join_inlist_pushdown_max_size = 0; + } + let task = Arc::new( + TaskContext::default() + .with_runtime(runtime) + .with_session_config(config), + ); + let stream = join.execute(0, task)?; + let batches = common::collect(stream).await?; + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); + let expression = dynamic.current()?; + let mut membership = None; + expression.apply(|node| { + if node.downcast_ref::().is_some() + || node.downcast_ref::().is_some() + { + membership = Some(Arc::clone(node)); + return Ok(TreeNodeRecursion::Stop); + } + Ok(TreeNodeRecursion::Continue) + })?; + let membership = membership.expect("dynamic membership expression"); + assert_eq!(membership.downcast_ref::().is_some(), in_list); + let bitmap = membership + .downcast_ref::() + .and_then(|lookup| lookup.pruning_bitmap().cloned()); + if !in_list { + assert!(bitmap.is_some(), "sparse keys should produce a bitmap"); + } + if in_list && key_count == 2 { + let list = membership.downcast_ref::().unwrap(); + let literal = list.list()[0].downcast_ref::().unwrap(); + let ScalarValue::Struct(value) = literal.value() else { + panic!("struct-key IN-list must retain real Struct literals"); + }; + let buffers = value + .columns() + .iter() + .map(|column| column.to_data().buffers()[0].as_ptr() as usize) + .collect::>(); + assert_eq!( + buffers, original_buffers, + "literals share producer buffers" + ); + } + let children = membership.children().into_iter().cloned().collect(); + let remapped = Arc::clone(&membership).with_new_children(children)?; + assert!(pool.reserved() > 0); + drop(batches); + drop(expression); + drop(dynamic); + drop(join); + let retained = pool.reserved(); + assert!( + retained > 0, + "{mode:?} in_list={in_list} keys={key_count}: detached expression lost producer charge" + ); + let result = remapped.evaluate(&probe)?.into_array(probe.num_rows())?; + assert_eq!( + result.as_any().downcast_ref::().unwrap(), + &BooleanArray::from(vec![true, true, false]) + ); + drop(membership); + assert_eq!( + pool.reserved(), + retained, + "remap shares producer charge once" + ); + drop(remapped); + if let Some(bitmap) = bitmap { + assert_eq!( + pool.reserved(), + bitmap.size(), + "bitmap alone retains only its own charge" + ); + assert!(bitmap.may_contain_range(0, 0)); + drop(bitmap); + } + assert_eq!(pool.reserved(), 0); + eprintln!( + "passed {mode:?} in_list={in_list} keys={key_count}, retained={retained}" + ); + } + } + Ok(()) + } + /// Runs a join whose 200 keys spread over 2M would size a pruning bitmap /// at the 128 KiB cap, and reports the bytes the build side reserved. /// @@ -6821,6 +6972,180 @@ mod tests { Ok(()) } + /// Ordinary LeftAnti joins preserve their NULL semantics in both build maps. + #[rstest] + #[tokio::test] + async fn left_anti_retains_null_equality_semantics( + #[values(false, true)] perfect_hash: bool, + ) -> Result<()> { + for null_equality in [ + NullEquality::NullEqualsNothing, + NullEquality::NullEqualsNull, + ] { + let left = build_table_two_cols( + ("a1", &vec![None, Some(1), Some(2)]), + ("b1", &vec![Some(10), Some(20), Some(30)]), + ); + let right = build_table_two_cols( + ("a2", &vec![None, Some(2)]), + ("b2", &vec![Some(40), Some(50)]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a2", &right.schema())?) as _, + )]; + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::LeftAnti, + None, + PartitionMode::CollectLeft, + null_equality, + false, + )?; + let output = + common::collect(join.execute(0, prepare_task_ctx(8192, perfect_hash))?) + .await?; + if null_equality == NullEquality::NullEqualsNull { + assert_batches_sorted_eq!( + [ + "+----+----+", + "| a1 | b1 |", + "+----+----+", + "| 1 | 20 |", + "+----+----+", + ], + &output + ); + } else { + assert_batches_sorted_eq!( + [ + "+----+----+", + "| a1 | b1 |", + "+----+----+", + "| | 10 |", + "| 1 | 20 |", + "+----+----+", + ], + &output + ); + } + } + Ok(()) + } + + /// Mutable visited-state charges must not follow a published build lease. + #[tokio::test] + async fn visited_bitmap_reservation_is_execution_local() -> Result<()> { + use crate::memory::MemoryStream; + let values: Vec = (0..1024).collect(); + let batch = build_table_i32(("a1", &values), ("b1", &values), ("c1", &values)); + let schema = batch.schema(); + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let metrics_set = ExecutionPlanMetricsSet::new(); + let data = collect_left_input( + HASH_JOIN_SEED.random_state().clone(), + Box::pin(MemoryStream::try_new(vec![batch], schema, None)?), + vec![Arc::new(Column::new("a1", 0))], + BuildProbeJoinMetrics::new(0, &metrics_set), + MemoryConsumer::new("visited lease").register(&pool), + true, + 1, + true, + Arc::new(ConfigOptions::default()), + NullEquality::NullEqualsNothing, + MetricBuilder::new(&metrics_set) + .counter(ARRAY_MAP_CREATED_COUNT_METRIC_NAME, 0), + BuildMode::Ordinary, + ) + .await?; + let producer_bytes = data.build.reservation.size(); + // Inspect the keepalive reservation to verify its execution-local charge. + #[expect(clippy::used_underscore_binding)] + let visited_bytes = data._visited_reservation.size(); + let bitmap_bytes = data.membership.bitmap_reserved_bytes(); + assert_eq!(visited_bytes, 1024 / 8); + assert_eq!( + pool.reserved(), + producer_bytes + visited_bytes + bitmap_bytes + ); + let producer = Arc::clone(&data.build); + drop(data); + assert_eq!( + pool.reserved(), + producer_bytes, + "visited and unshared pruning bitmap must release with execution" + ); + drop(producer); + assert_eq!(pool.reserved(), 0); + Ok(()) + } + + /// A build lease retains the batch, computed keys, map, and their reservation. + #[rstest] + #[tokio::test] + async fn build_lease_retains_computed_keys_and_map_reservation( + #[values(false, true)] perfect_hash: bool, + ) -> Result<()> { + use crate::memory::MemoryStream; + + let values: Vec = (0..1024).collect(); + let batch = build_table_i32(("a1", &values), ("b1", &values), ("c1", &values)); + let schema = batch.schema(); + let on: Vec = vec![Arc::new(BinaryExpr::new( + Arc::new(Column::new("a1", 0)), + Operator::Plus, + Arc::new(Literal::new(ScalarValue::Int32(Some(1)))), + ))]; + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let reservation = MemoryConsumer::new("build lease").register(&pool); + let metrics_set = ExecutionPlanMetricsSet::new(); + let metrics = BuildProbeJoinMetrics::new(0, &metrics_set); + let count = MetricBuilder::new(&metrics_set) + .counter(ARRAY_MAP_CREATED_COUNT_METRIC_NAME, 0); + let config = Arc::clone( + prepare_task_ctx(8192, perfect_hash) + .session_config() + .options(), + ); + let data = collect_left_input( + HASH_JOIN_SEED.random_state().clone(), + Box::pin(MemoryStream::try_new(vec![batch], schema, None)?), + on, + metrics, + reservation, + false, + 1, + false, + config, + NullEquality::NullEqualsNothing, + count, + BuildMode::Ordinary, + ) + .await?; + assert_eq!(matches!(data.map(), Map::ArrayMap(_)), perfect_hash); + assert!(!data.map().is_empty()); + let mut counter = RecordBatchMemoryCounter::new(); + let batch_bytes = counter.count_batch(data.batch()); + let key_bytes = counter.count_array(data.values()[0].as_ref()); + assert!(key_bytes > 0, "the computed key must own separate buffers"); + let reserved = pool.reserved(); + assert!( + reserved > batch_bytes + key_bytes, + "map storage must stay charged" + ); + let lease = Arc::clone(&data.build); + drop(data); + assert_eq!(pool.reserved(), reserved); + assert!(!lease.map.is_empty()); + assert_eq!(lease.values[0].len(), 1024); + drop(lease); + assert_eq!(pool.reserved(), 0); + Ok(()) + } + fn build_table_struct( struct_name: &str, field_name_and_values: (&str, &Vec>), diff --git a/datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs b/datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs index 48281e5f481b1..840ce6f9664fa 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs @@ -381,8 +381,32 @@ async fn prepared_build_reuses_data_with_independent_dynamic_filters() -> Result assert!(futures::poll!(Box::pin(dynamic.filter.wait_complete())).is_ready()); } assert_eq!(pool.reserved(), bytes); + let published = plans + .iter() + .map(|plan| { + plan.dynamic_filter + .as_ref() + .unwrap() + .filter + .current() + .unwrap() + }) + .collect::>(); drop(plans); drop(prepared); + assert_eq!( + pool.reserved(), + bytes, + "two published consumers share one prepared build charge" + ); + let mut published = published.into_iter(); + drop(published.next()); + assert_eq!( + pool.reserved(), + bytes, + "remaining consumer keeps producer charge" + ); + drop(published); assert_eq!(pool.reserved(), 0); Ok(()) } @@ -764,7 +788,7 @@ async fn prepared_inlist_shares_accounted_build_buffers() -> Result<()> { ] { let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 20)); let prepared = prepare(&base, batches.clone(), Arc::clone(&pool)).await?; - let PushdownStrategy::InList(values) = &prepared.build.membership else { + let PushdownStrategy::InList(values) = &prepared.membership else { panic!("expected IN-list membership"); }; let columns = prepared.build.batch.columns(); @@ -797,7 +821,7 @@ async fn prepared_inlist_shares_accounted_build_buffers() -> Result<()> { Arc::new(config), ) .await?; - let PushdownStrategy::Map(_, bitmap) = &map.build.membership else { + let PushdownStrategy::Map(_, bitmap) = &map.membership else { panic!("expected hash-map membership"); }; let bitmap_bytes = bitmap.as_ref().map_or(0, |bitmap| bitmap.size()); @@ -891,10 +915,7 @@ async fn prepared_byte_keys_use_hash_membership() -> Result<()> { ) .await?; assert!(prepared.build.bounds.is_none()); - assert!(matches!( - &prepared.build.membership, - PushdownStrategy::Map(_, _) - )); + assert!(matches!(&prepared.membership, PushdownStrategy::Map(_, _))); let task = with_probe_filter(base)? .builder() .with_prepared_build(prepared) diff --git a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs index 02eb96af97dfd..fdb877323b2f5 100644 --- a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs +++ b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs @@ -931,6 +931,81 @@ mod tests { expressions::{Column, Literal}, }; + fn test_build_lease( + pool: &Arc, + ) -> Arc { + use crate::joins::join_hash_map::JoinHashMapU32; + use datafusion_execution::memory_pool::MemoryConsumer; + let reservation = MemoryConsumer::new("publication test").register(pool); + reservation.try_grow(128).unwrap(); + let batch = RecordBatch::try_new( + test_probe_schema(), + vec![Arc::new(Int32Array::from(vec![1, 2]))], + ) + .unwrap(); + Arc::new(JoinBuildData { + map: Arc::new(Map::HashMap(Box::new(JoinHashMapU32::with_capacity(2)))), + values: vec![Arc::clone(batch.column(0))], + batch, + bounds: None, + reservation, + }) + } + + #[test] + fn publication_constructor_error_releases_producer_lease() -> Result<()> { + use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryPool}; + let pool: Arc = Arc::new(GreedyMemoryPool::new(1024)); + let producer = test_build_lease(&pool); + let array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2])); + let invalid_schema = + Schema::new(vec![Field::new("probe_key", DataType::Utf8, false)]); + let result = create_membership_predicate( + &test_on_right(), + PushdownStrategy::InList(array), + Some(Arc::clone(&producer)), + &HASH_JOIN_SEED, + &invalid_schema, + ); + assert!(result.is_err()); + assert_eq!( + pool.reserved(), + 128, + "error drops temporary lease but not original owner" + ); + drop(producer); + assert_eq!(pool.reserved(), 0); + Ok(()) + } + + #[tokio::test] + async fn partitioned_accumulator_retains_reported_build_until_drop() -> Result<()> { + use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryPool}; + let pool: Arc = Arc::new(GreedyMemoryPool::new(1024)); + let producer = test_build_lease(&pool); + let array: ArrayRef = Arc::clone(producer.batch.column(0)); + let acc = make_partitioned_expr_accumulator_for_test(2); + let mut report = + Box::pin(acc.report_build_data(PartitionBuildData::Partitioned { + partition_id: 0, + pushdown: PushdownStrategy::InList(array), + build_lease: Some(Arc::clone(&producer)), + bounds: no_bounds(), + keys_have_null: false, + })); + assert!(futures::poll!(&mut report).is_pending()); + drop(report); + drop(producer); + assert_eq!( + pool.reserved(), + 128, + "accumulator must retain an already reported partition after the stream is dropped" + ); + drop(acc); + assert_eq!(pool.reserved(), 0); + Ok(()) + } + fn test_on_right() -> Vec { vec![Arc::new(Column::new("probe_key", 0))] } @@ -1011,6 +1086,7 @@ mod tests { fn reported(pushdown: PushdownStrategy, bounds: PartitionBounds) -> PartitionStatus { PartitionStatus::Reported(PartitionData { pushdown, + build_lease: None, bounds, keys_have_null: false, }) @@ -1385,6 +1461,7 @@ mod tests { PartitionBuildData::Partitioned { partition_id: 0, pushdown: PushdownStrategy::Empty, + build_lease: None, bounds: PartitionBounds::new(vec![]), keys_have_null: false, }, diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs index 302ffb77dc528..3b5185909c7c7 100644 --- a/datafusion/physical-plan/src/joins/hash_join/stream.rs +++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs @@ -1084,6 +1084,7 @@ mod tests { PartitionBuildData::Partitioned { partition_id, pushdown: PushdownStrategy::Empty, + build_lease: None, bounds: PartitionBounds::new(vec![]), keys_have_null: false, } diff --git a/datafusion/pruning/src/pruning_predicate.rs b/datafusion/pruning/src/pruning_predicate.rs index 158e74a54aaed..4735fa375ced8 100644 --- a/datafusion/pruning/src/pruning_predicate.rs +++ b/datafusion/pruning/src/pruning_predicate.rs @@ -3608,6 +3608,85 @@ mod tests { assert_eq!(result, expected); } + #[test] + fn dynamic_inlist_snapshot_keeps_producer_lease_through_pruning() -> Result<()> { + use arrow::array::StructArray; + use datafusion_physical_expr::expressions::{InListExpr, Literal, NotExpr}; + use std::sync::atomic::{AtomicUsize, Ordering}; + + struct Producer { + columns: Vec, + live: Arc, + bytes: usize, + } + impl Drop for Producer { + fn drop(&mut self) { + self.live.fetch_sub(self.bytes, Ordering::SeqCst); + } + } + let live = Arc::new(AtomicUsize::new(128)); + let fields = vec![ + Arc::new(Field::new("c0", DataType::Int32, true)), + Arc::new(Field::new("c1", DataType::Int32, true)), + ]; + let columns: Vec = vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(Int32Array::from(vec![3, 4])), + ]; + let producer = Arc::new(Producer { + columns, + live: Arc::clone(&live), + bytes: 128, + }); + let values = Arc::new(StructArray::from( + fields + .into_iter() + .zip(producer.columns.iter().cloned()) + .collect::>(), + )) as ArrayRef; + let schema = Arc::new(Schema::new(vec![Field::new( + "k", + values.data_type().clone(), + false, + )])); + let column: PhysicalExprRef = Arc::new(phys_expr::Column::new("k", 0)); + let list: PhysicalExprRef = Arc::new( + InListExpr::try_new_from_array(Arc::clone(&column), values, false, &schema)? + .with_producer_lease(producer.clone()), + ); + assert!( + list.downcast_ref::() + .unwrap() + .list() + .iter() + .all(|item| item.downcast_ref::().is_some()) + ); + let negated: PhysicalExprRef = Arc::new(NotExpr::new(list)); + let dynamic: PhysicalExprRef = + Arc::new(DynamicFilterPhysicalExpr::new(vec![column], negated)); + let pruning = PruningPredicateBuilder::new() + .with_file_schema(schema) + .try_build(Arc::clone(&dynamic))?; + assert!( + pruning + .orig_expr() + .downcast_ref::() + .unwrap() + .negated(), + "snapshot must exercise NOT(IN) simplification" + ); + drop(dynamic); + drop(producer); + assert_eq!( + live.load(Ordering::SeqCst), + 128, + "pruning predicate lost producer lease during simplification" + ); + drop(pruning); + assert_eq!(live.load(Ordering::SeqCst), 0); + Ok(()) + } + #[test] fn row_group_predicate_lt_bool() -> Result<()> { let schema = Schema::new(vec![Field::new("c1", DataType::Boolean, false)]); @@ -7630,6 +7709,7 @@ mod tests { .try_build(case) .unwrap(); + assert!(predicate.can_be_inverted_for_full_match()); let statistics = TestStatistics::new().with( "b", ContainerStats::new_i32( @@ -7644,4 +7724,232 @@ mod tests { // other two intersect neither branch, nor the `ELSE`. assert_eq!(result, vec![false, true, true, false]); } + #[test] + fn computed_null_casts_block_full_match_only_when_not_total() -> Result<()> { + use arrow::compute::CastOptions; + let schema = Arc::new(Schema::new(vec![ + Field::new("x", DataType::Int32, false), + Field::new("d", DataType::Decimal128(10, 2), false), + Field::new("t", DataType::Timestamp(TimeUnit::Second, None), false), + ])); + let column = |name, index| { + Arc::new(phys_expr::Column::new(name, index)) as PhysicalExprRef + }; + let try_cast = |expr, target| { + Arc::new(phys_expr::TryCastExpr::new(expr, target)) as PhysicalExprRef + }; + let safe_cast = |expr, target| { + Arc::new(phys_expr::CastExpr::new( + expr, + target, + Some(CastOptions { + safe: true, + ..Default::default() + }), + )) as PhysicalExprRef + }; + let compare = |expr: PhysicalExprRef, value: ScalarValue| { + Arc::new(phys_expr::BinaryExpr::new( + expr, + Operator::Gt, + Arc::new(phys_expr::Literal::new(value)), + )) as PhysicalExprRef + }; + let build = |expr| { + PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(expr) + }; + let wide = |v| ScalarValue::Int64(Some(v)); + let narrow = |v| ScalarValue::Int8(Some(v)); + let i32_value = |v| ScalarValue::Int32(Some(v)); + for (label, expr, expected) in [ + ( + "TRY narrowing", + compare(try_cast(column("x", 0), DataType::Int8), narrow(0)), + false, + ), + ( + "safe narrowing", + compare(safe_cast(column("x", 0), DataType::Int8), narrow(0)), + false, + ), + ( + "TRY identity", + compare(try_cast(column("x", 0), DataType::Int32), i32_value(0)), + true, + ), + ( + "safe identity", + compare(safe_cast(column("x", 0), DataType::Int32), i32_value(0)), + true, + ), + ( + "TRY widening", + compare(try_cast(column("x", 0), DataType::Int64), wide(0)), + true, + ), + ( + "safe widening", + compare(safe_cast(column("x", 0), DataType::Int64), wide(0)), + true, + ), + ( + "nested unsafe below total", + compare( + try_cast(try_cast(column("x", 0), DataType::Int8), DataType::Int64), + wide(0), + ), + false, + ), + ( + "decimal scale change", + compare( + try_cast(column("d", 1), DataType::Decimal128(12, 3)), + ScalarValue::Decimal128(Some(0), 12, 3), + ), + false, + ), + ( + "timestamp unit change", + compare( + try_cast( + column("t", 2), + DataType::Timestamp(TimeUnit::Nanosecond, None), + ), + ScalarValue::TimestampNanosecond(Some(0), None), + ), + false, + ), + ( + "unknown source type", + compare(try_cast(column("missing", 99), DataType::Int8), narrow(0)), + false, + ), + ] { + assert_eq!( + build(expr)?.can_be_inverted_for_full_match(), + expected, + "{label}" + ); + } + Ok(()) + } + + #[test] + fn case_result_nullability_controls_full_match_but_not_forward_pruning() -> Result<()> + { + use arrow::compute::CastOptions; + let schema = Arc::new(Schema::new(vec![ + Field::new("when", DataType::Boolean, true), + Field::new("b", DataType::Int32, false), + Field::new("nullable", DataType::Boolean, true), + ])); + let col = |name, index| { + Arc::new(phys_expr::Column::new(name, index)) as PhysicalExprRef + }; + let boolean = |value| { + Arc::new(phys_expr::Literal::new(ScalarValue::Boolean(value))) + as PhysicalExprRef + }; + let false_lit = || boolean(Some(false)); + let when = || col("when", 0); + let compare = |left: PhysicalExprRef| { + Arc::new(phys_expr::BinaryExpr::new( + left, + Operator::Gt, + Arc::new(phys_expr::Literal::new(ScalarValue::Int32(Some(0)))), + )) as PhysicalExprRef + }; + let case = |then: PhysicalExprRef, + otherwise: Option| + -> Result { + Ok(Arc::new(phys_expr::CaseExpr::try_new( + None, + vec![(when(), then)], + otherwise, + )?)) + }; + let build = |expr| { + PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(expr) + }; + let nonnull = compare(col("b", 1)); + let nullable = compare(Arc::new(phys_expr::TryCastExpr::new( + col("b", 1), + DataType::Int32, + ))); + let safe_cast: PhysicalExprRef = Arc::new(phys_expr::BinaryExpr::new( + Arc::new(phys_expr::CastExpr::new( + col("b", 1), + DataType::Int8, + Some(CastOptions { + safe: true, + ..Default::default() + }), + )), + Operator::Gt, + Arc::new(phys_expr::Literal::new(ScalarValue::Int8(Some(0)))), + )); + assert!(!safe_cast.nullable(&schema)?); + let nested = case(Arc::clone(&nullable), Some(false_lit()))?; + let nullable_else = col("nullable", 2); + let invalid_column = col("missing", 99); + assert!(invalid_column.nullable(&schema).is_err()); + + for (label, then, otherwise, eligible) in [ + ( + "nonnull result, nullable WHEN", + Arc::clone(&nonnull), + Some(false_lit()), + true, + ), + ( + "nullable THEN", + Arc::clone(&nullable), + Some(false_lit()), + false, + ), + ( + "nullable ELSE", + Arc::clone(&nonnull), + Some(nullable_else), + false, + ), + ("nested nullable CASE", nested, Some(false_lit()), false), + ( + "safe CAST false-nullable", + safe_cast, + Some(false_lit()), + false, + ), + ( + "nullable error", + Arc::clone(&nonnull), + Some(invalid_column), + false, + ), + ] { + let predicate = build(case(then, otherwise)?)?; + assert_eq!( + predicate.can_be_inverted_for_full_match(), + eligible, + "{label}" + ); + } + // Missing ELSE retains the upstream conservative fallback. + let no_else = build(case(Arc::clone(&nonnull), None)?)?; + assert!(no_else.always_true()); + let stats = TestStatistics::new().with( + "b", + ContainerStats::new_i32(vec![Some(-2), Some(1)], vec![Some(-1), Some(2)]) + .with_null_counts([Some(0), Some(0)]) + .with_row_counts([Some(1), Some(1)]), + ); + let forward = build(case(nullable, Some(false_lit()))?)?; + assert!(!forward.can_be_inverted_for_full_match()); + assert_eq!(forward.prune(&stats)?, [false, true]); + Ok(()) + } } From 9df12a46e2dcccdccea3e6d955d208894a40b79f Mon Sep 17 00:00:00 2001 From: discord9 Date: Fri, 6 Nov 2026 21:25:03 +0800 Subject: [PATCH 19/19] test: align dynamic row-group pruning SLT metrics snapshot Record the final aggregate/bitmap plan metrics for the existing local regression, without changing product behavior. Signed-off-by: discord9 --- .../sqllogictest/test_files/dynamic_row_group_pruning.slt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt b/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt index 7c2e0f6baf4ca..963a4037a44aa 100644 --- a/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt +++ b/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt @@ -298,7 +298,7 @@ explain analyze select rgsel.b from dim join rgsel on dim.column1 = rgsel.b; Plan with Metrics 01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(column1@0, b@0)], projection=[b@1], metrics=[output_rows=3, elapsed_compute=, output_bytes=, output_batches=1, build_mem_used=, array_map_created_count=0, build_input_batches=1, build_input_rows=3, input_batches=2, input_rows=3, build_time=, join_time=, avg_fanout=100% (3/3), probe_hit_rate=100% (3/3)] 02)--DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] -03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_row_group_pruning/rgsel.parquet]]}, projection=[b], file_type=parquet, predicate=DynamicFilter [ b@1 >= 1000 AND b@1 <= 3050 AND hash_lookup ], dynamic_rg_pruning=eligible, pruning_predicate=b_null_count@1 != row_count@2 AND b_max@0 >= 1000 AND b_null_count@1 != row_count@2 AND b_min@3 <= 3050 AND b_null_count@1 != row_count@2 AND KEY_RANGE_BITMAP(b_min@3, b_max@0, 3/4096 buckets), required_guarantees=[], metrics=[output_rows=3, elapsed_compute=, output_bytes=, output_batches=2, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=4 total → 2 matched, row_groups_pruned_bloom_filter=2 total → 2 matched, page_index_pages_pruned=20 total → 3 matched, page_index_rows_pruned=200 total → 30 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, bytes_processed=, bytes_scanned=, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=3, pushdown_rows_pruned=27, row_groups_pruned_dynamic_filter=0, predicate_cache_inner_records=200, predicate_cache_records=13, bloom_filter_eval_time=, metadata_load_time=, page_index_eval_time=, row_pushdown_eval_time=, statistics_eval_time=, time_elapsed_opening=, time_elapsed_processing=, time_elapsed_scanning_total=, time_elapsed_scanning_until_data=, output_rows_skew=, scan_efficiency_ratio=] +03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_row_group_pruning/rgsel.parquet]]}, projection=[b], file_type=parquet, predicate=DynamicFilter [ b@1 >= 1000 AND b@1 <= 3050 AND hash_lookup ], dynamic_rg_pruning=eligible, pruning_predicate=b_null_count@1 != row_count@2 AND b_max@0 >= 1000 AND b_null_count@1 != row_count@2 AND b_min@3 <= 3050 AND b_null_count@1 != row_count@2 AND KEY_RANGE_BITMAP(b_min@3, b_max@0, 3/4096 buckets), required_guarantees=[], metrics=[output_rows=3, elapsed_compute=, output_bytes=, output_batches=2, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=4 total → 2 matched, row_groups_pruned_bloom_filter=2 total → 2 matched, page_index_pages_pruned=20 total → 3 matched, page_index_rows_pruned=200 total → 30 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, bytes_scanned=, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=3, pushdown_rows_pruned=27, row_groups_pruned_dynamic_filter=0, predicate_cache_inner_records=200, predicate_cache_records=, bloom_filter_eval_time=, metadata_load_time=, page_index_eval_time=, row_pushdown_eval_time=, statistics_eval_time=, time_elapsed_opening=, time_elapsed_processing=, time_elapsed_scanning_total=, time_elapsed_scanning_until_data=, output_rows_skew=, scan_efficiency_ratio=] statement ok RESET datafusion.optimizer.hash_join_inlist_pushdown_max_distinct_values;