diff --git a/Cargo.lock b/Cargo.lock index e5c5691f5d03c..9f0fa41743af7 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", @@ -2599,6 +2600,7 @@ dependencies = [ "insta", "itertools 0.15.0", "log", + "rstest", ] [[package]] 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/config.rs b/datafusion/common/src/config.rs index f5742f09f9b08..9f596ceefb7a7 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1282,13 +1282,19 @@ 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 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/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/common/src/utils/memory.rs b/datafusion/common/src/utils/memory.rs index 21c084119e120..4522ef35c95a5 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,237 @@ 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 + } + + /// 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 } -} -/// 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(); + } + } + + /// 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()), + } + } + + 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>>, +} - 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(); +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; + } - // Count all children `ArrayData` recursively - for child in array_data.child_data() { - count_array_data_memory_size(child, counted_buffers, total_size); + 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 +448,62 @@ 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, StructArray, new_null_array, + }; + use arrow::buffer::{NullBuffer, OffsetBuffer}; + 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_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![ @@ -359,6 +612,238 @@ 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) + .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![ 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/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..f2646855713ce --- /dev/null +++ b/datafusion/core/tests/parquet/string_in_list_pruning.rs @@ -0,0 +1,834 @@ +// 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 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::{ + 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; +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::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; + +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; + +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 +/// 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") + .tempfile() + .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Utf8, + false, + )])); + 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 + } else { + ROWS_PER_UNIT + }; + 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); + 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_file_layout(&metadata, total_rows, rows_per_group); + 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 + ); + 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); + } +} + +async fn scan( + file: &NamedTempFile, + list_size: usize, + max_in_list_size: Option, + page_pruning: bool, + negated: 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 op = if negated { "NOT IN" } else { "IN" }; + let sql = format!( + "SELECT value, count(*) AS n FROM t \ + WHERE value {op} ({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, + } +} + +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] { + // 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, 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, false).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, false).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); +} + +/// 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 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; +} + +#[tokio::test] +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; +} + +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("in_list_null_pruning") + .suffix(".parquet") + .tempfile() + .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) + .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 + // NOT IN (..., NULL) to an empty relation before the scan. + 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(), + file.as_file().metadata().unwrap().len(), + ); + let ctx = + SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); + + 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"), "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, + plan: plan_text, + metrics: MetricsFinder::find_metrics(plan.as_ref()).unwrap(), + }; + + if negated { + assert!( + output.batches.iter().all(|batch| batch.num_rows() == 0), + "type={name}, cap={max_in_list_size}" + ); + } else { + 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!( + output.counter("pushdown_rows_pruned"), + match (negated, max_in_list_size) { + (false, 0) => 7, + (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!( + output.pruned("row_groups_pruned_statistics"), + if max_in_list_size == 0 { + 0 + } else if negated { + 2 + } else { + 1 + }, + "type={name}, 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); + } + } +} + +#[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/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..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 { @@ -1105,10 +1108,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, ); @@ -1689,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, )) } @@ -1966,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 6bc1aca667981..236a08c5608af 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::{ @@ -31,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; @@ -137,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, @@ -490,6 +501,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 +541,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 +555,7 @@ impl<'a> PagesPruningStatistics<'a> { column_index, offset_index, page_offsets, + trusted_min_max, }) } @@ -563,6 +582,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 +599,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..56b8ce08b531f 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,12 @@ 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 = Arc::new(NotExpr::new(Arc::clone(predicate.orig_expr()))); @@ -375,13 +430,21 @@ 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; }; + // 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, + 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 +526,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 +534,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 +546,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 { @@ -603,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)); @@ -719,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/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 new file mode 100644 index 0000000000000..95291d48d384f --- /dev/null +++ b/datafusion/datasource-parquet/src/statistics_order_tests.rs @@ -0,0 +1,842 @@ +// 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_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; + 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 [ + 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)); +} 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/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-expr/src/projection.rs b/datafusion/physical-expr/src/projection.rs index 0e8876f017379..21e42903ada38 100644 --- a/datafusion/physical-expr/src/projection.rs +++ b/datafusion/physical-expr/src/projection.rs @@ -2490,6 +2490,80 @@ 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]); + 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 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/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/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 08d209003ad91..b1f677a8cb13b 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, @@ -69,15 +70,18 @@ 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; 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}; +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, @@ -101,14 +105,23 @@ 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); 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,8 +130,8 @@ fn try_create_array_map( perfect_hash_join_small_build_threshold: usize, perfect_hash_join_min_key_density: f64, null_equality: NullEquality, -) -> Result)>> { - if on_left.len() != 1 { +) -> Result> { + if !should_collect_min_max_for_perfect_hash(on_left, schema)? { return Ok(None); } @@ -182,56 +195,49 @@ 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((min_val, max_val))) +} - Ok(Some((array_map, batch, left_values))) +/// Immutable build buffers and their durable reservation. +pub(super) struct JoinBuildData { + /// Hash table with row indices into `batch`, also shared with dynamic filters. + pub(super) map: Arc, + /// The input rows for the build side. + pub(super) batch: RecordBatch, + /// Evaluated build-side key expressions. + pub(super) values: Vec, + /// Bounds computed from the build side; absent for an empty partition. + pub(super) bounds: Option, + // Declared last so accounted-for build allocations are released first. + pub(super) reservation: MemoryReservation, } -/// HashTable and input data for the left (build side) of a join +/// A build lease and mutable bookkeeping for one join execution. 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 - batch: RecordBatch, - /// The build side on expressions values - values: Vec, - /// Shared bitmap builder for visited left indices + /// 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, - /// 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 + /// 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, + /// 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, } 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. @@ -253,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,6 +272,15 @@ impl JoinLeftData { &self.membership } + pub(super) fn build_lease(&self) -> Arc { + Arc::clone(&self.build) + } + + /// 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` /// if caller is the last running thread pub(super) fn report_probe_completed(&self) -> bool { @@ -306,6 +321,7 @@ impl HashJoinExecBuilder { filter: None, join_type, left_fut: Default::default(), + prepared_build: None, random_state: HASH_JOIN_SEED, mode: PartitionMode::Auto, fetch: None, @@ -397,13 +413,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; @@ -419,10 +442,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(); @@ -451,6 +482,7 @@ impl HashJoinExecBuilder { filter, join_type, left_fut, + prepared_build, random_state, mode, metrics, @@ -498,6 +530,7 @@ impl HashJoinExecBuilder { join_type, join_schema, left_fut, + prepared_build, random_state, mode, metrics, @@ -528,6 +561,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(), @@ -572,7 +606,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, @@ -757,6 +791,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 @@ -800,6 +836,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) @@ -1226,6 +1263,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() @@ -1234,7 +1277,7 @@ impl DisplayAs for HashJoinExec { .join(", "); write!( f, - "HashJoinExec: mode={:?}, join_type={:?}, on=[{}]{}{}{}{}{}", + "HashJoinExec: mode={:?}, join_type={:?}, on=[{}]{}{}{}{}{}{}", self.mode, self.join_type, on, @@ -1243,6 +1286,7 @@ impl DisplayAs for HashJoinExec { display_null_equality, display_fetch, display_null_aware, + display_prepared, ) } DisplayFormatType::TreeRender => { @@ -1261,6 +1305,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")?; } @@ -1486,8 +1534,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 = @@ -1500,14 +1556,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 = @@ -1525,9 +1582,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 @@ -1610,11 +1668,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, @@ -1637,6 +1701,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 { @@ -1844,6 +1911,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())?; @@ -2145,6 +2217,12 @@ impl CollectLeftAccumulator { } } +#[derive(PartialEq, Eq)] +enum BuildMode { + Ordinary, + Prepared, +} + /// State for collecting the build-side data during hash join struct BuildSideState { batches: Vec, @@ -2185,6 +2263,9 @@ impl BuildSideState { } } +/// 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, @@ -2198,6 +2279,112 @@ 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))) + } +} + +/// 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: @@ -2216,6 +2403,13 @@ fn should_collect_min_max_for_perfect_hash( /// * `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 @@ -2229,7 +2423,7 @@ fn should_collect_min_max_for_perfect_hash( #[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, @@ -2239,13 +2433,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(), @@ -2253,41 +2449,53 @@ 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: _, } = state; + // 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 { Some(accumulators) if num_rows > 0 => { @@ -2300,8 +2508,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, @@ -2311,80 +2519,113 @@ 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()); - - (Map::ArrayMap(array_map), batch, left_value) + let batch = if prepared { + concat_batches(&schema, batches.iter())? } 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)?; - 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)?; - reservation.try_grow(estimated_hashtable_size)?; - metrics.build_mem_used.add(estimated_hashtable_size); - Box::new(JoinHashMapU32::with_capacity(num_rows)) - }; - - 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(); - } + 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)?; - // Merge all batches into a single batch, so we can directly index into the arrays - let batch = concat_batches(&schema, batches_iter.clone())?; + array_map_created_count.add(1); + metrics.build_mem_used.add(array_map.size()); - let left_values = evaluate_expressions_to_arrays(&on_left, &batch)?; + (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 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 + 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(); + } - (Map::HashMap(hashmap), batch, left_values) + // Merge all batches into a single batch, so we can directly index into the arrays + 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)?; + + (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 { + 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); - 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. @@ -2394,19 +2635,46 @@ 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(), + ) + }) + } + _ => None, + }; + // 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 + } + }); + PushdownStrategy::Map(Arc::clone(&map), pruning_bitmap) } }; @@ -2414,17 +2682,32 @@ async fn collect_left_input( bounds = None; } + if prepared { + drop(batches); + // 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, + input_bytes, + copy_bytes, + retained, + )?; + } + let data = JoinLeftData { - map, - batch, - values: left_values, 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), + membership, + _visited_reservation: visited_reservation, + build: Arc::new(JoinBuildData { + map, + batch, + values: left_values, + bounds, + reservation, + }), }; Ok(data) @@ -2482,7 +2765,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}; @@ -2493,6 +2777,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}; @@ -2504,6 +2791,254 @@ 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(()) + } + + /// 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. + /// + /// 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, @@ -6099,6 +6634,518 @@ 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(()) + } + + /// 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.rs b/datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs new file mode 100644 index 0000000000000..9ec405e919bab --- /dev/null +++ b/datafusion/physical-plan/src/joins/hash_join/exec/prepared.rs @@ -0,0 +1,369 @@ +// 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 { + // Membership may retain build buffers; drop it before the build reservation. + membership: PushdownStrategy, + 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() + self.membership.bitmap_reserved_bytes() + } + + /// 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), + visited_indices_bitmap: Mutex::new(BooleanBufferBuilder::new(0)), + 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(), + } + } + + /// 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, + count, + BuildMode::Prepared, + ) + .await?; + Ok(Arc::new(PreparedHashJoinBuild { + build: data.build, + membership: data.membership, + 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..840ce6f9664fa --- /dev/null +++ b/datafusion/physical-plan/src/joins/hash_join/exec/prepared/tests.rs @@ -0,0 +1,931 @@ +// 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); + 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(()) +} + +#[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.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 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 + .prepare_build( + input(batches, build.schema()), + Arc::clone(&pool), + Arc::new(config), + ) + .await?; + let PushdownStrategy::Map(_, bitmap) = &map.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); + 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.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..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 @@ -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. /// @@ -263,8 +264,12 @@ 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 + pruning_bitmap: Option>, } impl HashTableLookupExpr { /// Create a new HashTableLookupExpr @@ -274,6 +279,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 +288,30 @@ impl HashTableLookupExpr { random_state: SeededRandomState, map: Arc, description: String, + pruning_bitmap: Option>, ) -> Self { Self { 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() + } } impl std::fmt::Debug for HashTableLookupExpr { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { @@ -309,26 +331,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() @@ -353,12 +364,14 @@ 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), + build_lease: self.build_lease.clone(), + description: self.description.clone(), + pruning_bitmap: self.pruning_bitmap.clone(), + })) } fn data_type(&self, _input_schema: &Schema) -> Result { @@ -438,6 +451,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; @@ -448,6 +462,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)); @@ -726,6 +770,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup".to_string(), + None, ); let expr2 = HashTableLookupExpr::new( @@ -733,6 +778,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup".to_string(), + None, ); assert_eq!(expr1, expr2); @@ -751,6 +797,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup".to_string(), + None, ); let expr2 = HashTableLookupExpr::new( @@ -758,6 +805,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup".to_string(), + None, ); assert_ne!(expr1, expr2); @@ -774,6 +822,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup_one".to_string(), + None, ); let expr2 = HashTableLookupExpr::new( @@ -781,6 +830,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup_two".to_string(), + None, ); assert_ne!(expr1, expr2); @@ -800,6 +850,7 @@ mod tests { SeededRandomState::with_seed(1), hash_map1, "lookup".to_string(), + None, ); let expr2 = HashTableLookupExpr::new( @@ -807,6 +858,7 @@ mod tests { SeededRandomState::with_seed(1), hash_map2, "lookup".to_string(), + None, ); // Different Arc pointers means not equal (uses Arc::ptr_eq) @@ -824,6 +876,7 @@ mod tests { SeededRandomState::with_seed(1), Arc::clone(&hash_map), "lookup".to_string(), + None, ); let expr2 = HashTableLookupExpr::new( @@ -831,6 +884,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..fdb877323b2f5 100644 --- a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs +++ b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs @@ -26,11 +26,12 @@ 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, }; +use crate::joins::key_range_bitmap::KeyRangeBitmap; use crate::repartition::RangeExpr; use arrow::array::ArrayRef; use arrow::datatypes::{DataType, Field, Schema}; @@ -95,6 +96,7 @@ impl PartitionBounds { fn create_membership_predicate( on_right: &[PhysicalExprRef], pushdown: PushdownStrategy, + build_lease: Option>, random_state: &SeededRandomState, schema: &Schema, ) -> Result>> { @@ -128,20 +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) => 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, + ) + .with_build_lease(build_lease), + ) + as Arc)), // Empty partition - should not create a filter for this PushdownStrategy::Empty => Ok(None), } @@ -277,22 +284,34 @@ 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, } +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, }, @@ -303,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, @@ -489,6 +509,7 @@ impl SharedBuildAccumulator { PartitionBuildData::Partitioned { partition_id, pushdown, + build_lease, bounds, keys_have_null, }, @@ -502,6 +523,7 @@ impl SharedBuildAccumulator { } partitions[partition_id] = PartitionStatus::Reported(PartitionData { pushdown, + build_lease, bounds, keys_have_null, }); @@ -509,6 +531,7 @@ impl SharedBuildAccumulator { ( PartitionBuildData::CollectLeft { pushdown, + build_lease, bounds, keys_have_null, }, @@ -521,6 +544,7 @@ impl SharedBuildAccumulator { if matches!(data, PartitionStatus::Pending) { *data = PartitionStatus::Reported(PartitionData { pushdown, + build_lease, bounds, keys_have_null, }); @@ -616,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(), )?; @@ -664,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(), )?; @@ -905,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))] } @@ -985,6 +1086,7 @@ mod tests { fn reported(pushdown: PushdownStrategy, bounds: PartitionBounds) -> PartitionStatus { PartitionStatus::Reported(PartitionData { pushdown, + build_lease: None, bounds, keys_have_null: false, }) @@ -1359,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 686939537e73e..3b5185909c7c7 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 @@ -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, }, @@ -1082,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/physical-plan/src/joins/key_range_bitmap.rs b/datafusion/physical-plan/src/joins/key_range_bitmap.rs new file mode 100644 index 0000000000000..9a8767d858132 --- /dev/null +++ b/datafusion/physical-plan/src/joins/key_range_bitmap.rs @@ -0,0 +1,254 @@ +// 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 datafusion_execution::memory_pool::MemoryReservation; +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, + /// Reservation follows this bitmap even when a pruning expression retains it. + reservation: Option, +} + +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, + 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. + 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 e4f7e2e123e0e..f8e9788bc3356 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; @@ -48,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/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(()) +} diff --git a/datafusion/pruning/Cargo.toml b/datafusion/pruning/Cargo.toml index e6f4bb6f273c9..48b1ed009b28c 100644 --- a/datafusion/pruning/Cargo.toml +++ b/datafusion/pruning/Cargo.toml @@ -26,7 +26,17 @@ 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 } +rstest = { 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/benches/string_in_list_pruning.rs b/datafusion/pruning/benches/string_in_list_pruning.rs new file mode 100644 index 0000000000000..54d30922d7300 --- /dev/null +++ b/datafusion/pruning/benches/string_in_list_pruning.rs @@ -0,0 +1,473 @@ +// 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. +//! +//! 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. 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 +//! 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(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), + )) +} + +/// 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, + 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])), + } + } + + 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 { + 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, + 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, +} + +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 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={}", + in_list_predicate + .predicate_expr() + .to_string() + .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 every benchmark path does the same useful work, rather + // than comparing compact pruning with an always-true fallback. + let kept = (0..CONTAINERS) + .map(|index| index % 2 == 0) + .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(), + negated_kept + ); + assert_eq!( + 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, + } + } +} + +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 { + 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, max_in_list_size), + expression, + |bencher, expression| { + bencher.iter(|| { + black_box(build_predicate( + black_box(expression), + &case.schema, + max_in_list_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, 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, list_size), + predicate, + |bencher, predicate| { + bencher.iter(|| { + black_box(predicate.prune(black_box(&case.statistics)).unwrap()) + }); + }, + ); + } + } + 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); +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/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 2b334d2847980..3e7c2ad7aa7ab 100644 --- a/datafusion/pruning/src/lib.rs +++ b/datafusion/pruning/src/lib.rs @@ -18,7 +18,11 @@ #![cfg_attr(test, allow(clippy::needless_pass_by_value))] mod file_pruner; +mod in_list; +mod key_range_bitmap_expr; +mod primitive_in_list; mod pruning_predicate; +mod string_in_list; pub use file_pruner::FilePruner; pub use pruning_predicate::{ 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 3a63451495e4c..4735fa375ced8 100644 --- a/datafusion/pruning/src/pruning_predicate.rs +++ b/datafusion/pruning/src/pruning_predicate.rs @@ -22,6 +22,11 @@ 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}; + use arrow::array::AsArray; use arrow::{ array::{ArrayRef, BooleanArray, new_null_array}, @@ -48,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 @@ -378,6 +384,21 @@ 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, + /// 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 { + /// 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, } /// Build a pruning predicate from an optional predicate expression. @@ -441,9 +462,29 @@ 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`, 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 + /// 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. + /// + /// 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 + /// unsupported lists. /// /// Query engines typically pass /// `datafusion.execution.parquet.max_in_list_size` here. @@ -494,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`) @@ -502,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; } @@ -510,12 +556,45 @@ 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, &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. @@ -528,7 +607,10 @@ 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, }) } } @@ -701,6 +783,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 } @@ -1458,10 +1551,199 @@ 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 +/// 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(), + } + } +} + +/// 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`. +/// +/// `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. +/// +/// 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, +/// 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, + properties: &mut PruningExpressionProperties, +) -> Option> { + 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() { + DataType::Dictionary(_, value) => value.as_ref(), + data_type => data_type, + }; + if field.name() != column.name() { + return None; + } + 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() { + let literal = expr.downcast_ref::()?; + let value = unwrap_scalar(literal.value()); + if value.is_null() { + contains_null = true; + } else { + 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 || domain.is_empty()) { + properties.has_filter_semantics_only = true; + return Some(Arc::new(phys_expr::Literal::new(ScalarValue::Boolean( + Some(false), + )))); + } + 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)) + 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; + } + Some(Arc::new(phys_expr::BinaryExpr::new( + non_null, + Operator::And, + may_match, + ))) +} + +/// Default maximum number of entries in an `IN (...)` list eligible for +/// 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. pub const MAX_IN_LIST_SIZE: usize = 20; @@ -1497,14 +1779,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 @@ -1526,12 +1809,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, ) } } @@ -1544,15 +1829,17 @@ 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 -/// `unhandled_hook`. +/// `max_in_list_size` is the largest `IN (...)` list eligible for statistics +/// 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, 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)` @@ -1587,6 +1874,17 @@ 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. + // 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) = + build_compact_in_list_expr(in_list, schema, required_columns, properties) + { + 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 @@ -1616,11 +1914,63 @@ fn build_predicate_expression( required_columns, unhandled_hook, max_in_list_size, + properties, ); } else { 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); + } + // 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() + .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::() { @@ -1656,6 +2006,7 @@ fn build_predicate_expression( required_columns, unhandled_hook, max_in_list_size, + properties, ); let right_expr = build_predicate_expression( &right, @@ -1663,6 +2014,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) { @@ -1949,6 +2301,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)); @@ -2187,8 +2548,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, ScalarBuffer}, + datatypes::{Int32Type, Int64Type, TimeUnit, i256}, }; use datafusion_expr::expr::InList; use datafusion_expr::{BinaryExpr, Expr, cast, is_null, try_cast}; @@ -2197,7 +2562,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 @@ -3239,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)]); @@ -3463,13 +3911,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. @@ -3478,19 +3921,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(()) } @@ -3518,7 +3951,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 = @@ -3538,8 +3971,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) @@ -3550,8 +3981,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(()) } @@ -3580,93 +4011,1176 @@ mod tests { Ok(()) } - #[test] - fn row_group_predicate_cast_int_int() -> Result<()> { - let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]); - let expected_expr = "c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Int64) <= 1 AND 1 <= CAST(c1_max@1 AS Int64)"; - - // test cast(c1 as int64) = 1 - // test column on the left - let expr = cast(col("c1"), DataType::Int64).eq(lit(ScalarValue::Int64(Some(1)))); - let predicate_expr = - test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); - assert_eq!(predicate_expr.to_string(), expected_expr); - - // test column on the right - let expr = lit(ScalarValue::Int64(Some(1))).eq(cast(col("c1"), DataType::Int64)); - let predicate_expr = - test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); - assert_eq!(predicate_expr.to_string(), expected_expr); - - let expected_expr = - "c1_null_count@1 != row_count@2 AND TRY_CAST(c1_max@0 AS Int64) > 1"; - - // test column on the left - let expr = - try_cast(col("c1"), DataType::Int64).gt(lit(ScalarValue::Int64(Some(1)))); - let predicate_expr = - test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); - assert_eq!(predicate_expr.to_string(), expected_expr); - - // test column on the right - let expr = - lit(ScalarValue::Int64(Some(1))).lt(try_cast(col("c1"), DataType::Int64)); - let predicate_expr = - test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); - assert_eq!(predicate_expr.to_string(), expected_expr); + fn large_in_list_pruning_predicate( + expr: PhysicalExprRef, + schema: SchemaRef, + ) -> Result { + PruningPredicateBuilder::new() + .with_file_schema(schema) + .with_max_in_list_size(10_000) + .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 row_group_predicate_cast_string_string() -> Result<()> { - let schema = Schema::new(vec![Field::new("c1", DataType::Utf8View, false)]); - let expected_expr = "c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Utf8) <= 1 AND 1 <= CAST(c1_max@1 AS Utf8)"; + 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:?}"), + } + } - // test column on the left - let expr = cast(col("c1"), DataType::Utf8) - .eq(lit(ScalarValue::Utf8(Some("1".to_string())))); - let predicate_expr = - test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); - assert_eq!(predicate_expr.to_string(), expected_expr); + 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 column on the right - let expr = lit(ScalarValue::Utf8(Some("1".to_string()))) - .eq(cast(col("c1"), DataType::Utf8)); - let predicate_expr = - test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); - assert_eq!(predicate_expr.to_string(), expected_expr); + #[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 row_group_predicate_cast_string_int() -> Result<()> { - let schema = Schema::new(vec![Field::new("c1", DataType::Utf8View, false)]); - let expected_expr = "true"; - - // test column on the left - let expr = cast(col("c1"), DataType::Int32).eq(lit(ScalarValue::Int32(Some(1)))); - let predicate_expr = - test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); - assert_eq!(predicate_expr.to_string(), expected_expr); - - // test column on the right - let expr = lit(ScalarValue::Int32(Some(1))).eq(cast(col("c1"), DataType::Int32)); - let predicate_expr = - test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); - assert_eq!(predicate_expr.to_string(), expected_expr); - + 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 row_group_predicate_cast_int_string() -> Result<()> { - let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]); - let expected_expr = "true"; + 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 column on the left - let expr = cast(col("c1"), DataType::Utf8) + #[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 = [ + 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_in_list_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_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, _) = + 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_in_list_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::>(); + + 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}" + ); + assert!(predicate.can_be_inverted_for_full_match()); + assert_eq!(predicate.max_in_list_size(), limit); + } + + let default = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(expr)?; + assert!( + is_always_true(default.predicate_expr()), + "negated={negated}" + ); + } + Ok(()) + } + + #[test] + 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( + "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")], + [Some("middle"), Some("other")], + ) + .with_null_counts([Some(0); 2]) + .with_row_counts([Some(1); 2]), + ); + let positive_or = col("c1") + .in_list(with_null.clone(), false) + .or(col("c1").eq(lit("middle"))); + let predicate = large_in_list_pruning_predicate( + logical2physical(&positive_or, &schema), + Arc::clone(&schema), + )?; + assert_eq!(predicate.prune(&or_stats)?, [true, false]); + assert!(!predicate.can_be_inverted_for_full_match()); + + 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()); + + 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}"); + 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]); + } + } + + 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_in_list_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_in_list_pruning_predicate( + logical2physical(&col("c1").in_list(integer_values, false), &integer_schema), + integer_schema, + )?; + 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. + 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 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( + "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_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)", + "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_in_list_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. + /// 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)])); + let values = (0..21).map(|i| lit(format!("a{i:03}"))).collect::>(); + 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_in_list_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_in_list_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_in_list_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)]); + let expected_expr = "c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Int64) <= 1 AND 1 <= CAST(c1_max@1 AS Int64)"; + + // test cast(c1 as int64) = 1 + // test column on the left + let expr = cast(col("c1"), DataType::Int64).eq(lit(ScalarValue::Int64(Some(1)))); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + // test column on the right + let expr = lit(ScalarValue::Int64(Some(1))).eq(cast(col("c1"), DataType::Int64)); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + let expected_expr = + "c1_null_count@1 != row_count@2 AND TRY_CAST(c1_max@0 AS Int64) > 1"; + + // test column on the left + let expr = + try_cast(col("c1"), DataType::Int64).gt(lit(ScalarValue::Int64(Some(1)))); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + // test column on the right + let expr = + lit(ScalarValue::Int64(Some(1))).lt(try_cast(col("c1"), DataType::Int64)); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_cast_string_string() -> Result<()> { + let schema = Schema::new(vec![Field::new("c1", DataType::Utf8View, false)]); + let expected_expr = "c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Utf8) <= 1 AND 1 <= CAST(c1_max@1 AS Utf8)"; + + // test column on the left + let expr = cast(col("c1"), DataType::Utf8) + .eq(lit(ScalarValue::Utf8(Some("1".to_string())))); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + // test column on the right + let expr = lit(ScalarValue::Utf8(Some("1".to_string()))) + .eq(cast(col("c1"), DataType::Utf8)); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_cast_string_int() -> Result<()> { + let schema = Schema::new(vec![Field::new("c1", DataType::Utf8View, false)]); + let expected_expr = "true"; + + // test column on the left + let expr = cast(col("c1"), DataType::Int32).eq(lit(ScalarValue::Int32(Some(1)))); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + // test column on the right + let expr = lit(ScalarValue::Int32(Some(1))).eq(cast(col("c1"), DataType::Int32)); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_cast_int_string() -> Result<()> { + let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]); + let expected_expr = "true"; + + // test column on the left + let expr = cast(col("c1"), DataType::Utf8) .eq(lit(ScalarValue::Utf8(Some("1".to_string())))); let predicate_expr = test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); @@ -6015,12 +7529,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, ) } @@ -6061,4 +7577,379 @@ 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(); + + assert!(predicate.can_be_inverted_for_full_match()); + 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]); + } + #[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(()) + } } diff --git a/datafusion/pruning/src/string_in_list.rs b/datafusion/pruning/src/string_in_list.rs new file mode 100644 index 0000000000000..da1793f5d03a5 --- /dev/null +++ b/datafusion/pruning/src/string_in_list.rs @@ -0,0 +1,357 @@ +// 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; + +use crate::in_list::{SetMembership, not_in_may_match}; + +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, + } +} + +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, + } +} + +macro_rules! string_bytes { + ($value:expr) => { + ($value).as_bytes() + }; +} + +macro_rules! binary_bytes { + ($value:expr) => { + &($value)[..] + }; +} + +/// 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]>, + } + + 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(), + } + } + + /// 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(), + )))); + } + // 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 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}") + } + } + }; +} + +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, BinaryArray, StringArray}; + + 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)), + ] { + 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)); + } + } + + for data_type in [ + 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_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/dynamic_row_group_pruning.slt b/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt index c6700ebf0b97c..963a4037a44aa 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_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; + +statement ok +drop table dim; + statement ok drop table rgsel; diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 573fb04b3451b..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 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 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/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; 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 diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index c750a3e689951..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 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 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" |