diff --git a/Cargo.lock b/Cargo.lock index a09126ce1e84b..45cb27663871d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2582,6 +2582,7 @@ name = "datafusion-pruning" version = "55.0.0" dependencies = [ "arrow", + "criterion", "datafusion-common", "datafusion-datasource", "datafusion-expr", diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index b340ec86283d8..7087c57129658 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1368,13 +1368,15 @@ config_namespace! { /// parquet reader setting. 0 means no caching. pub max_predicate_cache_size: Option, default = None - /// Maximum number of values in an `IN (...)` list for which pruning will - /// occur. Longer lists will not be used to prune files, row groups, or - /// data pages. + /// Maximum number of input values in an `IN (...)` list eligible for + /// min/max pruning. Lists above this cap, or a cap of 0, skip this + /// rewrite; other predicates and Bloom-filter pruning remain available. /// - /// Higher values help in cases such as filtering on a list of - /// ~25-100 identifiers, but also make the predicate more expensive to - /// evaluate. Set to 0 to disable `IN (...)` list pruning entirely. + /// Within the cap, nonempty lists of at most 20 values use the existing + /// per-value rewrite. Larger positive, non-null literal string lists + /// on a string column use a compact sorted domain. Other lists retain + /// the existing per-value rewrite, so raising the cap can make those + /// predicates expensive to build and evaluate. /// /// Defaults to 20. pub max_in_list_size: usize, default = 20 diff --git a/datafusion/common/src/pruning.rs b/datafusion/common/src/pruning.rs index a36ac9f795b95..9f6a95c0978d4 100644 --- a/datafusion/common/src/pruning.rs +++ b/datafusion/common/src/pruning.rs @@ -67,6 +67,13 @@ pub trait PruningStatistics { /// returned array should have `null` in that row. If the minimum value is /// not known for any row, return `None`. /// + /// Each non-null entry must be a conservative lower bound for all non-null + /// values in its container, using Arrow's comparison order for the column + /// type. For strings, this is unsigned lexicographic UTF-8 byte order. + /// Bounds that cannot satisfy this contract must be reported as unknown. + /// Pruning relies on providers to uphold this contract; it does not validate + /// the ordering or bound guarantees of arbitrary statistics sources. + /// /// Note: the returned array must contain [`Self::num_containers`] rows fn min_values(&self, column: &Column) -> Option; @@ -74,6 +81,10 @@ pub trait PruningStatistics { /// /// See [`Self::min_values`] for when to return `None` and null values. /// + /// Each non-null entry must be a conservative upper bound for all non-null + /// values in its container, using the comparison order described in + /// [`Self::min_values`]. + /// /// Note: the returned array must contain [`Self::num_containers`] rows fn max_values(&self, column: &Column) -> Option; diff --git a/datafusion/core/tests/parquet/mod.rs b/datafusion/core/tests/parquet/mod.rs index 7066a4147c017..cf2e557abad5f 100644 --- a/datafusion/core/tests/parquet/mod.rs +++ b/datafusion/core/tests/parquet/mod.rs @@ -58,6 +58,7 @@ mod page_pruning; mod row_group_pruning; mod schema; mod schema_coercion; +mod string_in_list_pruning; mod utils; #[cfg(test)] diff --git a/datafusion/core/tests/parquet/string_in_list_pruning.rs b/datafusion/core/tests/parquet/string_in_list_pruning.rs new file mode 100644 index 0000000000000..e1005daaa1a38 --- /dev/null +++ b/datafusion/core/tests/parquet/string_in_list_pruning.rs @@ -0,0 +1,363 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! End-to-end coverage for compact, large string IN-list pruning. The positive +//! IN-list cases disable the row and Bloom filters to isolate min/max pruning. + +use std::sync::Arc; + +use arrow::array::StringArray; +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::record_batch::RecordBatch; +use arrow::util::pretty::pretty_format_batches; +use datafusion::datasource::listing::PartitionedFile; +use datafusion::datasource::object_store::ObjectStoreUrl; +use datafusion::datasource::physical_plan::ParquetSource; +use datafusion::datasource::source::DataSourceExec; +use datafusion::physical_plan::{ExecutionPlan, collect, displayable}; +use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; +use datafusion_common::config::TableParquetOptions; +use datafusion_common::{ScalarValue, assert_batches_eq}; +use datafusion_datasource::file_scan_config::FileScanConfigBuilder; +use datafusion_physical_expr::expressions::{col, in_list, lit}; +use datafusion_physical_plan::metrics::{MetricValue, MetricsSet}; +use object_store::path::Path; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::{EnabledStatistics, WriterProperties}; +use tempfile::NamedTempFile; + +use super::utils::MetricsFinder; + +const ROWS_PER_UNIT: usize = 16; +const UNITS: usize = 4; +const TOTAL_ROWS: usize = ROWS_PER_UNIT * UNITS; +const MATCHING_ROWS: usize = ROWS_PER_UNIT * 2; + +/// Write either four row groups or four pages in one row group. The second +/// unit lies in a gap between two members of every test IN list; an enclosing +/// min/max range for the list cannot prune it. +fn make_file(page_pruning: bool) -> NamedTempFile { + let mut file = tempfile::Builder::new() + .prefix("string_in_list_pruning") + .suffix(".parquet") + .tempfile() + .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Utf8, + false, + )])); + let values = ["v000000", "v000001", "v000010", "v999999"] + .into_iter() + .flat_map(|value| std::iter::repeat_n(value, ROWS_PER_UNIT)) + .collect::>(); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(StringArray::from(values))], + ) + .unwrap(); + let rows_per_group = if page_pruning { + TOTAL_ROWS + } else { + ROWS_PER_UNIT + }; + let properties = WriterProperties::builder() + .set_max_row_group_row_count(Some(rows_per_group)) + .set_data_page_row_count_limit(ROWS_PER_UNIT) + .set_write_batch_size(ROWS_PER_UNIT) + .set_dictionary_enabled(false) + .set_bloom_filter_enabled(false) + .set_statistics_enabled(EnabledStatistics::Page) + .build(); + let mut writer = ArrowWriter::try_new(&mut file, schema, Some(properties)).unwrap(); + writer.write(&batch).unwrap(); + let metadata = writer.close().unwrap(); + assert_eq!(metadata.num_row_groups(), TOTAL_ROWS / rows_per_group); + let offsets = metadata.offset_index().unwrap(); + for row_group in offsets { + assert_eq!( + row_group[0].page_locations().len(), + rows_per_group / ROWS_PER_UNIT + ); + } + file +} + +struct ScanOutput { + batches: Vec, + plan: String, + metrics: MetricsSet, +} + +impl ScanOutput { + fn counter(&self, name: &str) -> usize { + self.metrics + .sum(|metric| metric.value().name() == name) + .unwrap_or_else(|| panic!("missing {name}: {}", self.metrics)) + .as_usize() + } + + fn pruned(&self, name: &str) -> usize { + let value = self + .metrics + .sum(|metric| metric.value().name() == name) + .unwrap_or_else(|| panic!("missing {name}: {}", self.metrics)); + let MetricValue::PruningMetrics { + pruning_metrics, .. + } = value + else { + panic!("expected pruning metric {name}: {}", self.metrics); + }; + pruning_metrics.pruned() + } + + fn fully_matched(&self, name: &str) -> usize { + let value = self + .metrics + .sum(|metric| metric.value().name() == name) + .unwrap_or_else(|| panic!("missing {name}: {}", self.metrics)); + let MetricValue::PruningMetrics { + pruning_metrics, .. + } = value + else { + panic!("expected pruning metric {name}: {}", self.metrics); + }; + pruning_metrics.fully_matched() + } + + fn assert_results(&self) { + assert_batches_eq!( + [ + "+---------+----+", + "| value | n |", + "+---------+----+", + "| v000000 | 16 |", + "| v000010 | 16 |", + "+---------+----+", + ], + &self.batches + ); + assert_eq!(self.counter("predicate_evaluation_errors"), 0); + assert_eq!(self.counter("pushdown_rows_pruned"), 0); + assert_eq!(self.pruned("row_groups_pruned_bloom_filter"), 0); + } +} + +async fn scan( + file: &NamedTempFile, + list_size: usize, + max_in_list_size: Option, + page_pruning: bool, +) -> ScanOutput { + let mut config = SessionConfig::new() + .with_target_partitions(1) + .with_parquet_bloom_filter_pruning(false) + .with_parquet_page_index_pruning(page_pruning); + config.options_mut().execution.parquet.pushdown_filters = false; + if let Some(max_in_list_size) = max_in_list_size { + config.options_mut().execution.parquet.max_in_list_size = max_in_list_size; + } + let ctx = SessionContext::new_with_config(config); + ctx.register_parquet( + "t", + file.path().to_str().unwrap(), + ParquetReadOptions::default(), + ) + .await + .unwrap(); + let values = (0..list_size) + .map(|index| format!("'v{:06}'", index * 10)) + .collect::>() + .join(", "); + let sql = format!( + "SELECT value, count(*) AS n FROM t \ + WHERE value IN ({values}) GROUP BY value ORDER BY value" + ); + let plan = ctx + .sql(&sql) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let plan_text = displayable(plan.as_ref()).indent(true).to_string(); + let batches = collect(Arc::clone(&plan), ctx.task_ctx()).await.unwrap(); + let metrics = MetricsFinder::find_metrics(plan.as_ref()).unwrap(); + ScanOutput { + batches, + plan: plan_text, + metrics, + } +} + +async fn check_string_in_list_pruning(page_pruning: bool) { + let file = make_file(page_pruning); + for list_size in [20, 21, 256, 1024] { + // A zero cap provides a result-equivalence control that cannot use + // min/max IN-list pruning at either granularity. + let unpruned = scan(&file, list_size, Some(0), page_pruning).await; + unpruned.assert_results(); + assert!(!unpruned.plan.contains("IN_SET_INTERSECTS")); + assert_eq!(unpruned.pruned("row_groups_pruned_statistics"), 0); + assert_eq!(unpruned.pruned("page_index_rows_pruned"), 0); + assert_eq!(unpruned.counter("output_rows"), TOTAL_ROWS); + + let output = scan(&file, list_size, Some(list_size), page_pruning).await; + output.assert_results(); + assert_eq!( + pretty_format_batches(&output.batches).unwrap().to_string(), + pretty_format_batches(&unpruned.batches) + .unwrap() + .to_string() + ); + assert_eq!( + output.plan.contains("IN_SET_INTERSECTS"), + list_size > 20, + "list_size={list_size}, plan={}", + output.plan + ); + assert_eq!( + output.pruned("row_groups_pruned_statistics"), + if page_pruning { 0 } else { 2 }, + "list_size={list_size}, metrics={}", + output.metrics + ); + assert_eq!( + output.pruned("page_index_rows_pruned"), + if page_pruning { MATCHING_ROWS } else { 0 }, + "list_size={list_size}, metrics={}", + output.metrics + ); + assert_eq!(output.counter("output_rows"), MATCHING_ROWS); + } + + // The default remains 20: enabling the compact representation must not + // silently change the public cap's meaning. + let default = scan(&file, 21, None, page_pruning).await; + default.assert_results(); + assert!(!default.plan.contains("IN_SET_INTERSECTS")); + assert_eq!(default.pruned("row_groups_pruned_statistics"), 0); + assert_eq!(default.pruned("page_index_rows_pruned"), 0); + assert_eq!(default.counter("output_rows"), TOTAL_ROWS); +} + +#[tokio::test] +async fn string_in_list_row_group_pruning() { + check_string_in_list_pruning(false).await; +} + +#[tokio::test] +async fn string_in_list_page_pruning() { + check_string_in_list_pruning(true).await; +} + +#[tokio::test] +async fn string_not_in_list_with_null_does_not_bypass_row_filter() { + let mut file = tempfile::Builder::new() + .prefix("string_not_in_list_pruning") + .suffix(".parquet") + .tempfile() + .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Utf8, true)])); + // The first row group has a known zero null count, and every value lies + // in a gap in the IN list. Dropping the NULL list member while inverting + // NOT IN would incorrectly prove that this entire row group matches. + let values = vec![ + Some("v000001"), + Some("v000001"), + Some("v000001"), + Some("v000001"), + Some("v000000"), + Some("v000001"), + None, + Some("v999999"), + ]; + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(StringArray::from(values))], + ) + .unwrap(); + let properties = WriterProperties::builder() + .set_max_row_group_row_count(Some(4)) + .set_bloom_filter_enabled(false) + .build(); + let mut writer = + ArrowWriter::try_new(&mut file, Arc::clone(&schema), Some(properties)).unwrap(); + writer.write(&batch).unwrap(); + assert_eq!(writer.close().unwrap().num_row_groups(), 2); + + // Build the physical source directly so a logical optimizer cannot fold + // the SQL NOT IN (..., NULL) filter to an empty relation before the scan. + let mut list = (0..21) + .map(|index| lit(format!("v{:06}", index * 10))) + .collect::>(); + list.push(lit(ScalarValue::Utf8(None))); + let predicate = + in_list(col("value", &schema).unwrap(), list, &true, &schema).unwrap(); + let location = Path::from_filesystem_path(file.path()).unwrap(); + let partitioned_file = PartitionedFile::new( + location.to_string(), + file.as_file().metadata().unwrap().len(), + ); + let ctx = + SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); + + for max_in_list_size in [0, 32] { + let mut options = TableParquetOptions::default(); + options.global.max_in_list_size = max_in_list_size; + let source = Arc::new( + ParquetSource::new(Arc::clone(&schema)) + .with_table_parquet_options(options) + .with_predicate(Arc::clone(&predicate)) + .with_pushdown_filters(true) + .with_enable_page_index(false) + .with_bloom_filter_on_read(false), + ); + let config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + .with_file(partitioned_file.clone()) + .with_limit(Some(1)) + .build(); + let plan: Arc = + Arc::new(DataSourceExec::new(Arc::new(config))); + let plan_text = displayable(plan.as_ref()).indent(true).to_string(); + assert!(plan_text.contains("NOT IN"), "{plan_text}"); + let batches = collect(Arc::clone(&plan), ctx.task_ctx()).await.unwrap(); + let output = ScanOutput { + batches, + plan: plan_text, + metrics: MetricsFinder::find_metrics(plan.as_ref()).unwrap(), + }; + + assert_eq!( + output + .batches + .iter() + .map(RecordBatch::num_rows) + .sum::(), + 0, + "cap={max_in_list_size}, plan={}, metrics={}", + output.plan, + output.metrics + ); + assert_eq!(output.fully_matched("row_groups_pruned_statistics"), 0); + assert_eq!(output.pruned("row_groups_pruned_statistics"), 0); + assert_eq!(output.pruned("limit_pruned_row_groups"), 0); + assert_eq!(output.counter("pushdown_rows_pruned"), 8); + assert_eq!(output.counter("predicate_evaluation_errors"), 0); + } +} diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index 88241f23e79f8..25c3bc9a77851 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -290,9 +290,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 @@ -1085,7 +1084,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 { @@ -1814,10 +1817,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, )) } @@ -2091,7 +2096,7 @@ mod test { ); let page_pruning_predicate = predicate.map(|expr| { let predicate = logical2physical(&expr, &arrow_schema); - build_page_pruning_predicate(&predicate, &arrow_schema) + build_page_pruning_predicate(&predicate, &arrow_schema, MAX_IN_LIST_SIZE) }); let store: Arc = Arc::new(InMemory::new()); diff --git a/datafusion/datasource-parquet/src/page_filter.rs b/datafusion/datasource-parquet/src/page_filter.rs index 5372ae0a0007d..236a08c5608af 100644 --- a/datafusion/datasource-parquet/src/page_filter.rs +++ b/datafusion/datasource-parquet/src/page_filter.rs @@ -32,7 +32,7 @@ use arrow::{ use datafusion_common::ScalarValue; use datafusion_common::pruning::PruningStatistics; use datafusion_physical_expr::{PhysicalExpr, split_conjunction}; -use datafusion_pruning::{PruningPredicate, PruningPredicateBuilder}; +use datafusion_pruning::{MAX_IN_LIST_SIZE, PruningPredicate, PruningPredicateBuilder}; use log::{debug, trace}; use parquet::arrow::arrow_reader::statistics::StatisticsConverter; @@ -138,15 +138,25 @@ impl PagePruningResult { impl PagePruningAccessPlanFilter { /// Create a new [`PagePruningAccessPlanFilter`] from a physical - /// expression. + /// expression, using the default `IN (...)` pruning limit. #[expect(clippy::needless_pass_by_value)] pub fn new(expr: &Arc, schema: SchemaRef) -> Self { + Self::new_with_max_in_list_size(expr, &schema, MAX_IN_LIST_SIZE) + } + + /// Create a page filter using the same `IN (...)` limit as row-group pruning. + pub(crate) fn new_with_max_in_list_size( + expr: &Arc, + schema: &SchemaRef, + max_in_list_size: usize, + ) -> Self { // extract any single column predicates let predicates = split_conjunction(expr) .into_iter() .filter_map(|predicate| { let pp = match PruningPredicateBuilder::new() - .with_file_schema(Arc::clone(&schema)) + .with_file_schema(Arc::clone(schema)) + .with_max_in_list_size(max_in_list_size) .try_build(Arc::clone(predicate)) { Ok(pp) => pp, diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index c5506dd449681..28bdf71065427 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -486,9 +486,8 @@ impl ParquetSource { self.table_parquet_options.global.max_predicate_cache_size } - /// Return the maximum size of an `IN (...)` list that the pruning - /// predicate will rewrite into per-value statistics checks. Lists - /// longer than this skip container-level pruning. Reads from + /// Return the maximum size of an `IN (...)` list eligible for statistics + /// pruning. Longer lists skip container-level pruning. Reads from /// `datafusion.execution.parquet.max_in_list_size`. pub fn max_in_list_size(&self) -> usize { self.table_parquet_options.global.max_in_list_size diff --git a/datafusion/datasource-parquet/src/statistics_order_tests.rs b/datafusion/datasource-parquet/src/statistics_order_tests.rs index bf9df532f1b58..5081abf2cb47b 100644 --- a/datafusion/datasource-parquet/src/statistics_order_tests.rs +++ b/datafusion/datasource-parquet/src/statistics_order_tests.rs @@ -359,6 +359,74 @@ fn byte_array_order_preserves_matching_rows_at_every_pruning_level() { } } +#[test] +fn large_string_in_list_preserves_rows_with_untrusted_page_order() { + let max_in_list_size = MAX_IN_LIST_SIZE + 2; + for order in [ + StatisticsOrder::Modern, + StatisticsOrder::Missing, + StatisticsOrder::Unknown, + ] { + let file = TestFile::new(order); + // Only "az" occurs in the file. All other list members are above + // even the unsafe ["aé", "b"] interval in the first page's index. + let mut values = (0..=MAX_IN_LIST_SIZE) + .map(|index| lit(format!("z{index:03}"))) + .collect::>(); + values.push(lit("az")); + let physical = logical2physical(&col("s").in_list(values, false), &file.schema); + let predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&file.schema)) + .with_max_in_list_size(max_in_list_size) + .try_build(Arc::clone(&physical)) + .unwrap(); + assert!( + predicate + .predicate_expr() + .to_string() + .contains("IN_SET_INTERSECTS") + ); + + // Exercise the opener's configured-cap path, not the page filter's + // compatibility constructor that retains the default limit of 20. + let page_filter = crate::opener::build_page_pruning_predicate( + &physical, + &file.schema, + max_in_list_size, + ); + assert_eq!(page_filter.filter_number(), 1); + let all = ParquetAccessPlan::new_all(file.metadata.num_row_groups()); + assert_eq!(file.matching_rows(&physical, all.clone()), 1); + let file_metrics = metrics(); + let pages = page_filter.prune_plan_with_page_index( + all, + &file.schema, + file.metadata.file_metadata().schema_descr(), + &file.metadata, + &file_metrics, + ); + assert_eq!( + pages.row_group_indexes(), + if order == StatisticsOrder::Modern { + vec![0] + } else { + vec![0, 2] + }, + "order={order:?}", + ); + assert_eq!(file.matching_rows(&physical, pages), 1, "order={order:?}",); + assert_eq!( + file_metrics.page_index_rows_pruned.pruned(), + if order == StatisticsOrder::Modern { + 6 + } else { + 3 + }, + "order={order:?}", + ); + } +} + #[test] fn byte_array_order_keeps_null_counts_and_unrelated_column_bounds() { for order in [ diff --git a/datafusion/pruning/Cargo.toml b/datafusion/pruning/Cargo.toml index e6f4bb6f273c9..a703f68222f38 100644 --- a/datafusion/pruning/Cargo.toml +++ b/datafusion/pruning/Cargo.toml @@ -26,7 +26,12 @@ datafusion-physical-plan = { workspace = true } log = { workspace = true } [dev-dependencies] +criterion = { workspace = true } datafusion-expr = { workspace = true } datafusion-functions-nested = { workspace = true } insta = { workspace = true } itertools = { workspace = true } + +[[bench]] +harness = false +name = "string_in_list_pruning" diff --git a/datafusion/pruning/benches/string_in_list_pruning.rs b/datafusion/pruning/benches/string_in_list_pruning.rs new file mode 100644 index 0000000000000..d26112611bf3c --- /dev/null +++ b/datafusion/pruning/benches/string_in_list_pruning.rs @@ -0,0 +1,241 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Compare compact string IN-list pruning with per-value min/max expansion. +//! +//! Both cases raise `max_in_list_size` to the domain size. On a baseline +//! without compact pruning, `in_list` measures the ordinary raised-cap path. +//! The explicit `expanded_or` is a balanced tree of equalities, which produces +//! the same per-value statistics checks without making the baseline depend on +//! a deeply nested expression. Half of the statistics intervals hit a domain +//! member and half fall in a sparse gap. Bloom filters are not involved. +//! +//! Run with `cargo bench -p datafusion-pruning --bench string_in_list_pruning`. +//! The construction benchmarks reuse their input physical expressions; the +//! evaluation benchmarks reuse their already-built pruning predicates. + +use std::collections::HashSet; +use std::hint::black_box; +use std::sync::Arc; + +use arrow::array::{ArrayRef, BooleanArray, StringViewArray, UInt64Array}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use datafusion_common::{Column, ScalarValue}; +use datafusion_expr_common::operator::Operator; +use datafusion_physical_expr::PhysicalExprRef; +use datafusion_physical_expr::expressions::{BinaryExpr, col, in_list, lit}; +use datafusion_pruning::{PruningPredicate, PruningPredicateBuilder, PruningStatistics}; + +const DOMAIN_SIZES: [usize; 4] = [20, 21, 256, 1024]; +const CONTAINERS: usize = 4096; + +fn value(index: usize) -> String { + format!("key{index:08}") +} + +fn balanced_or(expressions: &[PhysicalExprRef]) -> PhysicalExprRef { + if expressions.len() == 1 { + return Arc::clone(&expressions[0]); + } + let middle = expressions.len() / 2; + Arc::new(BinaryExpr::new( + balanced_or(&expressions[..middle]), + Operator::Or, + balanced_or(&expressions[middle..]), + )) +} + +fn build_predicate( + expression: &PhysicalExprRef, + schema: &SchemaRef, + max_in_list_size: usize, +) -> PruningPredicate { + PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(schema)) + .with_max_in_list_size(max_in_list_size) + .try_build(Arc::clone(expression)) + .unwrap() +} + +struct IntervalStatistics { + min: ArrayRef, + max: ArrayRef, + null_counts: ArrayRef, + row_counts: ArrayRef, +} + +impl IntervalStatistics { + fn new(domain_size: usize) -> Self { + let min = StringViewArray::from_iter_values((0..CONTAINERS).map(|index| { + let start = (index / 2 % domain_size) * 10; + value(start + if index % 2 == 0 { 0 } else { 3 }) + })); + let max = StringViewArray::from_iter_values((0..CONTAINERS).map(|index| { + let start = (index / 2 % domain_size) * 10; + value(start + if index % 2 == 0 { 0 } else { 7 }) + })); + Self { + min: Arc::new(min), + max: Arc::new(max), + null_counts: Arc::new(UInt64Array::from(vec![0; CONTAINERS])), + row_counts: Arc::new(UInt64Array::from(vec![128; CONTAINERS])), + } + } +} + +impl PruningStatistics for IntervalStatistics { + fn min_values(&self, column: &Column) -> Option { + (column.name == "value").then(|| Arc::clone(&self.min)) + } + + fn max_values(&self, column: &Column) -> Option { + (column.name == "value").then(|| Arc::clone(&self.max)) + } + + fn num_containers(&self) -> usize { + CONTAINERS + } + + fn null_counts(&self, column: &Column) -> Option { + (column.name == "value").then(|| Arc::clone(&self.null_counts)) + } + + fn row_counts(&self) -> Option { + Some(Arc::clone(&self.row_counts)) + } + + fn contained( + &self, + _column: &Column, + _values: &HashSet, + ) -> Option { + None + } +} + +struct BenchmarkCase { + size: usize, + schema: SchemaRef, + in_list: PhysicalExprRef, + expanded_or: PhysicalExprRef, + in_list_predicate: PruningPredicate, + expanded_or_predicate: PruningPredicate, + statistics: IntervalStatistics, +} + +impl BenchmarkCase { + fn new(size: usize) -> Self { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Utf8View, + false, + )])); + let column = col("value", &schema).unwrap(); + let values = (0..size) + .map(|index| lit(ScalarValue::new_utf8view(value(index * 10)))) + .collect::>(); + let in_list = + in_list(Arc::clone(&column), values.clone(), &false, &schema).unwrap(); + let equalities = values + .into_iter() + .map(|value| { + Arc::new(BinaryExpr::new(Arc::clone(&column), Operator::Eq, value)) + as PhysicalExprRef + }) + .collect::>(); + let expanded_or = balanced_or(&equalities); + let in_list_predicate = build_predicate(&in_list, &schema, size); + let expanded_or_predicate = build_predicate(&expanded_or, &schema, size); + eprintln!( + "string_in_list_pruning: {size} values, compact={}", + in_list_predicate + .predicate_expr() + .to_string() + .contains("IN_SET_INTERSECTS") + ); + let statistics = IntervalStatistics::new(size); + + // Check that both benchmark paths do the same useful work, rather + // than comparing compact pruning with an always-true fallback. + let expected = (0..CONTAINERS) + .map(|index| index % 2 == 0) + .collect::>(); + assert_eq!(in_list_predicate.prune(&statistics).unwrap(), expected); + assert_eq!(expanded_or_predicate.prune(&statistics).unwrap(), expected); + + Self { + size, + schema, + in_list, + expanded_or, + in_list_predicate, + expanded_or_predicate, + statistics, + } + } +} + +fn criterion_benchmark(criterion: &mut Criterion) { + let cases = DOMAIN_SIZES.map(BenchmarkCase::new); + let mut construction = criterion.benchmark_group("string_in_list_pruning/construct"); + for case in &cases { + construction.throughput(Throughput::Elements(case.size as u64)); + for (name, expression) in [ + ("in_list", &case.in_list), + ("expanded_or", &case.expanded_or), + ] { + construction.bench_with_input( + BenchmarkId::new(name, case.size), + expression, + |bencher, expression| { + bencher.iter(|| { + black_box(build_predicate( + black_box(expression), + &case.schema, + case.size, + )) + }); + }, + ); + } + } + construction.finish(); + + let mut evaluation = criterion.benchmark_group("string_in_list_pruning/evaluate"); + evaluation.throughput(Throughput::Elements(CONTAINERS as u64)); + for case in &cases { + for (name, predicate) in [ + ("in_list", &case.in_list_predicate), + ("expanded_or", &case.expanded_or_predicate), + ] { + evaluation.bench_with_input( + BenchmarkId::new(name, case.size), + predicate, + |bencher, predicate| { + bencher.iter(|| { + black_box(predicate.prune(black_box(&case.statistics)).unwrap()) + }); + }, + ); + } + } + evaluation.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/pruning/src/lib.rs b/datafusion/pruning/src/lib.rs index 2b334d2847980..6bf1815900aa8 100644 --- a/datafusion/pruning/src/lib.rs +++ b/datafusion/pruning/src/lib.rs @@ -19,6 +19,7 @@ mod file_pruner; mod pruning_predicate; +mod string_in_list; pub use file_pruner::FilePruner; pub use pruning_predicate::{ diff --git a/datafusion/pruning/src/pruning_predicate.rs b/datafusion/pruning/src/pruning_predicate.rs index dff18173ae32a..4985cadac72b5 100644 --- a/datafusion/pruning/src/pruning_predicate.rs +++ b/datafusion/pruning/src/pruning_predicate.rs @@ -22,6 +22,8 @@ use std::collections::HashSet; use std::sync::Arc; +use crate::string_in_list::StringInListPruningExpr; + use arrow::array::AsArray; use arrow::{ array::{ArrayRef, BooleanArray, new_null_array}, @@ -441,9 +443,24 @@ impl<'a> PruningPredicateBuilder<'a> { self } - /// Cap on the size of `IN (...)` lists that will be rewritten into per- - /// value min/max statistics checks. Lists longer than this fall back to - /// the unhandled-predicate hook (typically "keep the container"). + /// Cap on the input length of `IN (...)` lists eligible for min/max pruning. + /// + /// For a nonempty list of `N` values (before deduplication) and cap `C`: + /// + /// | Condition | Pruning representation | + /// | --- | --- | + /// | `N <= min(20, C)` | Existing per-value rewrite | + /// | `20 < N <= C`, positive, non-null literal strings on a string column | Compact sorted domain | + /// | `20 < N <= C`, other lists | Existing per-value rewrite | + /// | `N > C` | Unhandled-predicate hook, normally "keep the container" | + /// + /// Empty lists also use the unhandled-predicate hook. A cap of zero disables + /// this `IN` min/max rewrite, not other predicates or literal/containment + /// pruning (such as Bloom filters). The default cap is [`MAX_IN_LIST_SIZE`] + /// (20), so the compact path requires an explicitly raised cap. + /// + /// Raising the cap can still build large comparison trees for non-string + /// lists, `NOT IN`, or lists containing NULL; their handling is unchanged. /// /// Query engines typically pass /// `datafusion.execution.parquet.max_in_list_size` here. @@ -1453,10 +1470,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; @@ -1492,14 +1550,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 @@ -1539,8 +1598,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, @@ -1582,6 +1642,16 @@ fn build_predicate_expression( } } if let Some(in_list) = expr.downcast_ref::() { + // Keep the existing expression shape for lists of at most 20 values. + // This lower bound is a scope/compatibility choice, not a measured + // performance threshold; compact pruning is opt-in via a raised cap. + if in_list.list().len() > MAX_IN_LIST_SIZE + && in_list.list().len() <= max_in_list_size + && let Some(pruning_expr) = + build_string_in_list_expr(in_list, schema, required_columns) + { + return pruning_expr; + } if !in_list.list().is_empty() && in_list.list().len() <= max_in_list_size { let eq_op = if in_list.negated() { Operator::NotEq @@ -2182,8 +2252,12 @@ mod tests { use arrow::array::Decimal128Array; use arrow::{ - array::{BinaryArray, Int32Array, Int64Array, StringArray, UInt64Array}, - datatypes::TimeUnit, + array::{ + BinaryArray, DictionaryArray, Int32Array, Int64Array, StringArray, + UInt64Array, + }, + buffer::NullBuffer, + datatypes::{Int32Type, TimeUnit}, }; use datafusion_expr::expr::InList; use datafusion_expr::{BinaryExpr, Expr, cast, is_null, try_cast}; @@ -3574,6 +3648,323 @@ mod tests { Ok(()) } + fn large_string_pruning_predicate( + expr: PhysicalExprRef, + schema: SchemaRef, + ) -> Result { + PruningPredicateBuilder::new() + .with_file_schema(schema) + .with_max_in_list_size(10_000) + .try_build(expr) + } + + #[test] + fn large_string_in_list_prunes_exact_intervals() -> Result<()> { + let types = [ + DataType::Utf8, + DataType::LargeUtf8, + DataType::Utf8View, + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + ]; + for count in [20, 21, 256, 10_000] { + for data_type in &types { + let schema = Arc::new(Schema::new(vec![Field::new( + "c1", + data_type.clone(), + true, + )])); + let values = (0..count) + .map(|i| { + let value = ScalarValue::from(format!("k{:06}", i * 10)) + .cast_to(data_type)?; + Ok(Arc::new(phys_expr::Literal::new(value)) as PhysicalExprRef) + }) + .collect::>>()?; + let expr = phys_expr::in_list( + Arc::new(phys_expr::Column::new("c1", 0)), + values, + &false, + &schema, + )?; + let predicate = + large_string_pruning_predicate(Arc::clone(&expr), schema)?; + let last_value = format!("k{:06}", (count - 1) * 10); + let stats = TestStatistics::new().with( + "c1", + ContainerStats::new_utf8( + [ + Some("k000000"), + Some("k000001"), + Some("k000010"), + Some("j"), + Some("z"), + None, + Some("k000020"), + Some("k000005"), + // Single bounds exclude the domain, but equality + // with either endpoint must remain eligible. + Some("z"), + None, + Some(last_value.as_str()), + None, + ], + [ + Some("k000000"), + Some("k000009"), + Some("k000010"), + Some("j9"), + Some("z9"), + Some("k000010"), + None, + Some("k000015"), + None, + Some("j9"), + None, + Some("k000000"), + ], + ) + .with_null_counts([Some(0); 12]) + .with_row_counts([Some(1); 12]), + ); + assert_eq!( + predicate.prune(&stats)?, + [ + true, false, true, false, false, true, true, true, false, false, + true, true, + ], + "count={count}, type={data_type:?}" + ); + assert!(Arc::ptr_eq(predicate.orig_expr(), &expr)); + assert_eq!(predicate.literal_guarantees().len(), 1); + assert_eq!(predicate.literal_guarantees()[0].literals.len(), count); + assert_eq!( + predicate.required_columns().single_column().unwrap().name(), + "c1" + ); + if count > MAX_IN_LIST_SIZE { + // The expression and statistics schema do not grow with the domain. + assert_eq!(predicate.required_columns.columns.len(), 4); + let mut nodes = 0; + predicate.predicate_expr().apply(|_| { + nodes += 1; + Ok(TreeNodeRecursion::Continue) + })?; + assert_eq!(nodes, 7); + } + } + } + Ok(()) + } + + #[test] + fn large_string_in_list_preserves_dictionary_nulls() -> Result<()> { + let data_type = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)); + let schema = + Arc::new(Schema::new(vec![Field::new("c1", data_type.clone(), true)])); + let values = (0..21) + .map(|i| { + let value = ScalarValue::from(format!("k{i:06}")).cast_to(&data_type)?; + Ok(Arc::new(phys_expr::Literal::new(value)) as PhysicalExprRef) + }) + .collect::>>()?; + let expr = phys_expr::in_list( + Arc::new(phys_expr::Column::new("c1", 0)), + values, + &false, + &schema, + )?; + let predicate = large_string_pruning_predicate(expr, schema)?; + // NULL dictionary values can have nonempty payloads. Neither payload + // may become a bound when a valid key references the NULL value. + let (offsets, values, _) = + StringArray::from(vec!["z", "zz", "", "", "k000010", "x"]).into_parts(); + let dictionary_values: ArrayRef = Arc::new(StringArray::new( + offsets, + values, + Some(NullBuffer::from(vec![false, true, false, true, true, true])), + )); + let bounds = |keys: Vec>| -> Result { + Ok(Arc::new(DictionaryArray::::try_new( + Int32Array::from(keys), + Arc::clone(&dictionary_values), + )?)) + }; + let stats = TestStatistics::new().with( + "c1", + ContainerStats::new() + .with_min(bounds(vec![ + Some(0), // NULL with payload "z" + Some(3), // empty string + Some(4), // matching singleton + Some(5), // disjoint singleton + None, // NULL keys + Some(2), // NULL values in both bounds + Some(1), // inverted known bounds + ])?) + .with_max(bounds(vec![ + Some(1), + Some(2), + Some(4), + Some(5), + None, + Some(2), + Some(3), + ])?) + .with_null_counts([Some(0); 7]) + .with_row_counts([Some(1); 7]), + ); + assert_eq!( + predicate.prune(&stats)?, + [true, true, true, false, true, true, true] + ); + Ok(()) + } + + #[test] + fn large_string_in_list_handles_unicode_and_unknown_bounds() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)])); + let mut values = (0..21).map(|i| lit(format!("m{i:03}"))).collect::>(); + values.extend([ + lit(""), + lit("az"), + lit("aé"), + lit("aé"), + lit("é"), + lit("🦀"), + ]); + let expr = col("c1").in_list(values, false); + let predicate = + large_string_pruning_predicate(logical2physical(&expr, &schema), schema)?; + let stats = TestStatistics::new().with( + "c1", + ContainerStats::new_utf8( + [ + Some(""), + Some("az"), + Some("a{"), + Some("aé"), + Some("ê"), + Some("z"), + None, + None, + ], + [ + Some(""), + Some("az"), + Some("aè"), + Some("aé"), + Some("🦀"), + Some("m"), + Some("z"), + None, + ], + ) + .with_null_counts([ + Some(0), + Some(0), + Some(0), + Some(0), + Some(0), + Some(0), + Some(0), + Some(1), + ]) + .with_row_counts([Some(1); 8]), + ); + assert_eq!( + predicate.prune(&stats)?, + [true, true, false, true, true, true, true, false] + ); + Ok(()) + } + + #[test] + fn large_string_in_list_respects_configured_limit() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)])); + let values = (0..21).map(|i| lit(format!("a{i:03}"))).collect::>(); + let expr = logical2physical(&col("c1").in_list(values, false), &schema); + + for (limit, compact) in [(0, false), (20, false), (21, true), (32, true)] { + let predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .with_max_in_list_size(limit) + .try_build(Arc::clone(&expr))?; + assert_eq!( + predicate + .predicate_expr() + .to_string() + .contains("IN_SET_INTERSECTS"), + compact, + "limit={limit}" + ); + assert_eq!(is_always_true(predicate.predicate_expr()), !compact); + } + + let default = PruningPredicateBuilder::new() + .with_file_schema(schema) + .try_build(expr)?; + assert!(is_always_true(default.predicate_expr())); + Ok(()) + } + + #[test] + fn large_string_in_list_keeps_null_and_not_in_semantics() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)])); + let values = (0..21).map(|i| lit(format!("a{i:03}"))).collect::>(); + let stats = TestStatistics::new().with( + "c1", + ContainerStats::new_utf8( + [Some("middle"), Some("other")], + [Some("middle"), Some("other")], + ) + .with_null_counts([Some(0); 2]) + .with_row_counts([Some(1); 2]), + ); + let positive = col("c1").in_list(values.clone(), false); + let predicate = large_string_pruning_predicate( + logical2physical(&positive.or(col("c1").eq(lit("middle"))), &schema), + Arc::clone(&schema), + )?; + assert_eq!(predicate.prune(&stats)?, [true, false]); + + let mut with_null = values.clone(); + with_null.push(lit(ScalarValue::Utf8(None))); + for expr in [ + col("c1").in_list(values, true), + col("c1").in_list(with_null.clone(), false), + col("c1").in_list(with_null.clone(), true), + ] { + let physical = logical2physical(&expr, &schema); + let default = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(Arc::clone(&physical))?; + assert!(is_always_true(default.predicate_expr()), "{expr}"); + assert_eq!(default.prune(&stats)?, [true, true]); + + // Raising the cap retains the existing per-value rewrite for + // NOT IN and lists containing NULL; neither uses the new path. + let raised = large_string_pruning_predicate(physical, Arc::clone(&schema))?; + assert!( + !raised + .predicate_expr() + .to_string() + .contains("IN_SET_INTERSECTS"), + "{expr}" + ); + } + + // Inverting NOT IN (..., NULL) must not prove a full match and bypass + // the original row filter, which returns UNKNOWN for both rows. + let not_in = logical2physical(&col("c1").in_list(with_null, true), &schema); + let batch = RecordBatch::try_new( + schema, + vec![Arc::new(StringArray::from(vec!["middle", "other"]))], + )?; + assert_eq!(not_in.evaluate(&batch)?.into_array(2)?.null_count(), 2); + Ok(()) + } + #[test] fn row_group_predicate_cast_int_int() -> Result<()> { let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]); diff --git a/datafusion/pruning/src/string_in_list.rs b/datafusion/pruning/src/string_in_list.rs new file mode 100644 index 0000000000000..bc5c8fa908e70 --- /dev/null +++ b/datafusion/pruning/src/string_in_list.rs @@ -0,0 +1,239 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::fmt::{self, Display, Formatter}; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +use arrow::array::{Array, AsArray, BooleanArray}; +use arrow::compute::cast; +use arrow::datatypes::{DataType, Schema}; +use arrow::record_batch::RecordBatch; +use datafusion_common::{Result, assert_eq_or_internal_err}; +use datafusion_physical_expr::{PhysicalExpr, PhysicalExprRef}; +use datafusion_physical_plan::ColumnarValue; + +/// Tests whether a sorted string domain intersects an inclusive statistics interval. +/// +/// [`PhysicalExpr::evaluate`] returns one nullable Boolean per min/max interval: +/// * `true`: the interval intersects the domain, so matching rows may exist. +/// * `false`: the available bounds prove the interval disjoint from the domain. +/// * `NULL`: incomplete, invalid, or unusable bounds prevent a safe decision. +/// +/// A single known bound can still prove disjointness. Otherwise, unknown results +/// keep the container eligible for reading. +/// +/// This expression is used only for pruning; the original IN remains the row filter. +#[derive(Debug, Eq)] +pub(crate) struct StringInListPruningExpr { + min: PhysicalExprRef, + max: PhysicalExprRef, + values: Arc<[String]>, +} + +impl StringInListPruningExpr { + pub(crate) fn new( + min: PhysicalExprRef, + max: PhysicalExprRef, + mut values: Vec, + ) -> Self { + values.sort_unstable(); + values.dedup(); + Self { + min, + max, + values: values.into(), + } + } +} + +impl PartialEq for StringInListPruningExpr { + fn eq(&self, other: &Self) -> bool { + self.min.eq(&other.min) && self.max.eq(&other.max) && self.values == other.values + } +} + +impl Hash for StringInListPruningExpr { + fn hash(&self, state: &mut H) { + self.min.hash(state); + self.max.hash(state); + self.values.hash(state); + } +} + +impl Display for StringInListPruningExpr { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!( + f, + "IN_SET_INTERSECTS({}, {}, {} values)", + self.min, + self.max, + self.values.len() + ) + } +} + +fn has_oversized_string_buffer(array: &dyn Array, limit: usize) -> bool { + match array.data_type() { + DataType::Utf8 => array.as_string::().values().len() >= limit, + DataType::LargeUtf8 => array.as_string::().values().len() >= limit, + DataType::Dictionary(_, _) => has_oversized_string_buffer( + array.as_any_dictionary().values().as_ref(), + limit, + ), + _ => false, + } +} + +impl PhysicalExpr for StringInListPruningExpr { + fn data_type(&self, _input_schema: &Schema) -> Result { + Ok(DataType::Boolean) + } + + fn nullable(&self, _input_schema: &Schema) -> Result { + Ok(true) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + // Normalize Utf8, LargeUtf8, Utf8View, and dictionary-encoded statistics. + let min = self.min.evaluate(batch)?.into_array(batch.num_rows())?; + let max = self.max.evaluate(batch)?.into_array(batch.num_rows())?; + // A short string slice can retain a buffer too large for Utf8View's + // u32 offsets. Avoid a panic in the cast and keep pruning conservative. + if has_oversized_string_buffer(min.as_ref(), u32::MAX as usize) + || has_oversized_string_buffer(max.as_ref(), u32::MAX as usize) + { + return Ok(ColumnarValue::Array(Arc::new(BooleanArray::new_null( + batch.num_rows(), + )))); + } + // Dictionary values can be NULL behind valid keys. Preserve their + // validity even if the view cast only carries the key nulls. + // TODO: Revisit this workaround once the Arrow dependency includes + // https://github.com/apache/arrow-rs/pull/10510. + let min_nulls = min.logical_nulls(); + let max_nulls = max.logical_nulls(); + let min = cast(&min, &DataType::Utf8View)?; + let max = cast(&max, &DataType::Utf8View)?; + let min = min.as_string_view(); + let max = max.as_string_view(); + let matches: BooleanArray = (0..batch.num_rows()) + .map(|i| { + let min = (min.is_valid(i) + && min_nulls.as_ref().is_none_or(|nulls| nulls.is_valid(i))) + .then(|| min.value(i).as_bytes()); + let max = (max.is_valid(i) + && max_nulls.as_ref().is_none_or(|nulls| nulls.is_valid(i))) + .then(|| max.value(i).as_bytes()); + match (min, max) { + (Some(min), Some(max)) => { + if min > max { + return None; + } + // Rust string ordering and these byte comparisons both use + // unsigned lexicographic UTF-8 order, as required by the + // PruningStatistics min/max contract. Parquet adapters mask + // bounds with unusable ordering; PartitionPruningStatistics + // uses actual Arrow partition values. PrunableStatistics + // trusts file providers' bounds: there is no ordering gate + // for arbitrary statistics providers here. + let index = self.values.partition_point(|v| v.as_bytes() < min); + Some(self.values.get(index).is_some_and(|v| v.as_bytes() <= max)) + } + // A missing bound makes that end of the interval unbounded. + // Exclude only when the whole domain lies beyond the known bound; + // gaps within the domain and equality cannot prove disjointness. + (Some(min), None) + if self.values.last().is_some_and(|v| v.as_bytes() < min) => + { + Some(false) + } + (None, Some(max)) + if self.values.first().is_some_and(|v| v.as_bytes() > max) => + { + Some(false) + } + _ => None, + } + }) + .collect(); + Ok(ColumnarValue::Array(Arc::new(matches))) + } + + fn children(&self) -> Vec<&PhysicalExprRef> { + vec![&self.min, &self.max] + } + + fn with_new_children( + self: Arc, + children: Vec, + ) -> Result { + assert_eq_or_internal_err!(children.len(), 2); + Ok(Arc::new(Self { + min: Arc::clone(&children[0]), + max: Arc::clone(&children[1]), + values: Arc::clone(&self.values), + })) + } + + fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "{self}") + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ArrayRef, StringArray}; + + #[test] + fn oversized_buffers_check_retained_data_not_visible_offsets() -> Result<()> { + // Exercise the size boundary without allocating a 4 GiB buffer. + let limit = 32; + let padding = "p".repeat(limit - 1); + let array: ArrayRef = Arc::new(StringArray::from(vec!["a", padding.as_str()])); + + for value_type in [DataType::Utf8, DataType::LargeUtf8] { + for data_type in [ + value_type.clone(), + DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(value_type.clone()), + ), + DataType::Dictionary( + Box::new(DataType::UInt64), + Box::new(value_type.clone()), + ), + ] { + let slice = cast(&array, &data_type)?.slice(0, 1); + assert!(has_oversized_string_buffer(slice.as_ref(), limit - 1)); + assert!(has_oversized_string_buffer(slice.as_ref(), limit)); + assert!(!has_oversized_string_buffer(slice.as_ref(), limit + 1)); + } + } + + // Already-normalized views do not have the byte-array cast limitation. + for data_type in [ + DataType::Utf8View, + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8View)), + ] { + let slice = cast(&array, &data_type)?.slice(0, 1); + assert!(!has_oversized_string_buffer(slice.as_ref(), limit)); + } + Ok(()) + } +} diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 261027cb92e25..48883070c3798 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -413,7 +413,7 @@ datafusion.execution.parquet.dictionary_page_size_limit 1048576 (writing) Sets b datafusion.execution.parquet.enable_page_index true (reading) If true, reads the Parquet data page level metadata (the Page Index), if present, to reduce the I/O and number of rows decoded. datafusion.execution.parquet.encoding NULL (writing) Sets default encoding for any column. Valid values are: plain, plain_dictionary, rle, bit_packed, delta_binary_packed, delta_length_byte_array, delta_byte_array, rle_dictionary, and byte_stream_split. These values are not case sensitive. If NULL, uses default parquet writer setting datafusion.execution.parquet.force_filter_selections false (reading) Force the use of RowSelections for filter results, when pushdown_filters is enabled. If false, the reader will automatically choose between a RowSelection and a Bitmap based on the number and pattern of selected rows. -datafusion.execution.parquet.max_in_list_size 20 Maximum number of values in an `IN (...)` list for which pruning will occur. Longer lists will not be used to prune files, row groups, or data pages. Higher values help in cases such as filtering on a list of ~25-100 identifiers, but also make the predicate more expensive to evaluate. Set to 0 to disable `IN (...)` list pruning entirely. Defaults to 20. +datafusion.execution.parquet.max_in_list_size 20 Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger positive, non-null literal string lists on a string column use a compact sorted domain. Other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. datafusion.execution.parquet.max_predicate_cache_size NULL (reading) The maximum predicate cache size, in bytes. When `pushdown_filters` is enabled, sets the maximum memory used to cache the results of predicate evaluation between filter evaluation and output generation. Decreasing this value will reduce memory usage, but may increase IO and CPU usage. None means use the default parquet reader setting. 0 means no caching. datafusion.execution.parquet.max_row_group_bytes NULL (writing) Target maximum size of each row group in bytes. When set, the writer flushes whenever either this limit or `max_row_group_size` is reached, whichever comes first. Useful for bounding writer memory on wide schemas where a row-count limit can map to very different byte sizes. Matches the behavior of `parquet.block.size` in parquet-mr. If `None` (the default), only the row-count limit applies. Currently only honored when `allow_single_file_parallelism` is `false`; by default the parallel file writer ignores this limit. datafusion.execution.parquet.max_row_group_size 1048576 (writing) Target maximum number of rows in each row group (defaults to 1M rows). Writing larger row groups requires more memory to write, but can get better compression and be faster to read. When `max_row_group_bytes` is also set, the writer flushes a row group when either limit is reached, whichever comes first. diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index dec013a88d38b..7411cdee6fda4 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -93,7 +93,7 @@ The following configuration settings are available: | datafusion.execution.parquet.coerce_int96_tz | NULL | (reading) Optional timezone applied to INT96 columns when `coerce_int96` is set. When `Some`, INT96 columns coerce to `Timestamp(, Some())` instead of the default `Timestamp(, None)`. Spark and other systems write INT96 values as UTC-adjusted instants, so callers that need the resulting Arrow type to be timezone-aware (e.g. for Spark `TimestampType` semantics) should set this to `"UTC"`. No effect when `coerce_int96` is `None`. | | datafusion.execution.parquet.bloom_filter_on_read | true | (reading) Use any available bloom filters when reading parquet files | | datafusion.execution.parquet.max_predicate_cache_size | NULL | (reading) The maximum predicate cache size, in bytes. When `pushdown_filters` is enabled, sets the maximum memory used to cache the results of predicate evaluation between filter evaluation and output generation. Decreasing this value will reduce memory usage, but may increase IO and CPU usage. None means use the default parquet reader setting. 0 means no caching. | -| datafusion.execution.parquet.max_in_list_size | 20 | Maximum number of values in an `IN (...)` list for which pruning will occur. Longer lists will not be used to prune files, row groups, or data pages. Higher values help in cases such as filtering on a list of ~25-100 identifiers, but also make the predicate more expensive to evaluate. Set to 0 to disable `IN (...)` list pruning entirely. Defaults to 20. | +| datafusion.execution.parquet.max_in_list_size | 20 | Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger positive, non-null literal string lists on a string column use a compact sorted domain. Other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. | | datafusion.execution.parquet.data_pagesize_limit | 1048576 | (writing) Sets best effort maximum size of data page in bytes | | datafusion.execution.parquet.write_batch_size | 1024 | (writing) Sets write_batch_size in rows | | datafusion.execution.parquet.writer_version | 1.0 | (writing) Sets parquet writer version valid values are "1.0" and "2.0" |