From 10aafae8e925e11dc7f8c96d93947314a09342f6 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Wed, 19 Aug 2026 23:23:27 -0700 Subject: [PATCH 1/2] fix: honor Parquet byte-array statistics ordering --- .../benches/parquet_metadata_statistics.rs | 9 +- datafusion/datasource-parquet/src/metadata.rs | 68 +- datafusion/datasource-parquet/src/mod.rs | 2 + .../datasource-parquet/src/opener/mod.rs | 5 +- .../datasource-parquet/src/page_filter.rs | 15 + .../datasource-parquet/src/push_decoder.rs | 5 + .../src/row_group_filter.rs | 121 ++- .../src/statistics_order_tests.rs | 717 ++++++++++++++++++ 8 files changed, 918 insertions(+), 24 deletions(-) create mode 100644 datafusion/datasource-parquet/src/statistics_order_tests.rs diff --git a/datafusion/datasource-parquet/benches/parquet_metadata_statistics.rs b/datafusion/datasource-parquet/benches/parquet_metadata_statistics.rs index 46ebd100fde88..82147abfd9a01 100644 --- a/datafusion/datasource-parquet/benches/parquet_metadata_statistics.rs +++ b/datafusion/datasource-parquet/benches/parquet_metadata_statistics.rs @@ -28,6 +28,7 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use criterion::{BatchSize, BenchmarkId, Criterion, criterion_group, criterion_main}; use datafusion_datasource_parquet::metadata::DFParquetMetadata; use parquet::arrow::ArrowSchemaConverter; +use parquet::basic::ColumnOrder; use parquet::data_type::ByteArray; use parquet::file::metadata::{ ColumnChunkMetaData, FileMetaData, ParquetMetaData, RowGroupMetaData, @@ -169,13 +170,19 @@ fn make_synthetic_metadata( }) .collect::>(); + // Model a modern writer so byte-array bounds have a trustworthy order. + let column_orders = schema_descr + .columns() + .iter() + .map(|column| ColumnOrder::TYPE_DEFINED_ORDER(column.sort_order())) + .collect(); let file_metadata = FileMetaData::new( 1, (spec.row_groups * ROWS_PER_GROUP) as i64, Some("datafusion parquet metadata benchmark".to_string()), None, schema_descr, - None, + Some(column_orders), ); ParquetMetaData::new(file_metadata, row_groups) diff --git a/datafusion/datasource-parquet/src/metadata.rs b/datafusion/datasource-parquet/src/metadata.rs index c01811c087386..8423894eaba84 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,54 @@ 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 byte-array bounds lack a recognized unsigned comparison order. +/// +/// 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. +pub(crate) fn has_untrusted_byte_array_order( + parquet_schema: &SchemaDescriptor, + column_orders: Option<&[ColumnOrder]>, + parquet_column_index: usize, +) -> bool { + let column = parquet_schema.column(parquet_column_index); + 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 +545,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_byte_array_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 b3ce024d66f1f..0b02fcb71d663 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -1105,10 +1105,9 @@ impl FiltersPreparedParquetOpen { // If there is a predicate that can be evaluated against the metadata if let Some(predicate) = self.pruning_predicate.as_ref().map(|p| p.as_ref()) { if prepared.enable_row_group_stats_pruning { - row_groups.prune_by_statistics( + row_groups.prune_by_statistics_with_metadata( &prepared.physical_file_schema, - loaded.reader_metadata.parquet_schema(), - rg_metadata, + &file_metadata, predicate, &prepared.file_metrics, ); diff --git a/datafusion/datasource-parquet/src/page_filter.rs b/datafusion/datasource-parquet/src/page_filter.rs index 6bc1aca667981..75acd431263e1 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_byte_array_order; use arrow::array::BooleanArray; use arrow::{ @@ -490,6 +491,7 @@ struct PagesPruningStatistics<'a> { column_index: &'a ParquetColumnIndex, offset_index: &'a ParquetOffsetIndex, page_offsets: &'a Vec, + trusted_min_max: bool, } impl<'a> PagesPruningStatistics<'a> { @@ -529,6 +531,12 @@ impl<'a> PagesPruningStatistics<'a> { return None; }; let page_offsets = offset_index_metadata.page_locations(); + let file_metadata = parquet_metadata.file_metadata(); + let trusted_min_max = !has_untrusted_byte_array_order( + file_metadata.schema_descr(), + file_metadata.column_orders().map(Vec::as_slice), + parquet_column_index, + ); Some(Self { row_group_index, @@ -537,6 +545,7 @@ impl<'a> PagesPruningStatistics<'a> { column_index, offset_index, page_offsets, + trusted_min_max, }) } @@ -563,6 +572,9 @@ impl<'a> PagesPruningStatistics<'a> { } impl PruningStatistics for PagesPruningStatistics<'_> { fn min_values(&self, _column: &datafusion_common::Column) -> Option { + if !self.trusted_min_max { + return None; + } match self.converter.data_page_mins( self.column_index, self.offset_index, @@ -577,6 +589,9 @@ impl PruningStatistics for PagesPruningStatistics<'_> { } fn max_values(&self, _column: &datafusion_common::Column) -> Option { + if !self.trusted_min_max { + return None; + } match self.converter.data_page_maxes( self.column_index, self.offset_index, diff --git a/datafusion/datasource-parquet/src/push_decoder.rs b/datafusion/datasource-parquet/src/push_decoder.rs index 74d8997198872..d91d86a657902 100644 --- a/datafusion/datasource-parquet/src/push_decoder.rs +++ b/datafusion/datasource-parquet/src/push_decoder.rs @@ -211,6 +211,11 @@ impl RowGroupPruner { .collect::>(); let stats = RowGroupPruningStatistics { parquet_schema: self.parquet_metadata.file_metadata().schema_descr(), + column_orders: self + .parquet_metadata + .file_metadata() + .column_orders() + .map(Vec::as_slice), row_group_metadatas, arrow_schema: self.arrow_schema.as_ref(), // Match the existing static row-group pruning behavior: when a diff --git a/datafusion/datasource-parquet/src/row_group_filter.rs b/datafusion/datasource-parquet/src/row_group_filter.rs index 9231d0232566b..c363a4ed47953 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_order, has_untrusted_byte_array_stats}; use arrow::array::{ArrayRef, BooleanArray, UInt64Array}; +use arrow::compute::nullif; use arrow::datatypes::Schema; use datafusion_common::pruning::PruningStatistics; use datafusion_common::{Column, Result, ScalarValue}; @@ -31,7 +33,8 @@ use datafusion_physical_expr::utils::collect_columns; use datafusion_physical_expr::{PhysicalExpr, PhysicalExprSimplifier}; use datafusion_pruning::{PruningPredicate, PruningPredicateBuilder}; use parquet::arrow::arrow_reader::statistics::StatisticsConverter; -use parquet::file::metadata::RowGroupMetaData; +use parquet::basic::ColumnOrder; +use parquet::file::metadata::{ParquetMetaData, RowGroupMetaData}; use parquet::schema::types::SchemaDescriptor; /// Reduces the [`ParquetAccessPlan`] based on row group level metadata. @@ -252,8 +255,10 @@ impl RowGroupAccessPlanFilter { /// /// Updates this set to mark row groups that should not be scanned /// - /// Note: This method currently ignores ColumnOrder - /// + /// This method has no file footer, so it cannot establish the sort order + /// of string or binary min/max statistics. Such bounds are ignored. Use + /// [`Self::prune_by_statistics_with_metadata`] when the full metadata is + /// available. Null counts and other columns' statistics remain usable. /// /// # Panics /// if `groups.len() != self.len()` @@ -264,6 +269,51 @@ impl RowGroupAccessPlanFilter { groups: &[RowGroupMetaData], predicate: &PruningPredicate, metrics: &ParquetFileMetrics, + ) { + self.prune_by_statistics_inner( + arrow_schema, + parquet_schema, + groups, + None, + predicate, + metrics, + ); + } + + /// Prune row groups using statistics and the comparison orders recorded + /// in the Parquet file footer. + /// + /// Unlike [`Self::prune_by_statistics`], this method can use string and + /// binary bounds when their ordering is known to match Arrow's ordering. + /// + /// # Panics + /// if `metadata.num_row_groups() != self.len()` + pub fn prune_by_statistics_with_metadata( + &mut self, + arrow_schema: &Schema, + metadata: &ParquetMetaData, + predicate: &PruningPredicate, + metrics: &ParquetFileMetrics, + ) { + let file_metadata = metadata.file_metadata(); + self.prune_by_statistics_inner( + arrow_schema, + file_metadata.schema_descr(), + metadata.row_groups(), + file_metadata.column_orders().map(Vec::as_slice), + predicate, + metrics, + ); + } + + fn prune_by_statistics_inner( + &mut self, + arrow_schema: &Schema, + parquet_schema: &SchemaDescriptor, + groups: &[RowGroupMetaData], + column_orders: Option<&[ColumnOrder]>, + predicate: &PruningPredicate, + metrics: &ParquetFileMetrics, ) { // scoped timer updates on drop let _timer_guard = metrics.statistics_eval_time.timer(); @@ -278,6 +328,7 @@ impl RowGroupAccessPlanFilter { let pruning_stats = RowGroupPruningStatistics { parquet_schema, + column_orders, row_group_metadatas, arrow_schema, // Preserve the existing row-group pruning behavior. This path only @@ -304,8 +355,7 @@ impl RowGroupAccessPlanFilter { // Check if any of the matched row groups are fully contained by the predicate self.identify_fully_matched_row_groups( &fully_contained_candidates_original_idx, - arrow_schema, - parquet_schema, + &pruning_stats, groups, predicate, metrics, @@ -330,8 +380,7 @@ impl RowGroupAccessPlanFilter { fn identify_fully_matched_row_groups( &mut self, candidate_row_group_indices: &[usize], - arrow_schema: &Schema, - parquet_schema: &SchemaDescriptor, + pruning_stats: &RowGroupPruningStatistics<'_>, groups: &[RowGroupMetaData], predicate: &PruningPredicate, metrics: &ParquetFileMetrics, @@ -339,6 +388,7 @@ impl RowGroupAccessPlanFilter { if candidate_row_group_indices.is_empty() { return; } + let arrow_schema = pruning_stats.arrow_schema; let mut inverted_expr: Arc = Arc::new(NotExpr::new(Arc::clone(predicate.orig_expr()))); @@ -381,7 +431,8 @@ impl RowGroupAccessPlanFilter { }; let inverted_pruning_stats = RowGroupPruningStatistics { - parquet_schema, + parquet_schema: pruning_stats.parquet_schema, + column_orders: pruning_stats.column_orders, row_group_metadatas: candidate_row_group_indices .iter() .map(|&i| &groups[i]) @@ -472,6 +523,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, @@ -479,14 +531,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, @@ -494,19 +543,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_byte_array_order(self.parquet_schema, self.column_orders, index) + }) { + return None; + } + Some(converter) + } + + fn mask_untrusted_byte_array_stats( + &self, + parquet_index: Option, + values: ArrayRef, + ) -> Option { + if !has_untrusted_byte_array_stats( + self.parquet_schema, + parquet_index, + self.metadata_iter(), + ) { + return Some(values); + } + // A file may mix legacy and modern row-group statistics. Keep the + // modern bounds usable rather than discarding this entire column. + let mask = BooleanArray::from_iter(self.metadata_iter().map(|group| { + has_untrusted_byte_array_stats(self.parquet_schema, parquet_index, [group]) + })); + nullif(values.as_ref(), &mask).ok() + } } impl PruningStatistics for RowGroupPruningStatistics<'_> { fn min_values(&self, column: &Column) -> Option { - self.statistics_converter(column) - .and_then(|c| Ok(c.row_group_mins(self.metadata_iter())?)) - .ok() + let converter = self.min_max_statistics_converter(column)?; + let values = converter.row_group_mins(self.metadata_iter()).ok()?; + self.mask_untrusted_byte_array_stats(converter.parquet_column_index(), values) } fn max_values(&self, column: &Column) -> Option { - self.statistics_converter(column) - .and_then(|c| Ok(c.row_group_maxes(self.metadata_iter())?)) - .ok() + let converter = self.min_max_statistics_converter(column)?; + let values = converter.row_group_maxes(self.metadata_iter()).ok()?; + self.mask_untrusted_byte_array_stats(converter.parquet_column_index(), values) } fn num_containers(&self) -> usize { diff --git a/datafusion/datasource-parquet/src/statistics_order_tests.rs b/datafusion/datasource-parquet/src/statistics_order_tests.rs new file mode 100644 index 0000000000000..f678fa6b47ee2 --- /dev/null +++ b/datafusion/datasource-parquet/src/statistics_order_tests.rs @@ -0,0 +1,717 @@ +// 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; +use crate::push_decoder::RowGroupPruner; +use crate::row_group_filter::RowGroupPruningStatistics; +use crate::{PagePruningAccessPlanFilter, ParquetAccessPlan, ParquetFileMetrics}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum StatisticsOrder { + Modern, + Deprecated, + Missing, + Unknown, +} + +struct TestFile { + bytes: Bytes, + schema: SchemaRef, + metadata: Arc, +} + +impl TestFile { + fn new(order: StatisticsOrder) -> Self { + let batch = record_batch!( + ( + "s", + Utf8, + vec![ + Some("aé"), + Some("az"), + Some("b"), + None, + None, + None, + Some("d"), + Some("e"), + Some("f"), + ] + ), + ("n", Int32, [1, 2, 3, 10, 11, 12, 20, 21, 22]) + ) + .unwrap(); + let schema = batch.schema(); + let properties = WriterProperties::builder() + .set_max_row_group_row_count(Some(3)) + .set_data_page_row_count_limit(3) + .set_write_batch_size(3) + .set_dictionary_enabled(false) + .set_statistics_enabled(EnabledStatistics::Page) + .build(); + let mut original = Vec::new(); + let mut writer = + ArrowWriter::try_new(&mut original, Arc::clone(&schema), Some(properties)) + .unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + let original = Bytes::from(original); + let metadata = read_metadata(&original); + assert_eq!(metadata.num_row_groups(), 3); + + if order == StatisticsOrder::Modern { + return Self { + bytes: original, + schema, + metadata: Arc::new(metadata), + }; + } + + // Signed-byte comparison gives ["aé", "b"] for the first row + // group's ["aé", "az", "b"]. The endpoints are not inverted in + // unsigned order, but the interval wrongly excludes "az". + let mut row_groups = metadata.row_groups().to_vec(); + let mut columns = row_groups[0].columns().to_vec(); + columns[0] = columns[0] + .clone() + .into_builder() + .set_statistics(ParquetStatistics::byte_array( + Some(ByteArray::from("aé")), + Some(ByteArray::from("b")), + None, + Some(0), + order == StatisticsOrder::Deprecated, + )) + .build() + .unwrap(); + row_groups[0] = row_groups[0] + .clone() + .into_builder() + .set_column_metadata(columns) + .build() + .unwrap(); + + let mut column_index = metadata.column_index().unwrap().clone(); + if matches!(order, StatisticsOrder::Missing | StatisticsOrder::Unknown) { + let mut index = ColumnIndexBuilder::new(PhysicalType::BYTE_ARRAY); + index.append(false, "aé".as_bytes().to_vec(), b"b".to_vec(), 0); + column_index[0][0] = index.build().unwrap(); + } + let metadata = metadata + .into_builder() + .set_row_groups(row_groups) + .set_column_index(Some(column_index)) + .build(); + + // Keep the real data pages, and serialize the replacement statistics + // and page indexes at their actual file offsets. + let mut bytes = Vec::new(); + let mut tracked = TrackedWrite::new(&mut bytes); + tracked + .write_all(&original[..footer_start(&original)]) + .unwrap(); + ParquetMetaDataWriter::new_with_tracked(tracked, &metadata) + .finish() + .unwrap(); + + // parquet-rs always writes TYPEORDER and cannot write an unknown + // ColumnOrder. The final Thrift field is column_orders (field 7): a + // two-element list of unions, followed by the FileMetaData STOP. + // Alter just that field to model old and future writers, then verify + // the decoded footer below. No data or page-index offsets change. + if matches!(order, StatisticsOrder::Missing | StatisticsOrder::Unknown) { + let end = bytes.len() - 8; + let encoded_orders = [0x19, 0x2c, 0x1c, 0, 0, 0x1c, 0, 0, 0]; + let start = end - encoded_orders.len(); + assert_eq!(&bytes[start..end], &encoded_orders); + let metadata_start = footer_start(&bytes); + if order == StatisticsOrder::Missing { + bytes.drain(start..end - 1); + let new_end = bytes.len() - 8; + let metadata_len = (new_end - metadata_start) as u32; + bytes[new_end..new_end + 4].copy_from_slice(&metadata_len.to_le_bytes()); + } else { + // Change the first union member from field 1 (TYPEORDER) to + // an unrecognized field 2. The numeric column stays known. + bytes[start + 2] = 0x2c; + } + } + + let bytes = Bytes::from(bytes); + let metadata = read_metadata(&bytes); + let expected_order = match order { + StatisticsOrder::Missing => ColumnOrder::UNDEFINED, + StatisticsOrder::Unknown => ColumnOrder::UNKNOWN, + _ => ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNSIGNED), + }; + assert_eq!(metadata.file_metadata().column_order(0), expected_order); + assert_eq!( + metadata + .row_group(0) + .column(0) + .statistics() + .unwrap() + .is_min_max_deprecated(), + order == StatisticsOrder::Deprecated, + ); + Self { + bytes, + schema, + metadata: Arc::new(metadata), + } + } + + fn predicate(&self, expr: &Expr) -> (Arc, PruningPredicate) { + let physical = logical2physical(expr, &self.schema); + let pruning = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&self.schema)) + .try_build(Arc::clone(&physical)) + .unwrap(); + (physical, pruning) + } + + fn statistics(&self) -> Statistics { + DFParquetMetadata::statistics_from_parquet_metadata(&self.metadata, &self.schema) + .unwrap() + } + + fn file_matches(&self, predicate: &PruningPredicate) -> bool { + let stats = PrunableStatistics::new( + vec![Arc::new(self.statistics())], + Arc::clone(&self.schema), + ); + predicate.prune(&stats).unwrap()[0] + } + + fn row_group_plan(&self, predicate: &PruningPredicate) -> ParquetAccessPlan { + let mut filter = RowGroupAccessPlanFilter::new(ParquetAccessPlan::new_all( + self.metadata.num_row_groups(), + )); + filter.prune_by_statistics_with_metadata( + &self.schema, + &self.metadata, + predicate, + &metrics(), + ); + filter.build() + } + + fn page_plan( + &self, + physical: &Arc, + plan: ParquetAccessPlan, + ) -> ParquetAccessPlan { + PagePruningAccessPlanFilter::new(physical, Arc::clone(&self.schema)) + .prune_plan_with_page_index( + plan, + &self.schema, + self.metadata.file_metadata().schema_descr(), + &self.metadata, + &metrics(), + ) + } + + fn matching_rows( + &self, + physical: &Arc, + plan: ParquetAccessPlan, + ) -> usize { + let mut builder = ParquetRecordBatchReaderBuilder::try_new(self.bytes.clone()) + .unwrap() + .with_row_groups(plan.row_group_indexes()); + if let Some(selection) = plan + .into_overall_row_selection(self.metadata.row_groups()) + .unwrap() + { + builder = builder.with_row_selection(selection); + } + builder + .build() + .unwrap() + .map(|batch| { + let batch = batch.unwrap(); + let matches = physical + .evaluate(&batch) + .unwrap() + .into_array(batch.num_rows()) + .unwrap(); + matches + .as_any() + .downcast_ref::() + .unwrap() + .true_count() + }) + .sum() + } +} + +fn footer_start(bytes: &[u8]) -> usize { + let end = bytes.len() - 8; + let metadata_len = u32::from_le_bytes(bytes[end..end + 4].try_into().unwrap()); + end - metadata_len as usize +} + +fn read_metadata(bytes: &Bytes) -> ParquetMetaData { + ParquetMetaDataReader::new() + .with_page_index_policy(PageIndexPolicy::Required) + .parse_and_finish(bytes) + .unwrap() +} + +fn metrics() -> ParquetFileMetrics { + ParquetFileMetrics::new( + 0, + "statistics-order.parquet", + &ExecutionPlanMetricsSet::new(), + ) +} + +#[test] +fn byte_array_order_preserves_matching_rows_at_every_pruning_level() { + for order in [ + StatisticsOrder::Modern, + StatisticsOrder::Deprecated, + StatisticsOrder::Missing, + StatisticsOrder::Unknown, + ] { + let file = TestFile::new(order); + let (physical, predicate) = file.predicate(&col("s").eq(lit("az"))); + let all = ParquetAccessPlan::new_all(file.metadata.num_row_groups()); + assert_eq!(file.matching_rows(&physical, all.clone()), 1); + assert!(file.file_matches(&predicate), "order={order:?}"); + let row_groups = file.row_group_plan(&predicate); + assert!(row_groups.should_scan(0), "order={order:?}"); + assert!(!row_groups.is_fully_matched(0), "order={order:?}"); + assert_eq!( + row_groups.row_group_indexes(), + if matches!(order, StatisticsOrder::Modern | StatisticsOrder::Deprecated) { + vec![0] + } else { + vec![0, 2] + }, + ); + assert_eq!(file.matching_rows(&physical, row_groups.clone()), 1); + assert_eq!( + file.matching_rows(&physical, file.page_plan(&physical, all)), + 1 + ); + assert_eq!( + file.matching_rows(&physical, file.page_plan(&physical, row_groups)), + 1, + "order={order:?}", + ); + + let mut runtime_pruner = RowGroupPruner::new( + physical, + Arc::clone(&file.schema), + Arc::clone(&file.metadata), + Count::new(), + Count::new(), + MAX_IN_LIST_SIZE, + ); + assert!(!runtime_pruner.should_prune(&[0]), "order={order:?}"); + assert!(runtime_pruner.should_prune(&[1]), "all-null row group"); + } +} + +#[test] +fn byte_array_order_keeps_null_counts_and_unrelated_column_bounds() { + for order in [ + StatisticsOrder::Deprecated, + StatisticsOrder::Missing, + StatisticsOrder::Unknown, + ] { + let file = TestFile::new(order); + let statistics = file.statistics(); + let string = &statistics.column_statistics[0]; + assert_eq!(string.min_value, Precision::Absent, "order={order:?}"); + assert_eq!(string.max_value, Precision::Absent, "order={order:?}"); + assert_eq!(string.null_count, Precision::Exact(3)); + assert_eq!( + statistics.column_statistics[1].min_value, + Precision::Exact(ScalarValue::Int32(Some(1))), + ); + assert_eq!( + statistics.column_statistics[1].max_value, + Precision::Exact(ScalarValue::Int32(Some(22))), + ); + + let (physical, predicate) = file.predicate(&col("s").is_null()); + assert_eq!(file.row_group_plan(&predicate).row_group_indexes(), vec![1]); + let page_plan = file.page_plan(&physical, ParquetAccessPlan::new_all(3)); + assert_eq!(page_plan.row_group_indexes(), vec![1]); + assert_eq!(file.matching_rows(&physical, page_plan), 3); + + let expr = col("s").eq(lit("az")).and(col("n").eq(lit(99))); + let (physical, predicate) = file.predicate(&expr); + assert!(!file.file_matches(&predicate)); + assert!( + file.row_group_plan(&predicate) + .row_group_indexes() + .is_empty() + ); + assert!( + file.page_plan(&physical, ParquetAccessPlan::new_all(3)) + .row_group_indexes() + .is_empty(), + ); + } +} + +#[test] +fn byte_array_order_modern_bounds_and_null_only_groups_remain_usable() { + let file = TestFile::new(StatisticsOrder::Modern); + let statistics = file.statistics(); + assert_eq!( + statistics.column_statistics[0].min_value, + Precision::Exact(ScalarValue::Utf8(Some("az".to_owned()))), + ); + let (physical, predicate) = file.predicate(&col("s").eq(lit("az"))); + assert_eq!(file.row_group_plan(&predicate).row_group_indexes(), vec![0]); + assert_eq!( + file.page_plan(&physical, ParquetAccessPlan::new_all(3)) + .row_group_indexes(), + vec![0], + ); + + // The compatibility API has no footer. It still uses null counts, but + // cannot assume that even modern-looking byte-array bounds are unsigned. + let mut filter = RowGroupAccessPlanFilter::new(ParquetAccessPlan::new_all(3)); + filter.prune_by_statistics( + &file.schema, + file.metadata.file_metadata().schema_descr(), + file.metadata.row_groups(), + &predicate, + &metrics(), + ); + assert_eq!(filter.build().row_group_indexes(), vec![0, 2]); +} + +#[test] +fn byte_array_order_guard_follows_parquet_type_not_arrow_representation() { + let file = TestFile::new(StatisticsOrder::Deprecated); + let metadata = file.metadata.file_metadata(); + let column = Column::from_name("s"); + for data_type in [ + DataType::Utf8, + DataType::LargeUtf8, + DataType::Utf8View, + DataType::Binary, + DataType::LargeBinary, + DataType::BinaryView, + DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)), + ] { + let schema = Schema::new(vec![Field::new("s", data_type, true)]); + let stats = RowGroupPruningStatistics { + parquet_schema: metadata.schema_descr(), + column_orders: metadata.column_orders().map(Vec::as_slice), + row_group_metadatas: file.metadata.row_groups().iter().collect(), + arrow_schema: &schema, + missing_null_counts_as_zero: true, + }; + for values in [stats.min_values(&column), stats.max_values(&column)] { + let values = values.unwrap(); + assert!(values.is_null(0)); + assert!(values.is_null(1)); + assert!(!values.is_null(2)); + } + assert!(stats.null_counts(&column).is_some()); + } +} + +fn single_column_metadata( + parquet_type: ParquetType, + statistics: ParquetStatistics, + order: Option, +) -> ParquetMetaData { + let physical_type = parquet_type.get_physical_type(); + let schema = Arc::new(SchemaDescriptor::new(Arc::new( + ParquetType::group_type_builder("schema") + .with_fields(vec![Arc::new(parquet_type)]) + .build() + .unwrap(), + ))); + let mut column_index = ColumnIndexBuilder::new(physical_type); + column_index.append( + false, + statistics.min_bytes_opt().unwrap().to_vec(), + statistics.max_bytes_opt().unwrap().to_vec(), + 0, + ); + let mut offset_index = OffsetIndexBuilder::new(); + offset_index.append_row_count(3); + offset_index.append_offset_and_size(0, 1); + let column = ColumnChunkMetaData::builder(schema.column(0)) + .set_num_values(3) + .set_statistics(statistics) + .build() + .unwrap(); + let group = RowGroupMetaData::builder(Arc::clone(&schema)) + .set_num_rows(3) + .set_column_metadata(vec![column]) + .build() + .unwrap(); + ParquetMetaData::new( + FileMetaData::new(1, 3, None, None, schema, order.map(|order| vec![order])), + vec![group], + ) + .into_builder() + .set_column_index(Some(vec![vec![column_index.build().unwrap()]])) + .set_offset_index(Some(vec![vec![offset_index.build()]])) + .build() +} + +#[test] +fn fixed_byte_array_and_uuid_orders_guard_bounds_but_not_null_counts() { + let schema = Arc::new(Schema::new(vec![Field::new( + "s", + DataType::FixedSizeBinary(16), + true, + )])); + let padded = |prefix: &[u8]| { + let mut value = vec![0; 16]; + value[..prefix.len()].copy_from_slice(prefix); + value + }; + let min = padded("aé".as_bytes()); + let max = padded(b"b"); + let value = ScalarValue::FixedSizeBinary(16, Some(padded(b"az"))); + let physical = logical2physical(&col("s").eq(lit(value)), &schema); + let predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(Arc::clone(&physical)) + .unwrap(); + + for logical_type in [None, Some(LogicalType::Uuid)] { + for deprecated in [false, true] { + for order in [ + None, + Some(ColumnOrder::UNKNOWN), + Some(ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::SIGNED)), + Some(ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNSIGNED)), + ] { + let parquet_type = ParquetType::primitive_type_builder( + "s", + PhysicalType::FIXED_LEN_BYTE_ARRAY, + ) + .with_length(16) + .with_logical_type(logical_type.clone()) + .build() + .unwrap(); + let statistics = ParquetStatistics::fixed_len_byte_array( + Some(FixedLenByteArray::from(min.clone())), + Some(FixedLenByteArray::from(max.clone())), + None, + Some(0), + deprecated, + ); + let metadata = single_column_metadata(parquet_type, statistics, order); + let trusted = + order == Some(ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNSIGNED)); + let statistics = DFParquetMetadata::statistics_from_parquet_metadata( + &metadata, &schema, + ) + .unwrap(); + assert_eq!( + statistics.column_statistics[0].min_value, + if trusted && !deprecated { + Precision::Exact(ScalarValue::FixedSizeBinary( + 16, + Some(min.clone()), + )) + } else { + Precision::Absent + }, + ); + assert_eq!( + statistics.column_statistics[0].null_count, + Precision::Exact(0), + ); + + let mut row_groups = + RowGroupAccessPlanFilter::new(ParquetAccessPlan::new_all(1)); + row_groups.prune_by_statistics_with_metadata( + &schema, + &metadata, + &predicate, + &metrics(), + ); + assert_eq!(row_groups.build().should_scan(0), !trusted || deprecated); + let pages = + PagePruningAccessPlanFilter::new(&physical, Arc::clone(&schema)) + .prune_plan_with_page_index( + ParquetAccessPlan::new_all(1), + &schema, + metadata.file_metadata().schema_descr(), + &metadata, + &metrics(), + ); + // Modern page indexes are independent of legacy row-group + // bounds, but still need a recognized footer order. + assert_eq!(pages.should_scan(0), !trusted); + } + } + } +} + +#[test] +fn signed_decimal_byte_array_statistics_remain_usable() { + let schema = Arc::new(Schema::new(vec![Field::new( + "s", + DataType::Decimal128(10, 0), + true, + )])); + for physical_type in [PhysicalType::BYTE_ARRAY, PhysicalType::FIXED_LEN_BYTE_ARRAY] { + let parquet_type = ParquetType::primitive_type_builder("s", physical_type) + .with_length(16) + .with_logical_type(Some(LogicalType::decimal(0, 10))) + .with_precision(10) + .with_scale(0) + .build() + .unwrap(); + let min = (-2_i128).to_be_bytes().to_vec(); + let max = 3_i128.to_be_bytes().to_vec(); + let statistics = match physical_type { + PhysicalType::BYTE_ARRAY => ParquetStatistics::byte_array( + Some(ByteArray::from(min)), + Some(ByteArray::from(max)), + None, + Some(0), + true, + ), + _ => ParquetStatistics::fixed_len_byte_array( + Some(FixedLenByteArray::from(min)), + Some(FixedLenByteArray::from(max)), + None, + Some(0), + true, + ), + }; + let metadata = single_column_metadata(parquet_type, statistics, None); + let statistics = + DFParquetMetadata::statistics_from_parquet_metadata(&metadata, &schema) + .unwrap(); + assert_eq!( + statistics.column_statistics[0].min_value, + Precision::Exact(ScalarValue::Decimal128(Some(-2), 10, 0)), + ); + let physical = logical2physical( + &col("s").eq(lit(ScalarValue::Decimal128(Some(4), 10, 0))), + &schema, + ); + let predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(physical) + .unwrap(); + let mut row_groups = RowGroupAccessPlanFilter::new(ParquetAccessPlan::new_all(1)); + row_groups.prune_by_statistics_with_metadata( + &schema, + &metadata, + &predicate, + &metrics(), + ); + assert!(!row_groups.build().should_scan(0)); + } +} + +#[test] +fn undefined_logical_byte_array_order_is_not_a_bound() { + let parquet_type = ParquetType::primitive_type_builder("s", PhysicalType::BYTE_ARRAY) + .with_logical_type(Some(LogicalType::_Unknown { field_id: 100 })) + .build() + .unwrap(); + let statistics = ParquetStatistics::byte_array( + Some(ByteArray::from("aé")), + Some(ByteArray::from("b")), + None, + Some(0), + false, + ); + let metadata = single_column_metadata( + parquet_type, + statistics, + Some(ColumnOrder::TYPE_DEFINED_ORDER(SortOrder::UNDEFINED)), + ); + let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, true)])); + let file_statistics = + DFParquetMetadata::statistics_from_parquet_metadata(&metadata, &schema).unwrap(); + assert_eq!( + file_statistics.column_statistics[0].min_value, + Precision::Absent + ); + assert_eq!( + file_statistics.column_statistics[0].max_value, + Precision::Absent + ); + assert_eq!( + file_statistics.column_statistics[0].null_count, + Precision::Exact(0) + ); + let physical = logical2physical(&col("s").eq(lit("az")), &schema); + let predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(Arc::clone(&physical)) + .unwrap(); + let mut row_groups = RowGroupAccessPlanFilter::new(ParquetAccessPlan::new_all(1)); + row_groups.prune_by_statistics_with_metadata( + &schema, + &metadata, + &predicate, + &metrics(), + ); + assert!(row_groups.build().should_scan(0)); + let pages = PagePruningAccessPlanFilter::new(&physical, Arc::clone(&schema)) + .prune_plan_with_page_index( + ParquetAccessPlan::new_all(1), + &schema, + metadata.file_metadata().schema_descr(), + &metadata, + &metrics(), + ); + assert!(pages.should_scan(0)); +} From b20b0e72659edbc67284e26442605647de122309 Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Wed, 19 Aug 2026 23:38:51 -0700 Subject: [PATCH 2/2] perf: use compact pruning for large string IN lists --- Cargo.lock | 1 + datafusion/core/tests/parquet/mod.rs | 1 + .../tests/parquet/string_in_list_pruning.rs | 363 ++++++++++++++++++ .../datasource-parquet/src/opener/mod.rs | 17 +- .../datasource-parquet/src/page_filter.rs | 16 +- datafusion/datasource-parquet/src/source.rs | 5 +- .../src/statistics_order_tests.rs | 68 ++++ datafusion/pruning/Cargo.toml | 5 + .../pruning/benches/string_in_list_pruning.rs | 241 ++++++++++++ datafusion/pruning/src/lib.rs | 1 + datafusion/pruning/src/pruning_predicate.rs | 316 ++++++++++++++- datafusion/pruning/src/string_in_list.rs | 134 +++++++ 12 files changed, 1144 insertions(+), 24 deletions(-) create mode 100644 datafusion/core/tests/parquet/string_in_list_pruning.rs create mode 100644 datafusion/pruning/benches/string_in_list_pruning.rs create mode 100644 datafusion/pruning/src/string_in_list.rs diff --git a/Cargo.lock b/Cargo.lock index 62b94e7e7f947..6b0fd3f60c837 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2588,6 +2588,7 @@ name = "datafusion-pruning" version = "55.0.0" dependencies = [ "arrow", + "criterion", "datafusion-common", "datafusion-datasource", "datafusion-expr", diff --git a/datafusion/core/tests/parquet/mod.rs b/datafusion/core/tests/parquet/mod.rs index 7066a4147c017..cf2e557abad5f 100644 --- a/datafusion/core/tests/parquet/mod.rs +++ b/datafusion/core/tests/parquet/mod.rs @@ -58,6 +58,7 @@ mod page_pruning; mod row_group_pruning; mod schema; mod schema_coercion; +mod string_in_list_pruning; mod utils; #[cfg(test)] diff --git a/datafusion/core/tests/parquet/string_in_list_pruning.rs b/datafusion/core/tests/parquet/string_in_list_pruning.rs new file mode 100644 index 0000000000000..e1005daaa1a38 --- /dev/null +++ b/datafusion/core/tests/parquet/string_in_list_pruning.rs @@ -0,0 +1,363 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! End-to-end coverage for compact, large string IN-list pruning. The positive +//! IN-list cases disable the row and Bloom filters to isolate min/max pruning. + +use std::sync::Arc; + +use arrow::array::StringArray; +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::record_batch::RecordBatch; +use arrow::util::pretty::pretty_format_batches; +use datafusion::datasource::listing::PartitionedFile; +use datafusion::datasource::object_store::ObjectStoreUrl; +use datafusion::datasource::physical_plan::ParquetSource; +use datafusion::datasource::source::DataSourceExec; +use datafusion::physical_plan::{ExecutionPlan, collect, displayable}; +use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; +use datafusion_common::config::TableParquetOptions; +use datafusion_common::{ScalarValue, assert_batches_eq}; +use datafusion_datasource::file_scan_config::FileScanConfigBuilder; +use datafusion_physical_expr::expressions::{col, in_list, lit}; +use datafusion_physical_plan::metrics::{MetricValue, MetricsSet}; +use object_store::path::Path; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::{EnabledStatistics, WriterProperties}; +use tempfile::NamedTempFile; + +use super::utils::MetricsFinder; + +const ROWS_PER_UNIT: usize = 16; +const UNITS: usize = 4; +const TOTAL_ROWS: usize = ROWS_PER_UNIT * UNITS; +const MATCHING_ROWS: usize = ROWS_PER_UNIT * 2; + +/// Write either four row groups or four pages in one row group. The second +/// unit lies in a gap between two members of every test IN list; an enclosing +/// min/max range for the list cannot prune it. +fn make_file(page_pruning: bool) -> NamedTempFile { + let mut file = tempfile::Builder::new() + .prefix("string_in_list_pruning") + .suffix(".parquet") + .tempfile() + .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Utf8, + false, + )])); + let values = ["v000000", "v000001", "v000010", "v999999"] + .into_iter() + .flat_map(|value| std::iter::repeat_n(value, ROWS_PER_UNIT)) + .collect::>(); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(StringArray::from(values))], + ) + .unwrap(); + let rows_per_group = if page_pruning { + TOTAL_ROWS + } else { + ROWS_PER_UNIT + }; + let properties = WriterProperties::builder() + .set_max_row_group_row_count(Some(rows_per_group)) + .set_data_page_row_count_limit(ROWS_PER_UNIT) + .set_write_batch_size(ROWS_PER_UNIT) + .set_dictionary_enabled(false) + .set_bloom_filter_enabled(false) + .set_statistics_enabled(EnabledStatistics::Page) + .build(); + let mut writer = ArrowWriter::try_new(&mut file, schema, Some(properties)).unwrap(); + writer.write(&batch).unwrap(); + let metadata = writer.close().unwrap(); + assert_eq!(metadata.num_row_groups(), TOTAL_ROWS / rows_per_group); + let offsets = metadata.offset_index().unwrap(); + for row_group in offsets { + assert_eq!( + row_group[0].page_locations().len(), + rows_per_group / ROWS_PER_UNIT + ); + } + file +} + +struct ScanOutput { + batches: Vec, + plan: String, + metrics: MetricsSet, +} + +impl ScanOutput { + fn counter(&self, name: &str) -> usize { + self.metrics + .sum(|metric| metric.value().name() == name) + .unwrap_or_else(|| panic!("missing {name}: {}", self.metrics)) + .as_usize() + } + + fn pruned(&self, name: &str) -> usize { + let value = self + .metrics + .sum(|metric| metric.value().name() == name) + .unwrap_or_else(|| panic!("missing {name}: {}", self.metrics)); + let MetricValue::PruningMetrics { + pruning_metrics, .. + } = value + else { + panic!("expected pruning metric {name}: {}", self.metrics); + }; + pruning_metrics.pruned() + } + + fn fully_matched(&self, name: &str) -> usize { + let value = self + .metrics + .sum(|metric| metric.value().name() == name) + .unwrap_or_else(|| panic!("missing {name}: {}", self.metrics)); + let MetricValue::PruningMetrics { + pruning_metrics, .. + } = value + else { + panic!("expected pruning metric {name}: {}", self.metrics); + }; + pruning_metrics.fully_matched() + } + + fn assert_results(&self) { + assert_batches_eq!( + [ + "+---------+----+", + "| value | n |", + "+---------+----+", + "| v000000 | 16 |", + "| v000010 | 16 |", + "+---------+----+", + ], + &self.batches + ); + assert_eq!(self.counter("predicate_evaluation_errors"), 0); + assert_eq!(self.counter("pushdown_rows_pruned"), 0); + assert_eq!(self.pruned("row_groups_pruned_bloom_filter"), 0); + } +} + +async fn scan( + file: &NamedTempFile, + list_size: usize, + max_in_list_size: Option, + page_pruning: bool, +) -> ScanOutput { + let mut config = SessionConfig::new() + .with_target_partitions(1) + .with_parquet_bloom_filter_pruning(false) + .with_parquet_page_index_pruning(page_pruning); + config.options_mut().execution.parquet.pushdown_filters = false; + if let Some(max_in_list_size) = max_in_list_size { + config.options_mut().execution.parquet.max_in_list_size = max_in_list_size; + } + let ctx = SessionContext::new_with_config(config); + ctx.register_parquet( + "t", + file.path().to_str().unwrap(), + ParquetReadOptions::default(), + ) + .await + .unwrap(); + let values = (0..list_size) + .map(|index| format!("'v{:06}'", index * 10)) + .collect::>() + .join(", "); + let sql = format!( + "SELECT value, count(*) AS n FROM t \ + WHERE value IN ({values}) GROUP BY value ORDER BY value" + ); + let plan = ctx + .sql(&sql) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let plan_text = displayable(plan.as_ref()).indent(true).to_string(); + let batches = collect(Arc::clone(&plan), ctx.task_ctx()).await.unwrap(); + let metrics = MetricsFinder::find_metrics(plan.as_ref()).unwrap(); + ScanOutput { + batches, + plan: plan_text, + metrics, + } +} + +async fn check_string_in_list_pruning(page_pruning: bool) { + let file = make_file(page_pruning); + for list_size in [20, 21, 256, 1024] { + // A zero cap provides a result-equivalence control that cannot use + // min/max IN-list pruning at either granularity. + let unpruned = scan(&file, list_size, Some(0), page_pruning).await; + unpruned.assert_results(); + assert!(!unpruned.plan.contains("IN_SET_INTERSECTS")); + assert_eq!(unpruned.pruned("row_groups_pruned_statistics"), 0); + assert_eq!(unpruned.pruned("page_index_rows_pruned"), 0); + assert_eq!(unpruned.counter("output_rows"), TOTAL_ROWS); + + let output = scan(&file, list_size, Some(list_size), page_pruning).await; + output.assert_results(); + assert_eq!( + pretty_format_batches(&output.batches).unwrap().to_string(), + pretty_format_batches(&unpruned.batches) + .unwrap() + .to_string() + ); + assert_eq!( + output.plan.contains("IN_SET_INTERSECTS"), + list_size > 20, + "list_size={list_size}, plan={}", + output.plan + ); + assert_eq!( + output.pruned("row_groups_pruned_statistics"), + if page_pruning { 0 } else { 2 }, + "list_size={list_size}, metrics={}", + output.metrics + ); + assert_eq!( + output.pruned("page_index_rows_pruned"), + if page_pruning { MATCHING_ROWS } else { 0 }, + "list_size={list_size}, metrics={}", + output.metrics + ); + assert_eq!(output.counter("output_rows"), MATCHING_ROWS); + } + + // The default remains 20: enabling the compact representation must not + // silently change the public cap's meaning. + let default = scan(&file, 21, None, page_pruning).await; + default.assert_results(); + assert!(!default.plan.contains("IN_SET_INTERSECTS")); + assert_eq!(default.pruned("row_groups_pruned_statistics"), 0); + assert_eq!(default.pruned("page_index_rows_pruned"), 0); + assert_eq!(default.counter("output_rows"), TOTAL_ROWS); +} + +#[tokio::test] +async fn string_in_list_row_group_pruning() { + check_string_in_list_pruning(false).await; +} + +#[tokio::test] +async fn string_in_list_page_pruning() { + check_string_in_list_pruning(true).await; +} + +#[tokio::test] +async fn string_not_in_list_with_null_does_not_bypass_row_filter() { + let mut file = tempfile::Builder::new() + .prefix("string_not_in_list_pruning") + .suffix(".parquet") + .tempfile() + .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Utf8, true)])); + // The first row group has a known zero null count, and every value lies + // in a gap in the IN list. Dropping the NULL list member while inverting + // NOT IN would incorrectly prove that this entire row group matches. + let values = vec![ + Some("v000001"), + Some("v000001"), + Some("v000001"), + Some("v000001"), + Some("v000000"), + Some("v000001"), + None, + Some("v999999"), + ]; + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(StringArray::from(values))], + ) + .unwrap(); + let properties = WriterProperties::builder() + .set_max_row_group_row_count(Some(4)) + .set_bloom_filter_enabled(false) + .build(); + let mut writer = + ArrowWriter::try_new(&mut file, Arc::clone(&schema), Some(properties)).unwrap(); + writer.write(&batch).unwrap(); + assert_eq!(writer.close().unwrap().num_row_groups(), 2); + + // Build the physical source directly so a logical optimizer cannot fold + // the SQL NOT IN (..., NULL) filter to an empty relation before the scan. + let mut list = (0..21) + .map(|index| lit(format!("v{:06}", index * 10))) + .collect::>(); + list.push(lit(ScalarValue::Utf8(None))); + let predicate = + in_list(col("value", &schema).unwrap(), list, &true, &schema).unwrap(); + let location = Path::from_filesystem_path(file.path()).unwrap(); + let partitioned_file = PartitionedFile::new( + location.to_string(), + file.as_file().metadata().unwrap().len(), + ); + let ctx = + SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); + + for max_in_list_size in [0, 32] { + let mut options = TableParquetOptions::default(); + options.global.max_in_list_size = max_in_list_size; + let source = Arc::new( + ParquetSource::new(Arc::clone(&schema)) + .with_table_parquet_options(options) + .with_predicate(Arc::clone(&predicate)) + .with_pushdown_filters(true) + .with_enable_page_index(false) + .with_bloom_filter_on_read(false), + ); + let config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + .with_file(partitioned_file.clone()) + .with_limit(Some(1)) + .build(); + let plan: Arc = + Arc::new(DataSourceExec::new(Arc::new(config))); + let plan_text = displayable(plan.as_ref()).indent(true).to_string(); + assert!(plan_text.contains("NOT IN"), "{plan_text}"); + let batches = collect(Arc::clone(&plan), ctx.task_ctx()).await.unwrap(); + let output = ScanOutput { + batches, + plan: plan_text, + metrics: MetricsFinder::find_metrics(plan.as_ref()).unwrap(), + }; + + assert_eq!( + output + .batches + .iter() + .map(RecordBatch::num_rows) + .sum::(), + 0, + "cap={max_in_list_size}, plan={}, metrics={}", + output.plan, + output.metrics + ); + assert_eq!(output.fully_matched("row_groups_pruned_statistics"), 0); + assert_eq!(output.pruned("row_groups_pruned_statistics"), 0); + assert_eq!(output.pruned("limit_pruned_row_groups"), 0); + assert_eq!(output.counter("pushdown_rows_pruned"), 8); + assert_eq!(output.counter("predicate_evaluation_errors"), 0); + } +} diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index 0b02fcb71d663..81fbc14265ec3 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -289,9 +289,8 @@ pub(super) struct ParquetMorselizer { /// Maximum size of the predicate cache, in bytes. If none, uses /// the arrow-rs default. pub max_predicate_cache_size: Option, - /// Maximum `IN (...)` list size that the pruning predicate will rewrite - /// into per-value statistics checks. Lists longer than this skip - /// container-level pruning. Sourced from + /// Maximum `IN (...)` list size eligible for statistics pruning. Longer + /// lists skip container-level pruning. Sourced from /// `datafusion.execution.parquet.max_in_list_size`. pub max_in_list_size: usize, /// Whether to read row groups in reverse order @@ -1062,7 +1061,11 @@ impl MetadataLoadedParquetOpen { // Only build page pruning predicate if page index is enabled let page_pruning_predicate = if prepared.enable_page_index { prepared.predicate.as_ref().and_then(|predicate| { - let p = build_page_pruning_predicate(predicate, &physical_file_schema); + let p = build_page_pruning_predicate( + predicate, + &physical_file_schema, + prepared.max_in_list_size, + ); (p.filter_number() > 0).then_some(p) }) } else { @@ -1688,10 +1691,12 @@ fn create_initial_plan( pub(crate) fn build_page_pruning_predicate( predicate: &Arc, file_schema: &SchemaRef, + max_in_list_size: usize, ) -> Arc { - Arc::new(PagePruningAccessPlanFilter::new( + Arc::new(PagePruningAccessPlanFilter::new_with_max_in_list_size( predicate, Arc::clone(file_schema), + max_in_list_size, )) } @@ -1965,7 +1970,7 @@ mod test { ); let page_pruning_predicate = predicate.map(|expr| { let predicate = logical2physical(&expr, &arrow_schema); - build_page_pruning_predicate(&predicate, &arrow_schema) + build_page_pruning_predicate(&predicate, &arrow_schema, MAX_IN_LIST_SIZE) }); let store: Arc = Arc::new(InMemory::new()); diff --git a/datafusion/datasource-parquet/src/page_filter.rs b/datafusion/datasource-parquet/src/page_filter.rs index 75acd431263e1..a48c9e335457a 100644 --- a/datafusion/datasource-parquet/src/page_filter.rs +++ b/datafusion/datasource-parquet/src/page_filter.rs @@ -32,7 +32,7 @@ use arrow::{ use datafusion_common::ScalarValue; use datafusion_common::pruning::PruningStatistics; use datafusion_physical_expr::{PhysicalExpr, split_conjunction}; -use datafusion_pruning::{PruningPredicate, PruningPredicateBuilder}; +use datafusion_pruning::{MAX_IN_LIST_SIZE, PruningPredicate, PruningPredicateBuilder}; use log::{debug, trace}; use parquet::arrow::arrow_reader::statistics::StatisticsConverter; @@ -138,15 +138,25 @@ impl PagePruningResult { impl PagePruningAccessPlanFilter { /// Create a new [`PagePruningAccessPlanFilter`] from a physical - /// expression. - #[expect(clippy::needless_pass_by_value)] + /// expression, using the default `IN (...)` pruning limit. 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. + #[expect(clippy::needless_pass_by_value)] + 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_max_in_list_size(max_in_list_size) .try_build(Arc::clone(predicate)) { Ok(pp) => pp, diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index 097b4563af5df..b29480f58379b 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -486,9 +486,8 @@ impl ParquetSource { self.table_parquet_options.global.max_predicate_cache_size } - /// Return the maximum size of an `IN (...)` list that the pruning - /// predicate will rewrite into per-value statistics checks. Lists - /// longer than this skip container-level pruning. Reads from + /// Return the maximum size of an `IN (...)` list eligible for statistics + /// pruning. Longer lists skip container-level pruning. Reads from /// `datafusion.execution.parquet.max_in_list_size`. pub fn max_in_list_size(&self) -> usize { self.table_parquet_options.global.max_in_list_size diff --git a/datafusion/datasource-parquet/src/statistics_order_tests.rs b/datafusion/datasource-parquet/src/statistics_order_tests.rs index f678fa6b47ee2..943ced0440cbf 100644 --- a/datafusion/datasource-parquet/src/statistics_order_tests.rs +++ b/datafusion/datasource-parquet/src/statistics_order_tests.rs @@ -359,6 +359,74 @@ fn byte_array_order_preserves_matching_rows_at_every_pruning_level() { } } +#[test] +fn large_string_in_list_preserves_rows_with_untrusted_page_order() { + let max_in_list_size = MAX_IN_LIST_SIZE + 2; + for order in [ + StatisticsOrder::Modern, + StatisticsOrder::Missing, + StatisticsOrder::Unknown, + ] { + let file = TestFile::new(order); + // Only "az" occurs in the file. All other list members are above + // even the unsafe ["aé", "b"] interval in the first page's index. + let mut values = (0..=MAX_IN_LIST_SIZE) + .map(|index| lit(format!("z{index:03}"))) + .collect::>(); + values.push(lit("az")); + let physical = logical2physical(&col("s").in_list(values, false), &file.schema); + let predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&file.schema)) + .with_max_in_list_size(max_in_list_size) + .try_build(Arc::clone(&physical)) + .unwrap(); + assert!( + predicate + .predicate_expr() + .to_string() + .contains("IN_SET_INTERSECTS") + ); + + // Exercise the opener's configured-cap path, not the page filter's + // compatibility constructor that retains the default limit of 20. + let page_filter = crate::opener::build_page_pruning_predicate( + &physical, + &file.schema, + max_in_list_size, + ); + assert_eq!(page_filter.filter_number(), 1); + let all = ParquetAccessPlan::new_all(file.metadata.num_row_groups()); + assert_eq!(file.matching_rows(&physical, all.clone()), 1); + let file_metrics = metrics(); + let pages = page_filter.prune_plan_with_page_index( + all, + &file.schema, + file.metadata.file_metadata().schema_descr(), + &file.metadata, + &file_metrics, + ); + assert_eq!( + pages.row_group_indexes(), + if order == StatisticsOrder::Modern { + vec![0] + } else { + vec![0, 2] + }, + "order={order:?}", + ); + assert_eq!(file.matching_rows(&physical, pages), 1, "order={order:?}",); + assert_eq!( + file_metrics.page_index_rows_pruned.pruned(), + if order == StatisticsOrder::Modern { + 6 + } else { + 3 + }, + "order={order:?}", + ); + } +} + #[test] fn byte_array_order_keeps_null_counts_and_unrelated_column_bounds() { for order in [ diff --git a/datafusion/pruning/Cargo.toml b/datafusion/pruning/Cargo.toml index e6f4bb6f273c9..a703f68222f38 100644 --- a/datafusion/pruning/Cargo.toml +++ b/datafusion/pruning/Cargo.toml @@ -26,7 +26,12 @@ datafusion-physical-plan = { workspace = true } log = { workspace = true } [dev-dependencies] +criterion = { workspace = true } datafusion-expr = { workspace = true } datafusion-functions-nested = { workspace = true } insta = { workspace = true } itertools = { workspace = true } + +[[bench]] +harness = false +name = "string_in_list_pruning" diff --git a/datafusion/pruning/benches/string_in_list_pruning.rs b/datafusion/pruning/benches/string_in_list_pruning.rs new file mode 100644 index 0000000000000..d26112611bf3c --- /dev/null +++ b/datafusion/pruning/benches/string_in_list_pruning.rs @@ -0,0 +1,241 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Compare compact string IN-list pruning with per-value min/max expansion. +//! +//! Both cases raise `max_in_list_size` to the domain size. On a baseline +//! without compact pruning, `in_list` measures the ordinary raised-cap path. +//! The explicit `expanded_or` is a balanced tree of equalities, which produces +//! the same per-value statistics checks without making the baseline depend on +//! a deeply nested expression. Half of the statistics intervals hit a domain +//! member and half fall in a sparse gap. Bloom filters are not involved. +//! +//! Run with `cargo bench -p datafusion-pruning --bench string_in_list_pruning`. +//! The construction benchmarks reuse their input physical expressions; the +//! evaluation benchmarks reuse their already-built pruning predicates. + +use std::collections::HashSet; +use std::hint::black_box; +use std::sync::Arc; + +use arrow::array::{ArrayRef, BooleanArray, StringViewArray, UInt64Array}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use datafusion_common::{Column, ScalarValue}; +use datafusion_expr_common::operator::Operator; +use datafusion_physical_expr::PhysicalExprRef; +use datafusion_physical_expr::expressions::{BinaryExpr, col, in_list, lit}; +use datafusion_pruning::{PruningPredicate, PruningPredicateBuilder, PruningStatistics}; + +const DOMAIN_SIZES: [usize; 4] = [20, 21, 256, 1024]; +const CONTAINERS: usize = 4096; + +fn value(index: usize) -> String { + format!("key{index:08}") +} + +fn balanced_or(expressions: &[PhysicalExprRef]) -> PhysicalExprRef { + if expressions.len() == 1 { + return Arc::clone(&expressions[0]); + } + let middle = expressions.len() / 2; + Arc::new(BinaryExpr::new( + balanced_or(&expressions[..middle]), + Operator::Or, + balanced_or(&expressions[middle..]), + )) +} + +fn build_predicate( + expression: &PhysicalExprRef, + schema: &SchemaRef, + max_in_list_size: usize, +) -> PruningPredicate { + PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(schema)) + .with_max_in_list_size(max_in_list_size) + .try_build(Arc::clone(expression)) + .unwrap() +} + +struct IntervalStatistics { + min: ArrayRef, + max: ArrayRef, + null_counts: ArrayRef, + row_counts: ArrayRef, +} + +impl IntervalStatistics { + fn new(domain_size: usize) -> Self { + let min = StringViewArray::from_iter_values((0..CONTAINERS).map(|index| { + let start = (index / 2 % domain_size) * 10; + value(start + if index % 2 == 0 { 0 } else { 3 }) + })); + let max = StringViewArray::from_iter_values((0..CONTAINERS).map(|index| { + let start = (index / 2 % domain_size) * 10; + value(start + if index % 2 == 0 { 0 } else { 7 }) + })); + Self { + min: Arc::new(min), + max: Arc::new(max), + null_counts: Arc::new(UInt64Array::from(vec![0; CONTAINERS])), + row_counts: Arc::new(UInt64Array::from(vec![128; CONTAINERS])), + } + } +} + +impl PruningStatistics for IntervalStatistics { + fn min_values(&self, column: &Column) -> Option { + (column.name == "value").then(|| Arc::clone(&self.min)) + } + + fn max_values(&self, column: &Column) -> Option { + (column.name == "value").then(|| Arc::clone(&self.max)) + } + + fn num_containers(&self) -> usize { + CONTAINERS + } + + fn null_counts(&self, column: &Column) -> Option { + (column.name == "value").then(|| Arc::clone(&self.null_counts)) + } + + fn row_counts(&self) -> Option { + Some(Arc::clone(&self.row_counts)) + } + + fn contained( + &self, + _column: &Column, + _values: &HashSet, + ) -> Option { + None + } +} + +struct BenchmarkCase { + size: usize, + schema: SchemaRef, + in_list: PhysicalExprRef, + expanded_or: PhysicalExprRef, + in_list_predicate: PruningPredicate, + expanded_or_predicate: PruningPredicate, + statistics: IntervalStatistics, +} + +impl BenchmarkCase { + fn new(size: usize) -> Self { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Utf8View, + false, + )])); + let column = col("value", &schema).unwrap(); + let values = (0..size) + .map(|index| lit(ScalarValue::new_utf8view(value(index * 10)))) + .collect::>(); + let in_list = + in_list(Arc::clone(&column), values.clone(), &false, &schema).unwrap(); + let equalities = values + .into_iter() + .map(|value| { + Arc::new(BinaryExpr::new(Arc::clone(&column), Operator::Eq, value)) + as PhysicalExprRef + }) + .collect::>(); + let expanded_or = balanced_or(&equalities); + let in_list_predicate = build_predicate(&in_list, &schema, size); + let expanded_or_predicate = build_predicate(&expanded_or, &schema, size); + eprintln!( + "string_in_list_pruning: {size} values, compact={}", + in_list_predicate + .predicate_expr() + .to_string() + .contains("IN_SET_INTERSECTS") + ); + let statistics = IntervalStatistics::new(size); + + // Check that both benchmark paths do the same useful work, rather + // than comparing compact pruning with an always-true fallback. + let expected = (0..CONTAINERS) + .map(|index| index % 2 == 0) + .collect::>(); + assert_eq!(in_list_predicate.prune(&statistics).unwrap(), expected); + assert_eq!(expanded_or_predicate.prune(&statistics).unwrap(), expected); + + Self { + size, + schema, + in_list, + expanded_or, + in_list_predicate, + expanded_or_predicate, + statistics, + } + } +} + +fn criterion_benchmark(criterion: &mut Criterion) { + let cases = DOMAIN_SIZES.map(BenchmarkCase::new); + let mut construction = criterion.benchmark_group("string_in_list_pruning/construct"); + for case in &cases { + construction.throughput(Throughput::Elements(case.size as u64)); + for (name, expression) in [ + ("in_list", &case.in_list), + ("expanded_or", &case.expanded_or), + ] { + construction.bench_with_input( + BenchmarkId::new(name, case.size), + expression, + |bencher, expression| { + bencher.iter(|| { + black_box(build_predicate( + black_box(expression), + &case.schema, + case.size, + )) + }); + }, + ); + } + } + construction.finish(); + + let mut evaluation = criterion.benchmark_group("string_in_list_pruning/evaluate"); + evaluation.throughput(Throughput::Elements(CONTAINERS as u64)); + for case in &cases { + for (name, predicate) in [ + ("in_list", &case.in_list_predicate), + ("expanded_or", &case.expanded_or_predicate), + ] { + evaluation.bench_with_input( + BenchmarkId::new(name, case.size), + predicate, + |bencher, predicate| { + bencher.iter(|| { + black_box(predicate.prune(black_box(&case.statistics)).unwrap()) + }); + }, + ); + } + } + evaluation.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/pruning/src/lib.rs b/datafusion/pruning/src/lib.rs index 2b334d2847980..6bf1815900aa8 100644 --- a/datafusion/pruning/src/lib.rs +++ b/datafusion/pruning/src/lib.rs @@ -19,6 +19,7 @@ mod file_pruner; mod pruning_predicate; +mod string_in_list; pub use file_pruner::FilePruner; pub use pruning_predicate::{ diff --git a/datafusion/pruning/src/pruning_predicate.rs b/datafusion/pruning/src/pruning_predicate.rs index 6df5cbc5cc135..6a63775fbede1 100644 --- a/datafusion/pruning/src/pruning_predicate.rs +++ b/datafusion/pruning/src/pruning_predicate.rs @@ -22,6 +22,8 @@ use std::collections::HashSet; use std::sync::Arc; +use crate::string_in_list::StringInListPruningExpr; + use arrow::array::AsArray; use arrow::{ array::{ArrayRef, BooleanArray, new_null_array}, @@ -441,9 +443,11 @@ 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 size of `IN (...)` lists eligible for statistics pruning. + /// Positive literal string lists larger than [`MAX_IN_LIST_SIZE`] use a + /// compact sorted domain; other eligible lists use per-value checks. + /// Lists longer than this cap fall back to the unhandled-predicate hook + /// (typically "keep the container"). /// /// Query engines typically pass /// `datafusion.execution.parquet.max_in_list_size` here. @@ -1456,10 +1460,51 @@ fn build_is_null_column_expr( } } -/// Default maximum number of entries in an `IN (...)` list that will be -/// rewritten into a chain of per-value min/max checks by -/// `build_predicate_expression`. Callers threading a [`PredicateRewriter`] -/// can override this via [`PredicateRewriter::with_max_in_list_size`], and +/// Keep large literal string domains compact instead of building an OR tree. +fn build_string_in_list_expr( + in_list: &phys_expr::InListExpr, + schema: &Schema, + required_columns: &mut RequiredColumns, +) -> Option> { + if in_list.negated() { + return None; + } + let column = in_list.expr().downcast_ref::()?; + let field = schema.fields().get(column.index())?; + let data_type = match field.data_type() { + DataType::Dictionary(_, value) => value.as_ref(), + data_type => data_type, + }; + if field.name() != column.name() || !data_type.is_string() { + return None; + } + // NULLs must remain unhandled: the inverse predicate is also used to prove + // that every row matches, and IN (..., NULL) can evaluate to UNKNOWN. + let values = in_list + .list() + .iter() + .map(|expr| extract_string_literal(expr).map(str::to_owned)) + .collect::>>()?; + let min = required_columns + .min_column_expr(column, in_list.expr(), field) + .ok()?; + let max = required_columns + .max_column_expr(column, in_list.expr(), field) + .ok()?; + let non_null = + build_is_null_column_expr(in_list.expr(), schema, required_columns, true)?; + let intersects = Arc::new(StringInListPruningExpr::new(min, max, values)); + Some(Arc::new(phys_expr::BinaryExpr::new( + non_null, + Operator::And, + intersects, + ))) +} + +/// Default maximum number of entries in an `IN (...)` list eligible for +/// statistics pruning. Eligible positive literal string lists above this +/// threshold use a compact sorted domain instead of per-value min/max checks. +/// Callers can raise the cap via [`PredicateRewriter::with_max_in_list_size`], and /// query engines can wire it from the /// `datafusion.execution.parquet.max_in_list_size` config option. pub const MAX_IN_LIST_SIZE: usize = 20; @@ -1495,14 +1540,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 + /// Set the maximum size of an `IN (...)` list eligible for statistics + /// pruning. Large positive literal string lists use a compact sorted + /// domain; other eligible lists use per-value checks. Longer lists fall back /// to the unhandled-predicate hook (typically "keep the container"), /// effectively skipping container-level pruning for large IN lists. /// /// 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 @@ -1542,8 +1588,9 @@ impl PredicateRewriter { /// /// Returns the pruning predicate as an [`PhysicalExpr`] /// -/// `max_in_list_size` is the largest `IN (...)` list that will be rewritten -/// into a chain of per-value statistics checks; longer lists fall back to +/// `max_in_list_size` is the largest `IN (...)` list eligible for statistics +/// pruning. Large positive literal string lists use a compact sorted domain; +/// other eligible lists use per-value checks. Longer lists fall back to /// `unhandled_hook`. fn build_predicate_expression( expr: &Arc, @@ -1585,6 +1632,13 @@ fn build_predicate_expression( } } if let Some(in_list) = expr.downcast_ref::() { + if in_list.list().len() > MAX_IN_LIST_SIZE + && in_list.list().len() <= max_in_list_size + && let Some(pruning_expr) = + build_string_in_list_expr(in_list, schema, required_columns) + { + return pruning_expr; + } if !in_list.list().is_empty() && in_list.list().len() <= max_in_list_size { let eq_op = if in_list.negated() { Operator::NotEq @@ -3577,6 +3631,244 @@ mod tests { Ok(()) } + fn large_string_pruning_predicate( + expr: PhysicalExprRef, + schema: SchemaRef, + ) -> Result { + PruningPredicateBuilder::new() + .with_file_schema(schema) + .with_max_in_list_size(10_000) + .try_build(expr) + } + + #[test] + fn large_string_in_list_prunes_exact_intervals() -> Result<()> { + let types = [ + DataType::Utf8, + DataType::LargeUtf8, + DataType::Utf8View, + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + ]; + for count in [20, 21, 256, 10_000] { + for data_type in &types { + let schema = Arc::new(Schema::new(vec![Field::new( + "c1", + data_type.clone(), + true, + )])); + let values = (0..count) + .map(|i| { + let value = ScalarValue::from(format!("k{:06}", i * 10)) + .cast_to(data_type)?; + Ok(Arc::new(phys_expr::Literal::new(value)) as PhysicalExprRef) + }) + .collect::>>()?; + let expr = phys_expr::in_list( + Arc::new(phys_expr::Column::new("c1", 0)), + values, + &false, + &schema, + )?; + let predicate = + large_string_pruning_predicate(Arc::clone(&expr), schema)?; + let stats = TestStatistics::new().with( + "c1", + ContainerStats::new_utf8( + [ + Some("k000000"), + Some("k000001"), + Some("k000010"), + Some("j"), + Some("z"), + None, + Some("k000020"), + Some("k000005"), + ], + [ + Some("k000000"), + Some("k000009"), + Some("k000010"), + Some("j9"), + Some("z9"), + Some("k000010"), + None, + Some("k000015"), + ], + ) + .with_null_counts([Some(0); 8]) + .with_row_counts([Some(1); 8]), + ); + assert_eq!( + predicate.prune(&stats)?, + [true, false, true, false, false, true, 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_handles_unicode_and_unknown_bounds() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)])); + let mut values = (0..21).map(|i| lit(format!("m{i:03}"))).collect::>(); + values.extend([ + lit(""), + lit("az"), + lit("aé"), + lit("aé"), + lit("é"), + lit("🦀"), + ]); + let expr = col("c1").in_list(values, false); + let predicate = + large_string_pruning_predicate(logical2physical(&expr, &schema), schema)?; + let stats = TestStatistics::new().with( + "c1", + ContainerStats::new_utf8( + [ + Some(""), + Some("az"), + Some("a{"), + Some("aé"), + Some("ê"), + Some("z"), + None, + None, + ], + [ + Some(""), + Some("az"), + Some("aè"), + Some("aé"), + Some("🦀"), + Some("m"), + Some("z"), + None, + ], + ) + .with_null_counts([ + Some(0), + Some(0), + Some(0), + Some(0), + Some(0), + Some(0), + Some(0), + Some(1), + ]) + .with_row_counts([Some(1); 8]), + ); + assert_eq!( + predicate.prune(&stats)?, + [true, true, false, true, true, true, true, false] + ); + Ok(()) + } + + #[test] + fn large_string_in_list_respects_configured_limit() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)])); + let values = (0..21).map(|i| lit(format!("a{i:03}"))).collect::>(); + let expr = logical2physical(&col("c1").in_list(values, false), &schema); + + for (limit, compact) in [(0, false), (20, false), (21, true), (32, true)] { + let predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .with_max_in_list_size(limit) + .try_build(Arc::clone(&expr))?; + assert_eq!( + predicate + .predicate_expr() + .to_string() + .contains("IN_SET_INTERSECTS"), + compact, + "limit={limit}" + ); + assert_eq!(is_always_true(predicate.predicate_expr()), !compact); + } + + let default = PruningPredicateBuilder::new() + .with_file_schema(schema) + .try_build(expr)?; + assert!(is_always_true(default.predicate_expr())); + Ok(()) + } + + #[test] + fn large_string_in_list_keeps_null_and_not_in_semantics() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)])); + let values = (0..21).map(|i| lit(format!("a{i:03}"))).collect::>(); + let stats = TestStatistics::new().with( + "c1", + ContainerStats::new_utf8( + [Some("middle"), Some("other")], + [Some("middle"), Some("other")], + ) + .with_null_counts([Some(0); 2]) + .with_row_counts([Some(1); 2]), + ); + let positive = col("c1").in_list(values.clone(), false); + let predicate = large_string_pruning_predicate( + logical2physical(&positive.or(col("c1").eq(lit("middle"))), &schema), + Arc::clone(&schema), + )?; + assert_eq!(predicate.prune(&stats)?, [true, false]); + + let mut with_null = values.clone(); + with_null.push(lit(ScalarValue::Utf8(None))); + for expr in [ + col("c1").in_list(values, true), + col("c1").in_list(with_null.clone(), false), + col("c1").in_list(with_null.clone(), true), + ] { + let physical = logical2physical(&expr, &schema); + let default = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(Arc::clone(&physical))?; + assert!(is_always_true(default.predicate_expr()), "{expr}"); + assert_eq!(default.prune(&stats)?, [true, true]); + + // Raising the cap retains the existing per-value rewrite for + // NOT IN and lists containing NULL; neither uses the new path. + let raised = large_string_pruning_predicate(physical, Arc::clone(&schema))?; + assert!( + !raised + .predicate_expr() + .to_string() + .contains("IN_SET_INTERSECTS"), + "{expr}" + ); + } + + // Inverting NOT IN (..., NULL) must not prove a full match and bypass + // the original row filter, which returns UNKNOWN for both rows. + let not_in = logical2physical(&col("c1").in_list(with_null, true), &schema); + let batch = RecordBatch::try_new( + schema, + vec![Arc::new(StringArray::from(vec!["middle", "other"]))], + )?; + assert_eq!(not_in.evaluate(&batch)?.into_array(2)?.null_count(), 2); + Ok(()) + } + #[test] fn row_group_predicate_cast_int_int() -> Result<()> { let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]); diff --git a/datafusion/pruning/src/string_in_list.rs b/datafusion/pruning/src/string_in_list.rs new file mode 100644 index 0000000000000..5b7e56cb70122 --- /dev/null +++ b/datafusion/pruning/src/string_in_list.rs @@ -0,0 +1,134 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::fmt::{self, Display, Formatter}; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +use arrow::array::{Array, AsArray, BooleanArray}; +use arrow::compute::cast; +use arrow::datatypes::{DataType, Schema}; +use arrow::record_batch::RecordBatch; +use datafusion_common::{Result, assert_eq_or_internal_err}; +use datafusion_physical_expr::{PhysicalExpr, PhysicalExprRef}; +use datafusion_physical_plan::ColumnarValue; + +/// Tests whether a sorted string domain intersects an inclusive statistics interval. +/// This expression is used only for pruning; the original IN remains the row filter. +#[derive(Debug, Eq)] +pub(crate) struct StringInListPruningExpr { + min: PhysicalExprRef, + max: PhysicalExprRef, + values: Arc<[String]>, +} + +impl StringInListPruningExpr { + pub(crate) fn new( + min: PhysicalExprRef, + max: PhysicalExprRef, + mut values: Vec, + ) -> Self { + values.sort_unstable(); + values.dedup(); + Self { + min, + max, + values: values.into(), + } + } +} + +impl PartialEq for StringInListPruningExpr { + fn eq(&self, other: &Self) -> bool { + self.min.eq(&other.min) && self.max.eq(&other.max) && self.values == other.values + } +} + +impl Hash for StringInListPruningExpr { + fn hash(&self, state: &mut H) { + self.min.hash(state); + self.max.hash(state); + self.values.hash(state); + } +} + +impl Display for StringInListPruningExpr { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!( + f, + "IN_SET_INTERSECTS({}, {}, {} values)", + self.min, + self.max, + self.values.len() + ) + } +} + +impl PhysicalExpr for StringInListPruningExpr { + fn data_type(&self, _input_schema: &Schema) -> Result { + Ok(DataType::Boolean) + } + + fn nullable(&self, _input_schema: &Schema) -> Result { + Ok(true) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + // Normalize Utf8, LargeUtf8, Utf8View, and dictionary-encoded statistics. + let min = self.min.evaluate(batch)?.into_array(batch.num_rows())?; + let max = self.max.evaluate(batch)?.into_array(batch.num_rows())?; + let min = cast(&min, &DataType::Utf8View)?; + let max = cast(&max, &DataType::Utf8View)?; + let min = min.as_string_view(); + let max = max.as_string_view(); + let matches: BooleanArray = (0..batch.num_rows()) + .map(|i| { + if min.is_null(i) || max.is_null(i) { + return None; + } + let min = min.value(i).as_bytes(); + let max = max.value(i).as_bytes(); + if min > max { + return None; + } + let index = self.values.partition_point(|v| v.as_bytes() < min); + Some(self.values.get(index).is_some_and(|v| v.as_bytes() <= 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 { + 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}") + } +}