diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 347721c43d7cf..042d3c330c231 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1392,12 +1392,14 @@ config_namespace! { /// rewrite; other predicates and Bloom-filter pruning remain available. /// /// Within the cap, nonempty lists of at most 20 values use the existing - /// per-value rewrite. Larger literal string lists on a string column use - /// a compact representation, for both `IN` and `NOT IN`, including lists - /// with NULL members. `NOT IN` with NULL and all-NULL `IN` lists cannot - /// match any rows. Other lists retain the existing per-value rewrite, so - /// raising the cap can make those predicates expensive to build and - /// evaluate. + /// per-value rewrite. Larger literal lists use a compact representation + /// when the column type is string, variable-length binary, integer, + /// decimal, date, time, timestamp, or duration. This applies to both `IN` + /// and `NOT IN`, including lists with NULL members. `NOT IN` with NULL and + /// all-NULL `IN` lists cannot match any rows. Compact lists containing NULL + /// do not use the fully-matched-row-group optimization. Floating-point and + /// other lists retain the existing per-value rewrite, so raising the cap + /// can make those predicates expensive to build and evaluate. /// /// Defaults to 20. pub max_in_list_size: usize, default = 20 diff --git a/datafusion/core/tests/parquet/string_in_list_pruning.rs b/datafusion/core/tests/parquet/string_in_list_pruning.rs index 52d2c58bdd943..f2646855713ce 100644 --- a/datafusion/core/tests/parquet/string_in_list_pruning.rs +++ b/datafusion/core/tests/parquet/string_in_list_pruning.rs @@ -15,12 +15,16 @@ // specific language governing permissions and limitations // under the License. -//! End-to-end coverage for compact, large string IN-list pruning. The `IN` and -//! `NOT IN` cases disable the row and Bloom filters to isolate min/max pruning. +//! End-to-end coverage for compact, large IN-list pruning. The string cases +//! exercise detailed edge conditions; representative ordered types verify the +//! same path against real Parquet row-group and page statistics. use std::sync::Arc; -use arrow::array::StringArray; +use arrow::array::{ + ArrayRef, BinaryArray, Date64Array, Decimal128Array, Int64Array, StringArray, + TimestampMicrosecondArray, +}; use arrow::datatypes::{DataType, Field, Schema}; use arrow::record_batch::RecordBatch; use arrow::util::pretty::pretty_format_batches; @@ -34,9 +38,11 @@ use datafusion_common::config::TableParquetOptions; use datafusion_common::{ScalarValue, assert_batches_eq}; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; use datafusion_physical_expr::expressions::{col, in_list, lit}; +use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::metrics::{MetricValue, MetricsSet}; use object_store::path::Path; use parquet::arrow::ArrowWriter; +use parquet::file::metadata::ParquetMetaData; use parquet::file::properties::{EnabledStatistics, WriterProperties}; use tempfile::NamedTempFile; @@ -47,6 +53,21 @@ const UNITS: usize = 4; const TOTAL_ROWS: usize = ROWS_PER_UNIT * UNITS; const MATCHING_ROWS: usize = ROWS_PER_UNIT * 2; +fn assert_file_layout( + metadata: &ParquetMetaData, + total_rows: usize, + rows_per_group: usize, +) { + assert_eq!(metadata.num_row_groups(), total_rows / rows_per_group); + let offsets = metadata.offset_index().unwrap(); + for row_group in offsets { + assert_eq!( + row_group[0].page_locations().len(), + rows_per_group / ROWS_PER_UNIT + ); + } +} + /// Write either four row groups or four pages in one row group. Each unit holds /// a single repeated value, so `NOT IN` can exclude the two units whose value is /// a list member. The second unit lies in a gap between two members of every @@ -116,14 +137,7 @@ fn write_file_with_truncation( let mut writer = ArrowWriter::try_new(&mut file, schema, Some(properties)).unwrap(); writer.write(&batch).unwrap(); let metadata = writer.close().unwrap(); - assert_eq!(metadata.num_row_groups(), total_rows / rows_per_group); - let offsets = metadata.offset_index().unwrap(); - for row_group in offsets { - assert_eq!( - row_group[0].page_locations().len(), - rows_per_group / ROWS_PER_UNIT - ); - } + assert_file_layout(&metadata, total_rows, rows_per_group); file } @@ -257,6 +271,93 @@ async fn scan( } } +fn repeated_ordered_values(scale: i64) -> Vec { + [0, 1, 2, 100] + .into_iter() + .flat_map(|value| std::iter::repeat_n(value * scale, ROWS_PER_UNIT)) + .collect() +} + +fn write_ordered_file( + schema: &Arc, + values: ArrayRef, + page_pruning: bool, +) -> NamedTempFile { + let mut file = tempfile::Builder::new() + .prefix("ordered_in_list_pruning") + .suffix(".parquet") + .tempfile() + .unwrap(); + let rows_per_group = if page_pruning { + TOTAL_ROWS + } else { + ROWS_PER_UNIT + }; + let properties = WriterProperties::builder() + .set_max_row_group_row_count(Some(rows_per_group)) + .set_data_page_row_count_limit(ROWS_PER_UNIT) + .set_write_batch_size(ROWS_PER_UNIT) + .set_dictionary_enabled(false) + .set_bloom_filter_enabled(false) + .set_statistics_enabled(EnabledStatistics::Page) + .build(); + let batch = RecordBatch::try_new(Arc::clone(schema), vec![values]).unwrap(); + let mut writer = + ArrowWriter::try_new(&mut file, Arc::clone(schema), Some(properties)).unwrap(); + writer.write(&batch).unwrap(); + let metadata = writer.close().unwrap(); + assert_file_layout(&metadata, TOTAL_ROWS, rows_per_group); + file +} + +async fn scan_ordered( + file: &NamedTempFile, + schema: &Arc, + list: &[ScalarValue], + max_in_list_size: usize, + page_pruning: bool, + negated: bool, +) -> ScanOutput { + let predicate = in_list( + col("value", schema).unwrap(), + list.iter().cloned().map(lit).collect(), + &negated, + schema, + ) + .unwrap(); + let mut options = TableParquetOptions::default(); + options.global.max_in_list_size = max_in_list_size; + let source = Arc::new( + ParquetSource::new(Arc::clone(schema)) + .with_table_parquet_options(options) + .with_predicate(Arc::clone(&predicate)) + .with_pushdown_filters(false) + .with_enable_page_index(page_pruning) + .with_bloom_filter_on_read(false), + ); + let location = Path::from_filesystem_path(file.path()).unwrap(); + let partitioned_file = PartitionedFile::new( + location.to_string(), + file.as_file().metadata().unwrap().len(), + ); + let config = FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + .with_file(partitioned_file) + .build(); + let scan: Arc = Arc::new(DataSourceExec::new(Arc::new(config))); + let plan: Arc = + Arc::new(FilterExec::try_new(predicate, scan).unwrap()); + let plan_text = displayable(plan.as_ref()).indent(true).to_string(); + let ctx = + SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); + let batches = collect(Arc::clone(&plan), ctx.task_ctx()).await.unwrap(); + let metrics = MetricsFinder::find_metrics(plan.as_ref()).unwrap(); + ScanOutput { + batches, + plan: plan_text, + metrics, + } +} + async fn check_string_in_list_pruning(page_pruning: bool) { let file = make_file(page_pruning); for list_size in [20, 21, 256, 1024] { @@ -413,6 +514,118 @@ async fn check_string_not_in_list_with_truncated_bounds(page_pruning: bool) { assert_eq!(output.counter("output_rows"), ROWS_PER_UNIT); } +#[tokio::test] +async fn ordered_in_list_parquet_pruning() { + let decimal_values = Decimal128Array::from_iter_values( + repeated_ordered_values(1).into_iter().map(i128::from), + ) + .with_precision_and_scale(18, 2) + .unwrap(); + let cases = [ + ( + "binary", + Arc::new(BinaryArray::from_iter_values( + repeated_ordered_values(1) + .into_iter() + .map(|value| [0xff, value as u8]), + )) as ArrayRef, + (0..21) + .map(|value| ScalarValue::Binary(Some(vec![0xff, value * 2]))) + .collect::>(), + ), + ( + "int64", + Arc::new(Int64Array::from(repeated_ordered_values(1))) as ArrayRef, + (0..21) + .map(|value| ScalarValue::Int64(Some(value * 2))) + .collect(), + ), + ( + "decimal128", + Arc::new(decimal_values) as ArrayRef, + (0..21) + .map(|value| ScalarValue::Decimal128(Some(value * 2), 18, 2)) + .collect(), + ), + ( + "date64", + Arc::new(Date64Array::from(repeated_ordered_values(86_400_000))) as ArrayRef, + (0..21) + .map(|value| ScalarValue::Date64(Some(value * 2 * 86_400_000))) + .collect(), + ), + ( + "timestamp_microsecond_utc", + Arc::new( + TimestampMicrosecondArray::from(repeated_ordered_values(1_000_000)) + .with_timezone("UTC"), + ) as ArrayRef, + (0..21) + .map(|value| { + ScalarValue::TimestampMicrosecond( + Some(value * 2 * 1_000_000), + Some("UTC".into()), + ) + }) + .collect(), + ), + ]; + + for (name, values, list) in cases { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + values.data_type().clone(), + false, + )])); + for page_pruning in [false, true] { + let file = write_ordered_file(&schema, Arc::clone(&values), page_pruning); + for negated in [false, true] { + let control = + scan_ordered(&file, &schema, &list, 0, page_pruning, negated).await; + let output = + scan_ordered(&file, &schema, &list, 32, page_pruning, negated).await; + control.assert_no_filter_interference(); + output.assert_no_filter_interference(); + assert!(!control.plan.contains("IN_SET_INTERSECTS")); + assert!(!control.plan.contains("NOT_IN_SET_MAY_MATCH")); + assert_eq!( + pretty_format_batches(&output.batches).unwrap().to_string(), + pretty_format_batches(&control.batches).unwrap().to_string(), + "type={name}, page_pruning={page_pruning}, negated={negated}" + ); + assert_eq!( + output + .batches + .iter() + .map(RecordBatch::num_rows) + .sum::(), + MATCHING_ROWS, + "type={name}, page_pruning={page_pruning}, negated={negated}" + ); + assert!( + output.plan.contains(if negated { + "NOT_IN_SET_MAY_MATCH" + } else { + "IN_SET_INTERSECTS" + }), + "type={name}, page_pruning={page_pruning}, negated={negated}, plan={}", + output.plan + ); + assert_eq!( + output.pruned("row_groups_pruned_statistics"), + if page_pruning { 0 } else { 2 }, + "type={name}, page_pruning={page_pruning}, negated={negated}" + ); + assert_eq!( + output.pruned("page_index_rows_pruned"), + if page_pruning { MATCHING_ROWS } else { 0 }, + "type={name}, page_pruning={page_pruning}, negated={negated}" + ); + } + } + } +} + #[tokio::test] async fn string_in_list_row_group_pruning() { check_string_in_list_pruning(false).await; @@ -439,33 +652,23 @@ async fn string_not_in_list_with_truncated_bounds() { check_string_not_in_list_with_truncated_bounds(true).await; } -#[tokio::test] -async fn string_in_list_with_null_preserves_filter_semantics() { +async fn check_in_list_with_null_preserves_filter_semantics( + name: &str, + values: ArrayRef, + list: Vec, + expected: ScalarValue, +) { let mut file = tempfile::Builder::new() - .prefix("string_in_list_null_pruning") + .prefix("in_list_null_pruning") .suffix(".parquet") .tempfile() .unwrap(); - let schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Utf8, true)])); - // The first row group has a known zero null count, and every value lies - // in a gap in the IN list. The matching value in the second row group is - // deliberately not first, so incorrectly bypassing the row filter changes - // the result when the scan has a limit. - let values = vec![ - Some("v000001"), - Some("v000001"), - Some("v000001"), - Some("v000001"), - Some("v000001"), - Some("v000000"), - None, - Some("v999999"), - ]; - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(StringArray::from(values))], - ) - .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + values.data_type().clone(), + true, + )])); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![values]).unwrap(); let properties = WriterProperties::builder() .set_max_row_group_row_count(Some(4)) .set_bloom_filter_enabled(false) @@ -477,10 +680,7 @@ async fn string_in_list_with_null_preserves_filter_semantics() { // Build the physical source directly so a logical optimizer cannot fold // NOT IN (..., NULL) to an empty relation before the scan. - let mut list = (0..21) - .map(|index| lit(format!("v{:06}", index * 10))) - .collect::>(); - list.push(lit(ScalarValue::Utf8(None))); + let list = list.into_iter().map(lit).collect::>(); let location = Path::from_filesystem_path(file.path()).unwrap(); let partitioned_file = PartitionedFile::new( location.to_string(), @@ -516,7 +716,12 @@ async fn string_in_list_with_null_preserves_filter_semantics() { let plan: Arc = Arc::new(DataSourceExec::new(Arc::new(config))); let plan_text = displayable(plan.as_ref()).indent(true).to_string(); - assert!(plan_text.contains("IN"), "{plan_text}"); + assert!(plan_text.contains("IN"), "type={name}, plan={plan_text}"); + assert_eq!( + plan_text.contains("IN_SET_INTERSECTS"), + !negated && max_in_list_size == 32, + "type={name}, negated={negated}, cap={max_in_list_size}, plan={plan_text}" + ); let batches = collect(Arc::clone(&plan), ctx.task_ctx()).await.unwrap(); let output = ScanOutput { batches, @@ -525,17 +730,29 @@ async fn string_in_list_with_null_preserves_filter_semantics() { }; if negated { - assert!(output.batches.iter().all(|batch| batch.num_rows() == 0)); + assert!( + output.batches.iter().all(|batch| batch.num_rows() == 0), + "type={name}, cap={max_in_list_size}" + ); } else { - assert_batches_eq!( - [ - "+---------+", - "| value |", - "+---------+", - "| v000000 |", - "+---------+", - ], - &output.batches + assert_eq!( + output + .batches + .iter() + .map(RecordBatch::num_rows) + .sum::(), + 1, + "type={name}, cap={max_in_list_size}" + ); + let batch = output + .batches + .iter() + .find(|batch| batch.num_rows() > 0) + .unwrap(); + assert_eq!( + ScalarValue::try_from_array(batch.column(0), 0).unwrap(), + expected, + "type={name}, cap={max_in_list_size}" ); } assert_eq!( @@ -545,7 +762,8 @@ async fn string_in_list_with_null_preserves_filter_semantics() { (false, _) => 3, (true, 0) => 8, (true, _) => 0, - } + }, + "type={name}, negated={negated}, cap={max_in_list_size}" ); assert_eq!(output.fully_matched("row_groups_pruned_statistics"), 0); assert_eq!( @@ -557,7 +775,7 @@ async fn string_in_list_with_null_preserves_filter_semantics() { } else { 1 }, - "negated={negated}, cap={max_in_list_size}, metrics={}", + "type={name}, negated={negated}, cap={max_in_list_size}, metrics={}", output.metrics ); assert_eq!(output.pruned("limit_pruned_row_groups"), 0); @@ -565,3 +783,52 @@ async fn string_in_list_with_null_preserves_filter_semantics() { } } } + +#[tokio::test] +async fn in_list_with_null_preserves_filter_semantics() { + // The first row group has a known zero null count, and every value lies + // in a gap in the IN list. The matching value in the second row group is + // deliberately not first, so incorrectly bypassing the row filter changes + // the result when the scan has a limit. + let mut string_list = (0..21) + .map(|index| ScalarValue::Utf8(Some(format!("v{:06}", index * 10)))) + .collect::>(); + string_list.push(ScalarValue::Utf8(None)); + check_in_list_with_null_preserves_filter_semantics( + "string", + Arc::new(StringArray::from(vec![ + Some("v000001"), + Some("v000001"), + Some("v000001"), + Some("v000001"), + Some("v000001"), + Some("v000000"), + None, + Some("v999999"), + ])) as ArrayRef, + string_list, + ScalarValue::Utf8(Some("v000000".into())), + ) + .await; + + let mut int64_list = (0..21) + .map(|index| ScalarValue::Int64(Some(index * 10))) + .collect::>(); + int64_list.push(ScalarValue::Int64(None)); + check_in_list_with_null_preserves_filter_semantics( + "int64", + Arc::new(Int64Array::from(vec![ + Some(1), + Some(1), + Some(1), + Some(1), + Some(1), + Some(0), + None, + Some(999_999), + ])) as ArrayRef, + int64_list, + ScalarValue::Int64(Some(0)), + ) + .await; +} diff --git a/datafusion/pruning/Cargo.toml b/datafusion/pruning/Cargo.toml index a703f68222f38..a914a0a079bcd 100644 --- a/datafusion/pruning/Cargo.toml +++ b/datafusion/pruning/Cargo.toml @@ -35,3 +35,7 @@ itertools = { workspace = true } [[bench]] harness = false name = "string_in_list_pruning" + +[[bench]] +harness = false +name = "primitive_in_list_pruning" diff --git a/datafusion/pruning/benches/primitive_in_list_pruning.rs b/datafusion/pruning/benches/primitive_in_list_pruning.rs new file mode 100644 index 0000000000000..830cf4d0b6875 --- /dev/null +++ b/datafusion/pruning/benches/primitive_in_list_pruning.rs @@ -0,0 +1,285 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Compare primitive IN-list pruning with per-value min/max expansion. +//! +//! The domain and container matrices match `string_in_list_pruning`. The +//! compact form uses a typed contiguous domain, while `expanded_or` and +//! `expanded_and` measure balanced per-value comparison trees. +//! +//! Run with `cargo bench -p datafusion-pruning --bench primitive_in_list_pruning`. + +use std::collections::HashSet; +use std::hint::black_box; +use std::sync::Arc; + +use arrow::array::{ArrayRef, BooleanArray, Int64Array, UInt64Array}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use datafusion_common::{Column, ScalarValue}; +use datafusion_expr_common::operator::Operator; +use datafusion_physical_expr::PhysicalExprRef; +use datafusion_physical_expr::expressions::{ + BinaryExpr, col, in_list as make_in_list, lit, +}; +use datafusion_pruning::{PruningPredicate, PruningPredicateBuilder, PruningStatistics}; + +const DOMAIN_SIZES: [usize; 9] = [1, 2, 4, 8, 16, 20, 21, 256, 1024]; +const CONTAINER_COUNTS: [usize; 3] = [16, 256, 4096]; + +fn sampled_domain_index( + pair_index: usize, + pair_count: usize, + domain_size: usize, +) -> usize { + let sampled_positions = pair_count.min(domain_size); + let position = pair_index % sampled_positions; + if sampled_positions == 1 { + 0 + } else { + position * (domain_size - 1) / (sampled_positions - 1) + } +} + +fn balanced(expressions: &[PhysicalExprRef], op: Operator) -> PhysicalExprRef { + if expressions.len() == 1 { + return Arc::clone(&expressions[0]); + } + let middle = expressions.len() / 2; + Arc::new(BinaryExpr::new( + balanced(&expressions[..middle], op), + op, + balanced(&expressions[middle..], op), + )) +} + +fn expanded( + column: &PhysicalExprRef, + values: &[PhysicalExprRef], + op: Operator, + combine: Operator, +) -> PhysicalExprRef { + let comparisons = values + .iter() + .map(|value| { + Arc::new(BinaryExpr::new(Arc::clone(column), op, Arc::clone(value))) + as PhysicalExprRef + }) + .collect::>(); + balanced(&comparisons, combine) +} + +fn build_predicate( + expression: &PhysicalExprRef, + schema: &SchemaRef, + max_in_list_size: usize, +) -> PruningPredicate { + PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(schema)) + .with_max_in_list_size(max_in_list_size) + .try_build(Arc::clone(expression)) + .unwrap() +} + +struct IntervalStatistics { + min: ArrayRef, + max: ArrayRef, + null_counts: ArrayRef, + row_counts: ArrayRef, +} + +impl IntervalStatistics { + fn new(domain_size: usize, container_count: usize) -> Self { + let pair_count = container_count.div_ceil(2); + let min = Int64Array::from_iter_values((0..container_count).map(|index| { + let start = + sampled_domain_index(index / 2, pair_count, domain_size) as i64 * 10; + start + if index % 2 == 0 { 0 } else { 3 } + })); + let max = Int64Array::from_iter_values((0..container_count).map(|index| { + let start = + sampled_domain_index(index / 2, pair_count, domain_size) as i64 * 10; + start + if index % 2 == 0 { 0 } else { 7 } + })); + Self { + min: Arc::new(min), + max: Arc::new(max), + null_counts: Arc::new(UInt64Array::from(vec![0; container_count])), + row_counts: Arc::new(UInt64Array::from(vec![128; container_count])), + } + } +} + +impl PruningStatistics for IntervalStatistics { + fn min_values(&self, column: &Column) -> Option { + (column.name == "value").then(|| Arc::clone(&self.min)) + } + + fn max_values(&self, column: &Column) -> Option { + (column.name == "value").then(|| Arc::clone(&self.max)) + } + + fn num_containers(&self) -> usize { + self.min.len() + } + + fn null_counts(&self, column: &Column) -> Option { + (column.name == "value").then(|| Arc::clone(&self.null_counts)) + } + + fn row_counts(&self) -> Option { + Some(Arc::clone(&self.row_counts)) + } + + fn contained( + &self, + _column: &Column, + _values: &HashSet, + ) -> Option { + None + } +} + +struct BenchmarkCase { + size: usize, + schema: SchemaRef, + in_list: PhysicalExprRef, + expanded_or: PhysicalExprRef, + not_in_list: PhysicalExprRef, + expanded_and: PhysicalExprRef, + in_list_predicate: PruningPredicate, + expanded_or_predicate: PruningPredicate, + not_in_list_predicate: PruningPredicate, + expanded_and_predicate: PruningPredicate, +} + +impl BenchmarkCase { + fn new(size: usize) -> Self { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int64, + false, + )])); + let column = col("value", &schema).unwrap(); + let values = (0..size) + .map(|index| lit(ScalarValue::Int64(Some(index as i64 * 10)))) + .collect::>(); + let in_list = + make_in_list(Arc::clone(&column), values.clone(), &false, &schema).unwrap(); + let not_in_list = + make_in_list(Arc::clone(&column), values.clone(), &true, &schema).unwrap(); + let expanded_or = expanded(&column, &values, Operator::Eq, Operator::Or); + let expanded_and = expanded(&column, &values, Operator::NotEq, Operator::And); + let in_list_predicate = build_predicate(&in_list, &schema, size); + let expanded_or_predicate = build_predicate(&expanded_or, &schema, size); + let not_in_list_predicate = build_predicate(¬_in_list, &schema, size); + let expanded_and_predicate = build_predicate(&expanded_and, &schema, size); + Self { + size, + schema, + in_list, + expanded_or, + not_in_list, + expanded_and, + in_list_predicate, + expanded_or_predicate, + not_in_list_predicate, + expanded_and_predicate, + } + } +} + +fn assert_equivalent_results(case: &BenchmarkCase, statistics: &IntervalStatistics) { + let expected = (0..statistics.num_containers()) + .map(|index| index % 2 == 0) + .collect::>(); + assert_eq!(case.in_list_predicate.prune(statistics).unwrap(), expected); + assert_eq!( + case.expanded_or_predicate.prune(statistics).unwrap(), + expected + ); + let negated = expected.into_iter().map(|value| !value).collect::>(); + assert_eq!( + case.not_in_list_predicate.prune(statistics).unwrap(), + negated + ); + assert_eq!( + case.expanded_and_predicate.prune(statistics).unwrap(), + negated + ); +} + +fn criterion_benchmark(criterion: &mut Criterion) { + let cases = DOMAIN_SIZES.map(BenchmarkCase::new); + let mut construction = + criterion.benchmark_group("primitive_in_list_pruning/construct"); + for case in &cases { + for (name, expression) in [ + ("in_list", &case.in_list), + ("expanded_or", &case.expanded_or), + ("not_in_list", &case.not_in_list), + ("expanded_and", &case.expanded_and), + ] { + construction.throughput(Throughput::Elements(case.size as u64)); + construction.bench_with_input( + BenchmarkId::new(name, case.size), + expression, + |bencher, expression| { + bencher.iter(|| { + black_box(build_predicate( + black_box(expression), + &case.schema, + case.size, + )) + }); + }, + ); + } + } + construction.finish(); + + for container_count in CONTAINER_COUNTS { + let mut evaluation = criterion.benchmark_group(format!( + "primitive_in_list_pruning/evaluate/{container_count}_containers" + )); + evaluation.throughput(Throughput::Elements(container_count as u64)); + for case in &cases { + let statistics = IntervalStatistics::new(case.size, container_count); + assert_equivalent_results(case, &statistics); + for (name, predicate) in [ + ("in_list", &case.in_list_predicate), + ("expanded_or", &case.expanded_or_predicate), + ("not_in_list", &case.not_in_list_predicate), + ("expanded_and", &case.expanded_and_predicate), + ] { + evaluation.bench_with_input( + BenchmarkId::new(name, case.size), + predicate, + |bencher, predicate| { + bencher.iter(|| { + black_box(predicate.prune(black_box(&statistics)).unwrap()) + }); + }, + ); + } + } + evaluation.finish(); + } +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/pruning/src/in_list.rs b/datafusion/pruning/src/in_list.rs new file mode 100644 index 0000000000000..7f4b6b5aaf93b --- /dev/null +++ b/datafusion/pruning/src/in_list.rs @@ -0,0 +1,85 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::cmp::Ordering; + +use datafusion_common::ScalarValue; + +/// Which `IN` form a sorted domain is pruning for. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub(crate) enum SetMembership { + /// `col IN (...)`. A row matches only where the domain intersects the + /// interval, so a disjoint interval excludes every row. + In, + /// `col NOT IN (...)`. Overlap proves nothing here: values outside the + /// domain still satisfy the predicate. An interval excludes every row only + /// when it holds a single value that the domain contains. + NotIn, +} + +impl SetMembership { + pub(crate) fn display_name(self) -> &'static str { + match self { + Self::In => "IN_SET_INTERSECTS", + Self::NotIn => "NOT_IN_SET_MAY_MATCH", + } + } + + pub(crate) fn compare_bytes(self, left: &[u8], right: &[u8]) -> Ordering { + match self { + // IN uses this order for interval searches. + Self::In => left.cmp(right), + // NOT IN only needs exact membership. Reject impossible lengths + // before comparing bytes that may have a long common prefix. + Self::NotIn => left.len().cmp(&right.len()).then_with(|| left.cmp(right)), + } + } +} + +/// Evaluates `NOT IN` against an inclusive statistics interval. +/// +/// A container can be excluded only when both bounds identify one value in the +/// domain. A known bound outside the domain proves the container may match and +/// lets an enclosing Boolean expression short-circuit. +/// +/// This remains safe when Parquet truncates byte-array statistics. A truncated +/// minimum is no greater than the true minimum, and a truncated maximum is no +/// less than the true maximum. Therefore, equal stored bounds prove that the +/// true minimum and maximum are also equal. +#[inline(always)] +pub(crate) fn not_in_may_match( + min: Option<&T>, + max: Option<&T>, + contains: impl Fn(&T) -> bool, +) -> Option { + match (min, max) { + (Some(min), Some(max)) if contains(min) => Some(min != max), + (Some(_), Some(_)) => Some(true), + (Some(bound), None) | (None, Some(bound)) if !contains(bound) => Some(true), + _ => None, + } +} + +/// Removes scalar wrappers that do not change the represented value. +pub(crate) fn unwrap_scalar(value: &ScalarValue) -> &ScalarValue { + match value { + ScalarValue::Dictionary(_, value) | ScalarValue::RunEndEncoded(_, _, value) => { + unwrap_scalar(value) + } + value => value, + } +} diff --git a/datafusion/pruning/src/lib.rs b/datafusion/pruning/src/lib.rs index 6bf1815900aa8..88a7fdda5e733 100644 --- a/datafusion/pruning/src/lib.rs +++ b/datafusion/pruning/src/lib.rs @@ -18,6 +18,8 @@ #![cfg_attr(test, allow(clippy::needless_pass_by_value))] mod file_pruner; +mod in_list; +mod primitive_in_list; mod pruning_predicate; mod string_in_list; diff --git a/datafusion/pruning/src/primitive_in_list.rs b/datafusion/pruning/src/primitive_in_list.rs new file mode 100644 index 0000000000000..6b21a437862a6 --- /dev/null +++ b/datafusion/pruning/src/primitive_in_list.rs @@ -0,0 +1,430 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::fmt::{self, Display, Formatter}; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +use arrow::array::{Array, AsArray, BooleanArray}; +use arrow::compute::cast; +use arrow::datatypes::{ + ArrowPrimitiveType, DataType, Date32Type, Date64Type, Decimal32Type, Decimal64Type, + Decimal128Type, Decimal256Type, DurationMicrosecondType, DurationMillisecondType, + DurationNanosecondType, DurationSecondType, Int8Type, Int16Type, Int32Type, + Int64Type, Schema, Time32MillisecondType, Time32SecondType, Time64MicrosecondType, + Time64NanosecondType, TimeUnit, TimestampMicrosecondType, TimestampMillisecondType, + TimestampNanosecondType, TimestampSecondType, UInt8Type, UInt16Type, UInt32Type, + UInt64Type, +}; +use arrow::record_batch::RecordBatch; +use datafusion_common::{Result, ScalarValue, assert_eq_or_internal_err}; +use datafusion_physical_expr::{PhysicalExpr, PhysicalExprRef}; +use datafusion_physical_plan::ColumnarValue; + +use crate::in_list::{SetMembership, not_in_may_match}; + +/// A typed, non-null primitive IN-list domain under construction. +pub(crate) struct PrimitiveInListDomain { + data_type: DataType, + values: PrimitiveValues, +} + +macro_rules! define_primitive_values { + ($( + $variant:ident, $arrow_type:ty, $data_pattern:pat, + $scalar_pattern:pat, $native:expr $(, if $guard:expr)?; + )+) => { + enum PrimitiveValues { + $( + $variant(Vec<<$arrow_type as ArrowPrimitiveType>::Native>), + )+ + } + + impl PrimitiveInListDomain { + pub(crate) fn new( + data_type: &DataType, + capacity: usize, + ) -> Option { + // Parameter bindings in these patterns are used by `push` guards, + // but construction only selects the matching storage type. + #[expect( + unused_variables, + reason = "parameter bindings are reused by the generated push guards" + )] + let values = match data_type { + $( + $data_pattern => PrimitiveValues::$variant( + Vec::with_capacity(capacity), + ), + )+ + _ => return None, + }; + Some(Self { + data_type: data_type.clone(), + values, + }) + } + + /// Adds one non-null scalar whose logical type matches the domain. + pub(crate) fn push(&mut self, value: &ScalarValue) -> Option<()> { + match (&self.data_type, &mut self.values, value) { + $( + ($data_pattern, PrimitiveValues::$variant(values), $scalar_pattern) + $(if $guard)? => + { + values.push($native); + Some(()) + } + )+ + _ => None, + } + } + + pub(crate) fn is_empty(&self) -> bool { + match &self.values { + $(PrimitiveValues::$variant(values) => values.is_empty(),)+ + } + } + + pub(crate) fn into_expr( + self, + membership: SetMembership, + min: PhysicalExprRef, + max: PhysicalExprRef, + ) -> PhysicalExprRef { + let data_type = self.data_type; + match self.values { + $( + PrimitiveValues::$variant(values) => Arc::new( + PrimitiveInListPruningExpr::<$arrow_type>::new( + membership, + data_type, + min, + max, + values, + ), + ), + )+ + } + } + } + }; +} + +// This single list defines both supported types and expression dispatch. +define_primitive_values! { + Int8, Int8Type, DataType::Int8, + ScalarValue::Int8(Some(value)), *value; + Int16, Int16Type, DataType::Int16, + ScalarValue::Int16(Some(value)), *value; + Int32, Int32Type, DataType::Int32, + ScalarValue::Int32(Some(value)), *value; + Int64, Int64Type, DataType::Int64, + ScalarValue::Int64(Some(value)), *value; + UInt8, UInt8Type, DataType::UInt8, + ScalarValue::UInt8(Some(value)), *value; + UInt16, UInt16Type, DataType::UInt16, + ScalarValue::UInt16(Some(value)), *value; + UInt32, UInt32Type, DataType::UInt32, + ScalarValue::UInt32(Some(value)), *value; + UInt64, UInt64Type, DataType::UInt64, + ScalarValue::UInt64(Some(value)), *value; + Decimal32, Decimal32Type, DataType::Decimal32(_precision, scale), + ScalarValue::Decimal32(Some(value), _value_precision, value_scale), *value, + if scale == value_scale; + Decimal64, Decimal64Type, DataType::Decimal64(_precision, scale), + ScalarValue::Decimal64(Some(value), _value_precision, value_scale), *value, + if scale == value_scale; + Decimal128, Decimal128Type, DataType::Decimal128(_precision, scale), + ScalarValue::Decimal128(Some(value), _value_precision, value_scale), *value, + if scale == value_scale; + Decimal256, Decimal256Type, DataType::Decimal256(_precision, scale), + ScalarValue::Decimal256(Some(value), _value_precision, value_scale), *value, + if scale == value_scale; + Date32, Date32Type, DataType::Date32, + ScalarValue::Date32(Some(value)), *value; + Date64, Date64Type, DataType::Date64, + ScalarValue::Date64(Some(value)), *value; + Time32Second, Time32SecondType, DataType::Time32(TimeUnit::Second), + ScalarValue::Time32Second(Some(value)), *value; + Time32Millisecond, Time32MillisecondType, + DataType::Time32(TimeUnit::Millisecond), + ScalarValue::Time32Millisecond(Some(value)), *value; + Time64Microsecond, Time64MicrosecondType, + DataType::Time64(TimeUnit::Microsecond), + ScalarValue::Time64Microsecond(Some(value)), *value; + Time64Nanosecond, Time64NanosecondType, + DataType::Time64(TimeUnit::Nanosecond), + ScalarValue::Time64Nanosecond(Some(value)), *value; + TimestampSecond, TimestampSecondType, + DataType::Timestamp(TimeUnit::Second, _timezone), + ScalarValue::TimestampSecond(Some(value), _value_timezone), *value; + TimestampMillisecond, TimestampMillisecondType, + DataType::Timestamp(TimeUnit::Millisecond, _timezone), + ScalarValue::TimestampMillisecond(Some(value), _value_timezone), *value; + TimestampMicrosecond, TimestampMicrosecondType, + DataType::Timestamp(TimeUnit::Microsecond, _timezone), + ScalarValue::TimestampMicrosecond(Some(value), _value_timezone), *value; + TimestampNanosecond, TimestampNanosecondType, + DataType::Timestamp(TimeUnit::Nanosecond, _timezone), + ScalarValue::TimestampNanosecond(Some(value), _value_timezone), *value; + DurationSecond, DurationSecondType, DataType::Duration(TimeUnit::Second), + ScalarValue::DurationSecond(Some(value)), *value; + DurationMillisecond, DurationMillisecondType, + DataType::Duration(TimeUnit::Millisecond), + ScalarValue::DurationMillisecond(Some(value)), *value; + DurationMicrosecond, DurationMicrosecondType, + DataType::Duration(TimeUnit::Microsecond), + ScalarValue::DurationMicrosecond(Some(value)), *value; + DurationNanosecond, DurationNanosecondType, + DataType::Duration(TimeUnit::Nanosecond), + ScalarValue::DurationNanosecond(Some(value)), *value; +} + +/// Tests an inclusive statistics interval against a sorted primitive domain. +struct PrimitiveInListPruningExpr +where + T::Native: Eq + Hash + Ord, +{ + membership: SetMembership, + // Retain decimal metadata and timestamp timezone for statistics casts. + data_type: DataType, + min: PhysicalExprRef, + max: PhysicalExprRef, + values: Arc<[T::Native]>, +} + +impl fmt::Debug for PrimitiveInListPruningExpr +where + T::Native: Eq + Hash + Ord, +{ + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("PrimitiveInListPruningExpr") + .field("membership", &self.membership) + .field("data_type", &self.data_type) + .field("min", &self.min) + .field("max", &self.max) + .field("values", &self.values) + .finish() + } +} + +impl PrimitiveInListPruningExpr +where + T::Native: Eq + Hash + Ord, +{ + fn new( + membership: SetMembership, + data_type: DataType, + min: PhysicalExprRef, + max: PhysicalExprRef, + mut values: Vec, + ) -> Self { + values.sort_unstable(); + values.dedup(); + Self { + membership, + data_type, + min, + max, + values: values.into(), + } + } + + fn contains(&self, value: T::Native) -> bool { + self.values.binary_search(&value).is_ok() + } + + fn may_match(&self, min: Option, max: Option) -> Option { + if self.membership == SetMembership::NotIn { + return not_in_may_match(min.as_ref(), max.as_ref(), |value| { + self.contains(*value) + }); + } + match (min, max) { + (Some(min), Some(max)) => { + // Inverted statistics are unusable, not proof that the domain + // and interval are disjoint. + if min > max { + return None; + } + let index = self.values.partition_point(|value| *value < min); + Some(self.values.get(index).is_some_and(|value| *value <= max)) + } + (Some(min), None) if self.values.last().is_some_and(|value| *value < min) => { + Some(false) + } + (None, Some(max)) + if self.values.first().is_some_and(|value| *value > max) => + { + Some(false) + } + _ => None, + } + } +} + +impl PartialEq for PrimitiveInListPruningExpr +where + T::Native: Eq + Hash + Ord, +{ + fn eq(&self, other: &Self) -> bool { + self.membership == other.membership + && self.data_type == other.data_type + && self.min.eq(&other.min) + && self.max.eq(&other.max) + && self.values == other.values + } +} + +impl Eq for PrimitiveInListPruningExpr where + T::Native: Eq + Hash + Ord +{ +} + +impl Hash for PrimitiveInListPruningExpr +where + T::Native: Eq + Hash + Ord, +{ + fn hash(&self, state: &mut H) { + self.membership.hash(state); + self.data_type.hash(state); + self.min.hash(state); + self.max.hash(state); + self.values.hash(state); + } +} + +impl Display for PrimitiveInListPruningExpr +where + T::Native: Eq + Hash + Ord, +{ + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!( + f, + "{}({}, {}, {} values)", + self.membership.display_name(), + self.min, + self.max, + self.values.len() + ) + } +} + +impl PhysicalExpr for PrimitiveInListPruningExpr +where + T::Native: Eq + Hash + Ord, +{ + fn data_type(&self, _input_schema: &Schema) -> Result { + Ok(DataType::Boolean) + } + + fn nullable(&self, _input_schema: &Schema) -> Result { + Ok(true) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + let min = self.min.evaluate(batch)?.into_array(batch.num_rows())?; + let max = self.max.evaluate(batch)?.into_array(batch.num_rows())?; + // Unlike view casts, dictionary-to-primitive casts propagate both key + // and value NULLs, so no logical-null intersection is necessary here. + let min = cast(&min, &self.data_type)?; + let max = cast(&max, &self.data_type)?; + let min = min.as_primitive::(); + let max = max.as_primitive::(); + let matches: BooleanArray = (0..batch.num_rows()) + .map(|index| { + let min = min.is_valid(index).then(|| min.value(index)); + let max = max.is_valid(index).then(|| max.value(index)); + self.may_match(min, max) + }) + .collect(); + Ok(ColumnarValue::Array(Arc::new(matches))) + } + + fn children(&self) -> Vec<&PhysicalExprRef> { + vec![&self.min, &self.max] + } + + fn with_new_children( + self: Arc, + children: Vec, + ) -> Result { + assert_eq_or_internal_err!(children.len(), 2); + Ok(Arc::new(Self { + membership: self.membership, + data_type: self.data_type.clone(), + min: Arc::clone(&children[0]), + max: Arc::clone(&children[1]), + values: Arc::clone(&self.values), + })) + } + + fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "{self}") + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::in_list::unwrap_scalar; + use arrow::datatypes::Field; + + #[test] + fn domain_accepts_compatible_non_null_values() { + let mut domain = PrimitiveInListDomain::new(&DataType::Int64, 4).unwrap(); + assert!(domain.push(&ScalarValue::Int64(None)).is_none()); + assert!(domain.push(&ScalarValue::Int32(Some(1))).is_none()); + assert!(domain.is_empty()); + + let run_end_encoded = ScalarValue::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int16, false)), + Arc::new(Field::new("values", DataType::Int64, true)), + Box::new(ScalarValue::Int64(Some(1))), + ); + assert!(domain.push(unwrap_scalar(&run_end_encoded)).is_some()); + assert!(!domain.is_empty()); + + let mut decimal = + PrimitiveInListDomain::new(&DataType::Decimal128(10, 2), 2).unwrap(); + assert!( + decimal + .push(&ScalarValue::Decimal128(Some(1), 12, 2)) + .is_some() + ); + assert!( + decimal + .push(&ScalarValue::Decimal128(Some(1), 10, 3)) + .is_none() + ); + + let mut timestamp = PrimitiveInListDomain::new( + &DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + 1, + ) + .unwrap(); + assert!( + timestamp + .push(&ScalarValue::TimestampMicrosecond( + Some(1), + Some("Asia/Kolkata".into()), + )) + .is_some() + ); + } +} diff --git a/datafusion/pruning/src/pruning_predicate.rs b/datafusion/pruning/src/pruning_predicate.rs index c3362c63299e1..b49b72058e0cd 100644 --- a/datafusion/pruning/src/pruning_predicate.rs +++ b/datafusion/pruning/src/pruning_predicate.rs @@ -22,7 +22,9 @@ use std::collections::HashSet; use std::sync::Arc; -use crate::string_in_list::{SetMembership, StringInListPruningExpr}; +use crate::in_list::{SetMembership, unwrap_scalar}; +use crate::primitive_in_list::PrimitiveInListDomain; +use crate::string_in_list::{BinaryInListPruningExpr, StringInListPruningExpr}; use arrow::array::AsArray; use arrow::{ @@ -461,8 +463,8 @@ impl<'a> PruningPredicateBuilder<'a> { /// | Condition | Pruning representation | /// | --- | --- | /// | `N <= min(20, C)` | Existing per-value rewrite | - /// | `20 < N <= C`, literal strings with optional NULLs on a string column | Compact pruning expression | - /// | `20 < N <= C`, other lists | Existing per-value rewrite | + /// | `20 < N <= C`, supported ordered literals with optional NULLs | Compact pruning expression | + /// | `20 < N <= C`, unsupported lists | Existing per-value rewrite | /// | `N > C` | Unhandled-predicate hook, normally "keep the container" | /// /// Empty lists also use the unhandled-predicate hook. A cap of zero disables @@ -470,9 +472,13 @@ impl<'a> PruningPredicateBuilder<'a> { /// pruning (such as Bloom filters). The default cap is [`MAX_IN_LIST_SIZE`] /// (20), so the compact path requires an explicitly raised cap. /// + /// Compact domains support strings, binary values, integers, decimals, + /// dates, times, timestamps, and durations. Floating-point values are + /// excluded because NaN and signed zero do not follow the required order. + /// /// The compact form covers `IN` and `NOT IN` alike, including lists with /// NULL members. Raising the cap can still build large comparison trees for - /// other eligible lists. + /// unsupported lists. /// /// Query engines typically pass /// `datafusion.execution.parquet.max_in_list_size` here. @@ -1497,8 +1503,25 @@ fn build_is_null_column_expr( } } -/// Keep large literal string lists compact instead of building a per-value -/// tree: an OR tree for `IN`, an AND chain for `NOT IN`. +/// Values collected for one supported compact IN-list representation. +enum CompactInListDomain { + String(Vec), + Binary(Vec>), + Primitive(PrimitiveInListDomain), +} + +impl CompactInListDomain { + fn is_empty(&self) -> bool { + match self { + Self::String(values) => values.is_empty(), + Self::Binary(values) => values.is_empty(), + Self::Primitive(values) => values.is_empty(), + } + } +} + +/// Keep large literal lists of supported ordered types compact instead of +/// building a per-value tree: an OR tree for `IN`, an AND chain for `NOT IN`. /// /// `IN` excludes a container whose interval is disjoint from the domain. That /// matches the per-value OR tree except for inverted bounds, which the compact @@ -1510,8 +1533,9 @@ fn build_is_null_column_expr( /// everywhere, including absent and inverted bounds. /// /// A NULL list member makes `NOT IN` and an all-NULL `IN` list never TRUE. For -/// other `IN` lists, NULL does not change which rows can make the predicate TRUE. -fn build_string_in_list_expr( +/// other `IN` lists, NULL does not change which rows can make the predicate TRUE, +/// but prevents inversion from proving that a container is fully matched. +fn build_compact_in_list_expr( in_list: &phys_expr::InListExpr, schema: &Schema, required_columns: &mut RequiredColumns, @@ -1528,42 +1552,79 @@ fn build_string_in_list_expr( DataType::Dictionary(_, value) => value.as_ref(), data_type => data_type, }; - if field.name() != column.name() || !data_type.is_string() { + if field.name() != column.name() { return None; } - let mut values = Vec::with_capacity(in_list.list().len()); + let mut domain = if data_type.is_string() { + CompactInListDomain::String(Vec::with_capacity(in_list.list().len())) + } else if matches!( + data_type, + DataType::Binary | DataType::LargeBinary | DataType::BinaryView + ) { + CompactInListDomain::Binary(Vec::with_capacity(in_list.list().len())) + } else { + CompactInListDomain::Primitive(PrimitiveInListDomain::new( + data_type, + in_list.list().len(), + )?) + }; let mut contains_null = false; for expr in in_list.list() { - if let Some(value) = extract_string_literal(expr) { - values.push(value.to_owned()); - } else if expr - .downcast_ref::() - .is_some_and(|literal| literal.value().is_null()) - { + let literal = expr.downcast_ref::()?; + let value = unwrap_scalar(literal.value()); + if value.is_null() { contains_null = true; } else { - return None; + match &mut domain { + CompactInListDomain::String(values) => { + values.push(unpack_string(value)?.to_owned()); + } + CompactInListDomain::Binary(values) => { + values.push(extract_binary(value)?.into()); + } + CompactInListDomain::Primitive(values) => values.push(value)?, + } } } // Pruning asks only whether the predicate can be TRUE. UNKNOWN and FALSE // both reject a row, so NOT IN with a NULL member and an all-NULL IN list // can never match. IN can otherwise ignore NULL and search its non-null domain. - if contains_null && (membership == SetMembership::NotIn || values.is_empty()) { + if contains_null && (membership == SetMembership::NotIn || domain.is_empty()) { properties.has_filter_semantics_only = true; return Some(Arc::new(phys_expr::Literal::new(ScalarValue::Boolean( Some(false), )))); } - let min = required_columns - .min_column_expr(column, in_list.expr(), field) - .ok()?; - let max = required_columns - .max_column_expr(column, in_list.expr(), field) - .ok()?; - let non_null = - build_is_null_column_expr(in_list.expr(), schema, required_columns, true)?; - let may_match = Arc::new(StringInListPruningExpr::new(membership, min, max, values)); + // Roll back appended statistics columns if the compact rewrite cannot be + // completed. `RequiredColumns::stat_column_expr` only appends entries. + let required_columns_len = required_columns.columns.len(); + let statistics = (|| { + let min = required_columns + .min_column_expr(column, in_list.expr(), field) + .ok()?; + let max = required_columns + .max_column_expr(column, in_list.expr(), field) + .ok()?; + let non_null = + build_is_null_column_expr(in_list.expr(), schema, required_columns, true)?; + Some((min, max, non_null)) + })(); + let Some((min, max, non_null)) = statistics else { + required_columns.columns.truncate(required_columns_len); + return None; + }; + let may_match = match domain { + CompactInListDomain::String(values) => { + Arc::new(StringInListPruningExpr::new(membership, min, max, values)) + as PhysicalExprRef + } + CompactInListDomain::Binary(values) => { + Arc::new(BinaryInListPruningExpr::new(membership, min, max, values)) + as PhysicalExprRef + } + CompactInListDomain::Primitive(values) => values.into_expr(membership, min, max), + }; if contains_null { properties.has_filter_semantics_only = true; } @@ -1575,9 +1636,9 @@ fn build_string_in_list_expr( } /// Default maximum number of entries in an `IN (...)` list eligible for -/// statistics pruning. Eligible literal string lists above this threshold use a -/// compact sorted domain instead of per-value min/max checks, for both `IN` and -/// `NOT IN`. +/// statistics pruning. Eligible literal lists of supported ordered types above +/// this threshold use a compact sorted domain instead of per-value min/max +/// checks, for both `IN` and `NOT IN`. /// Callers can raise the cap via [`PredicateRewriter::with_max_in_list_size`], and /// query engines can wire it from the /// `datafusion.execution.parquet.max_in_list_size` config option. @@ -1665,9 +1726,9 @@ impl PredicateRewriter { /// Returns the pruning predicate as an [`PhysicalExpr`] /// /// `max_in_list_size` is the largest `IN (...)` list eligible for statistics -/// pruning. Large literal string lists use a compact representation, for both -/// `IN` and `NOT IN`; other eligible lists use per-value checks. Longer lists -/// fall back to `unhandled_hook`. +/// pruning. Large literal lists of supported ordered types use a compact +/// representation for both `IN` and `NOT IN`; unsupported lists use per-value +/// checks. Longer lists fall back to `unhandled_hook`. fn build_predicate_expression( expr: &Arc, schema: &SchemaRef, @@ -1716,7 +1777,7 @@ fn build_predicate_expression( if in_list.list().len() > MAX_IN_LIST_SIZE && in_list.list().len() <= max_in_list_size && let Some(pruning_expr) = - build_string_in_list_expr(in_list, schema, required_columns, properties) + build_compact_in_list_expr(in_list, schema, required_columns, properties) { return pruning_expr; } @@ -2085,6 +2146,15 @@ fn extract_string_literal(expr: &Arc) -> Option<&str> { None } +fn extract_binary(value: &ScalarValue) -> Option<&[u8]> { + match value { + ScalarValue::Binary(value) + | ScalarValue::LargeBinary(value) + | ScalarValue::BinaryView(value) => value.as_deref(), + _ => None, + } +} + /// Wrap a string in a `Literal` whose `ScalarValue` matches `target_type` fn string_literal_as(value: String, target_type: &DataType) -> Arc { let utf8 = ScalarValue::Utf8(Some(value)); @@ -2327,8 +2397,8 @@ mod tests { BinaryArray, DictionaryArray, Int32Array, Int64Array, StringArray, UInt64Array, }, - buffer::NullBuffer, - datatypes::{Int32Type, TimeUnit}, + buffer::{NullBuffer, ScalarBuffer}, + datatypes::{Int32Type, Int64Type, TimeUnit, i256}, }; use datafusion_expr::expr::InList; use datafusion_expr::{BinaryExpr, Expr, cast, is_null, try_cast}; @@ -3602,13 +3672,8 @@ mod tests { Ok(()) } - // With the configurable cap, a caller that raises - // `max_in_list_size` above the default gets the IN list rewritten - // into a per-value min/max chain instead of falling through to `true`. - // This verifies both `PredicateRewriter::with_max_in_list_size` and the - // recursive OR path inside `build_predicate_expression`. #[test] - fn row_group_predicate_in_list_rewritten_at_raised_cap() -> Result<()> { + fn row_group_predicate_in_list_compacted_at_raised_cap() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, false)])); // 25 items — above the default 20, below a raised cap of 32. @@ -3617,19 +3682,9 @@ mod tests { let rewriter = PredicateRewriter::new().with_max_in_list_size(32); let predicate_expr = rewriter.rewrite_predicate_to_statistics_predicate(&physical, &schema); - // At the raised cap, IN is rewritten into per-value min/max checks - // OR'd together; the resulting predicate must not collapse to - // `true` (which is what the default cap produces). - assert_ne!( - predicate_expr.to_string(), - "true", - "IN(25) with raised cap must rewrite into a statistics-based predicate, not fall through to `true`" - ); - // Sanity: the rewritten predicate references per-value literals. assert!( - predicate_expr.to_string().contains(" <= 1 ") - && predicate_expr.to_string().contains(" <= 25 "), - "rewritten predicate should include per-value bounds for each IN entry, got: {predicate_expr}" + predicate_expr.to_string().contains("IN_SET_INTERSECTS"), + "raised-cap predicate should use compact pruning, got: {predicate_expr}" ); Ok(()) } @@ -3657,7 +3712,7 @@ mod tests { // The high-level [`PruningPredicateBuilder`] should thread // `max_in_list_size` all the way through: a 25-item IN with the default // cap must fall through to the unhandled hook (`predicate_expr = true`), - // while a raised cap produces a real per-value statistics predicate. + // while a raised cap produces a compact statistics predicate. #[test] fn pruning_predicate_builder_threads_max_in_list_size() -> Result<()> { let schema = @@ -3677,8 +3732,6 @@ mod tests { "default cap must fall through to `true` for 25-item IN" ); - // Raising the cap produces a real statistics predicate with per- - // value bounds. let raised_pp = PruningPredicateBuilder::new() .with_file_schema(Arc::clone(&schema)) .with_max_in_list_size(32) @@ -3689,8 +3742,8 @@ mod tests { "raised cap must produce a real statistics predicate for 25-item IN" ); assert!( - raised_expr.contains(" <= 1 ") && raised_expr.contains(" <= 25 "), - "raised-cap predicate should include per-value bounds, got: {raised_expr}" + raised_expr.contains("IN_SET_INTERSECTS"), + "raised-cap predicate should use compact pruning, got: {raised_expr}" ); Ok(()) } @@ -3719,7 +3772,7 @@ mod tests { Ok(()) } - fn large_string_pruning_predicate( + fn large_in_list_pruning_predicate( expr: PhysicalExprRef, schema: SchemaRef, ) -> Result { @@ -3729,6 +3782,388 @@ mod tests { .try_build(expr) } + fn assert_per_value_fallback( + predicate: &PruningPredicate, + expected_literals: &[ScalarValue], + ) -> Result<()> { + let expression = predicate.predicate_expr(); + let display = expression.to_string(); + assert_ne!(display, "true"); + assert!(!display.contains("IN_SET_INTERSECTS")); + assert!(!display.contains("NOT_IN_SET_MAY_MATCH")); + assert!(display.contains("_min")); + assert!(display.contains("_max")); + let mut nodes = 0; + let mut literals = vec![]; + expression.apply(|expr| { + nodes += 1; + if let Some(literal) = expr.downcast_ref::() { + literals.push(literal.value().clone()); + } + Ok(TreeNodeRecursion::Continue) + })?; + assert!(nodes > 7, "expected a per-value tree, got {expression}"); + for expected in expected_literals { + assert!( + literals.contains(expected), + "expected literal {expected} in per-value tree {expression}" + ); + } + Ok(()) + } + + #[test] + fn compact_extraction_unwraps_encoded_binary_literals() { + let value = ScalarValue::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int16, false)), + Arc::new(Field::new("values", DataType::Binary, true)), + Box::new(ScalarValue::Binary(Some(vec![1, 2]))), + ); + assert_eq!(extract_binary(unwrap_scalar(&value)), Some(&[1, 2][..])); + } + + fn ordered_scalar(data_type: &DataType, value: Option) -> ScalarValue { + match data_type { + DataType::Int8 => ScalarValue::Int8(value.map(|value| value as i8)), + DataType::Int16 => ScalarValue::Int16(value.map(|value| value as i16)), + DataType::Int32 => ScalarValue::Int32(value.map(|value| value as i32)), + DataType::Int64 => ScalarValue::Int64(value), + DataType::UInt8 => ScalarValue::UInt8(value.map(|value| value as u8)), + DataType::UInt16 => ScalarValue::UInt16(value.map(|value| value as u16)), + DataType::UInt32 => ScalarValue::UInt32(value.map(|value| value as u32)), + DataType::UInt64 => ScalarValue::UInt64(value.map(|value| value as u64)), + DataType::Decimal32(precision, scale) => ScalarValue::Decimal32( + value.map(|value| value as i32), + *precision, + *scale, + ), + DataType::Decimal64(precision, scale) => { + ScalarValue::Decimal64(value, *precision, *scale) + } + DataType::Decimal128(precision, scale) => { + ScalarValue::Decimal128(value.map(i128::from), *precision, *scale) + } + DataType::Decimal256(precision, scale) => { + ScalarValue::Decimal256(value.map(i256::from), *precision, *scale) + } + DataType::Date32 => ScalarValue::Date32(value.map(|value| value as i32)), + DataType::Date64 => ScalarValue::Date64(value), + DataType::Time32(TimeUnit::Second) => { + ScalarValue::Time32Second(value.map(|value| value as i32)) + } + DataType::Time32(TimeUnit::Millisecond) => { + ScalarValue::Time32Millisecond(value.map(|value| value as i32)) + } + DataType::Time64(TimeUnit::Microsecond) => { + ScalarValue::Time64Microsecond(value) + } + DataType::Time64(TimeUnit::Nanosecond) => { + ScalarValue::Time64Nanosecond(value) + } + DataType::Timestamp(TimeUnit::Second, timezone) => { + ScalarValue::TimestampSecond(value, timezone.clone()) + } + DataType::Timestamp(TimeUnit::Millisecond, timezone) => { + ScalarValue::TimestampMillisecond(value, timezone.clone()) + } + DataType::Timestamp(TimeUnit::Microsecond, timezone) => { + ScalarValue::TimestampMicrosecond(value, timezone.clone()) + } + DataType::Timestamp(TimeUnit::Nanosecond, timezone) => { + ScalarValue::TimestampNanosecond(value, timezone.clone()) + } + DataType::Duration(TimeUnit::Second) => ScalarValue::DurationSecond(value), + DataType::Duration(TimeUnit::Millisecond) => { + ScalarValue::DurationMillisecond(value) + } + DataType::Duration(TimeUnit::Microsecond) => { + ScalarValue::DurationMicrosecond(value) + } + DataType::Duration(TimeUnit::Nanosecond) => { + ScalarValue::DurationNanosecond(value) + } + DataType::Binary => { + ScalarValue::Binary(value.map(|value| vec![0xff, value as u8])) + } + DataType::LargeBinary => { + ScalarValue::LargeBinary(value.map(|value| vec![0xff, value as u8])) + } + DataType::BinaryView => { + ScalarValue::BinaryView(value.map(|value| vec![0xff, value as u8])) + } + DataType::Dictionary(key_type, value_type) => ScalarValue::Dictionary( + key_type.clone(), + Box::new(ordered_scalar(value_type, value)), + ), + data_type => panic!("unsupported test type {data_type:?}"), + } + } + + fn ordered_statistics(data_type: &DataType) -> ContainerStats { + let statistics_type = match data_type { + DataType::Dictionary(_, value_type) => value_type.as_ref(), + data_type => data_type, + }; + let min = [Some(0), Some(1), Some(1), Some(50), None, Some(41), Some(4)]; + let max = [Some(0), Some(1), Some(2), Some(51), Some(2), None, Some(2)]; + ContainerStats::new() + .with_min( + ScalarValue::iter_to_array( + min.into_iter() + .map(|value| ordered_scalar(statistics_type, value)), + ) + .unwrap(), + ) + .with_max( + ScalarValue::iter_to_array( + max.into_iter() + .map(|value| ordered_scalar(statistics_type, value)), + ) + .unwrap(), + ) + .with_null_counts([Some(0); 7]) + .with_row_counts([Some(1); 7]) + } + + #[test] + fn large_ordered_in_lists_use_compact_pruning() -> Result<()> { + let types = [ + DataType::Int8, + DataType::Int16, + DataType::Int32, + DataType::Int64, + DataType::UInt8, + DataType::UInt16, + DataType::UInt32, + DataType::UInt64, + DataType::Decimal32(9, 2), + DataType::Decimal64(18, 3), + DataType::Decimal128(30, 4), + DataType::Decimal256(60, 5), + DataType::Date32, + DataType::Date64, + DataType::Time32(TimeUnit::Second), + DataType::Time32(TimeUnit::Millisecond), + DataType::Time64(TimeUnit::Microsecond), + DataType::Time64(TimeUnit::Nanosecond), + DataType::Timestamp(TimeUnit::Second, None), + DataType::Timestamp(TimeUnit::Millisecond, None), + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + DataType::Timestamp(TimeUnit::Nanosecond, None), + DataType::Duration(TimeUnit::Second), + DataType::Duration(TimeUnit::Millisecond), + DataType::Duration(TimeUnit::Microsecond), + DataType::Duration(TimeUnit::Nanosecond), + DataType::Binary, + DataType::LargeBinary, + DataType::BinaryView, + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Int64)), + DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Binary)), + ]; + + for data_type in types { + let schema = + Arc::new(Schema::new(vec![Field::new("c1", data_type.clone(), true)])); + let values = (0..21) + .map(|index| { + Arc::new(phys_expr::Literal::new(ordered_scalar( + &data_type, + Some(index * 2), + ))) as PhysicalExprRef + }) + .collect::>(); + let stats = TestStatistics::new().with("c1", ordered_statistics(&data_type)); + + for (negated, expected) in [ + (false, [true, false, true, false, true, false, true]), + (true, [false, true, true, true, true, true, true]), + ] { + let expr = phys_expr::in_list( + Arc::new(phys_expr::Column::new("c1", 0)), + values.clone(), + &negated, + &schema, + )?; + let predicate = large_in_list_pruning_predicate(expr, schema.clone())?; + assert!( + predicate.predicate_expr().to_string().contains(if negated { + "NOT_IN_SET_MAY_MATCH" + } else { + "IN_SET_INTERSECTS" + }), + "type={data_type:?}, negated={negated}" + ); + assert_eq!( + predicate.prune(&stats)?, + expected, + "type={data_type:?}, negated={negated}" + ); + assert_eq!(predicate.required_columns.columns.len(), 4); + } + } + Ok(()) + } + + #[test] + fn large_float_in_list_keeps_per_value_rewrite() -> Result<()> { + let data_type = DataType::Float64; + let schema = + Arc::new(Schema::new(vec![Field::new("c1", data_type.clone(), true)])); + let values = (0..21) + .map(|value| { + Arc::new(phys_expr::Literal::new(ScalarValue::Float64(Some( + value as f64, + )))) as PhysicalExprRef + }) + .collect::>(); + for negated in [false, true] { + let expr = phys_expr::in_list( + Arc::new(phys_expr::Column::new("c1", 0)), + values.clone(), + &negated, + &schema, + )?; + let predicate = large_in_list_pruning_predicate(expr, Arc::clone(&schema))?; + assert_per_value_fallback( + &predicate, + &[ + ScalarValue::Float64(Some(0.0)), + ScalarValue::Float64(Some(20.0)), + ], + )?; + } + Ok(()) + } + + #[test] + fn large_fixed_size_binary_in_list_keeps_per_value_rewrite() -> Result<()> { + let data_type = DataType::FixedSizeBinary(2); + let schema = + Arc::new(Schema::new(vec![Field::new("c1", data_type.clone(), true)])); + let values = (0..21) + .map(|value| { + Arc::new(phys_expr::Literal::new(ScalarValue::FixedSizeBinary( + 2, + Some(vec![0xff, value]), + ))) as PhysicalExprRef + }) + .collect::>(); + for negated in [false, true] { + let expr = phys_expr::in_list( + Arc::new(phys_expr::Column::new("c1", 0)), + values.clone(), + &negated, + &schema, + )?; + let predicate = large_in_list_pruning_predicate(expr, Arc::clone(&schema))?; + assert_per_value_fallback( + &predicate, + &[ + ScalarValue::FixedSizeBinary(2, Some(vec![0xff, 0])), + ScalarValue::FixedSizeBinary(2, Some(vec![0xff, 20])), + ], + )?; + } + Ok(()) + } + + #[test] + fn large_primitive_in_list_preserves_dictionary_nulls() -> Result<()> { + let data_type = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Int64)); + let schema = + Arc::new(Schema::new(vec![Field::new("c1", data_type.clone(), true)])); + let values = (0..21) + .map(|value| { + Arc::new(phys_expr::Literal::new(ordered_scalar( + &data_type, + Some(value * 2), + ))) as PhysicalExprRef + }) + .collect::>(); + let expr = phys_expr::in_list( + Arc::new(phys_expr::Column::new("c1", 0)), + values, + &false, + &schema, + )?; + let predicate = large_in_list_pruning_predicate(expr, schema)?; + + // The invalid dictionary values have nonzero payloads outside the + // domain. They must remain missing bounds after dictionary casting. + let dictionary_values = Arc::new(arrow::array::PrimitiveArray::::new( + ScalarBuffer::from(vec![50, 51, -1, -1]), + Some(NullBuffer::from(vec![false, true, true, false])), + )); + let dictionary = |keys| -> Result { + Ok(Arc::new(DictionaryArray::::try_new( + Int32Array::from(keys), + dictionary_values.clone(), + )?)) + }; + let stats = TestStatistics::new().with( + "c1", + ContainerStats::new() + .with_min(dictionary(vec![Some(0), Some(2)])?) + .with_max(dictionary(vec![Some(1), Some(3)])?) + .with_null_counts([Some(0); 2]) + .with_row_counts([Some(1); 2]), + ); + assert_eq!(predicate.prune(&stats)?, [true, true]); + Ok(()) + } + + #[test] + fn large_binary_in_list_preserves_dictionary_nulls() -> Result<()> { + let data_type = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Binary)); + let schema = + Arc::new(Schema::new(vec![Field::new("c1", data_type.clone(), true)])); + let values = (0..21) + .map(|value| { + Arc::new(phys_expr::Literal::new(ordered_scalar( + &data_type, + Some(value * 2), + ))) as PhysicalExprRef + }) + .collect::>(); + let expr = phys_expr::in_list( + Arc::new(phys_expr::Column::new("c1", 0)), + values, + &false, + &schema, + )?; + let predicate = large_in_list_pruning_predicate(expr, schema)?; + + // BinaryView casts currently lose NULLs in dictionary values. Use + // nonempty invalid payloads so the test fails if they become bounds. + let payloads = [vec![0xff, 50], vec![0xff, 51], vec![0], vec![0]]; + let (offsets, values, _) = + BinaryArray::from_iter_values(payloads.iter().map(Vec::as_slice)) + .into_parts(); + let dictionary_values: ArrayRef = Arc::new(BinaryArray::new( + offsets, + values, + Some(NullBuffer::from(vec![false, true, true, false])), + )); + let dictionary = |keys| -> Result { + Ok(Arc::new(DictionaryArray::::try_new( + Int32Array::from(keys), + Arc::clone(&dictionary_values), + )?)) + }; + let stats = TestStatistics::new().with( + "c1", + ContainerStats::new() + .with_min(dictionary(vec![Some(0), Some(2)])?) + .with_max(dictionary(vec![Some(1), Some(3)])?) + .with_null_counts([Some(0); 2]) + .with_row_counts([Some(1); 2]), + ); + assert_eq!(predicate.prune(&stats)?, [true, true]); + Ok(()) + } + #[test] fn large_string_in_list_prunes_exact_intervals() -> Result<()> { let types = [ @@ -3758,7 +4193,7 @@ mod tests { &schema, )?; let predicate = - large_string_pruning_predicate(Arc::clone(&expr), schema)?; + large_in_list_pruning_predicate(Arc::clone(&expr), schema)?; let last_value = format!("k{:06}", (count - 1) * 10); let stats = TestStatistics::new().with( "c1", @@ -3845,7 +4280,7 @@ mod tests { &false, &schema, )?; - let predicate = large_string_pruning_predicate(expr, schema)?; + let predicate = large_in_list_pruning_predicate(expr, schema)?; // NULL dictionary values can have nonempty payloads. Neither payload // may become a bound when a valid key references the NULL value. let (offsets, values, _) = @@ -3906,7 +4341,7 @@ mod tests { ]); let expr = col("c1").in_list(values, false); let predicate = - large_string_pruning_predicate(logical2physical(&expr, &schema), schema)?; + large_in_list_pruning_predicate(logical2physical(&expr, &schema), schema)?; let stats = TestStatistics::new().with( "c1", ContainerStats::new_utf8( @@ -3992,7 +4427,7 @@ mod tests { } #[test] - fn large_string_in_list_compacts_null_literals() -> Result<()> { + fn large_in_list_compacts_null_literals() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)])); let values = (0..21).map(|i| lit(format!("a{i:03}"))).collect::>(); let stats = TestStatistics::new().with( @@ -4020,7 +4455,7 @@ mod tests { let positive_or = col("c1") .in_list(with_null.clone(), false) .or(col("c1").eq(lit("middle"))); - let predicate = large_string_pruning_predicate( + let predicate = large_in_list_pruning_predicate( logical2physical(&positive_or, &schema), Arc::clone(&schema), )?; @@ -4037,7 +4472,7 @@ mod tests { assert_eq!(default.prune(&stats)?, [true, true]); assert!(default.can_be_inverted_for_full_match()); - let raised = large_string_pruning_predicate(physical, Arc::clone(&schema))?; + let raised = large_in_list_pruning_predicate(physical, Arc::clone(&schema))?; assert!(!raised.can_be_inverted_for_full_match()); if negated { assert!(is_always_false(raised.predicate_expr()), "{expr}"); @@ -4058,7 +4493,7 @@ mod tests { std::iter::repeat_n(lit(ScalarValue::Utf8(None)), 21).collect::>(); for negated in [false, true] { let expr = col("c1").in_list(all_null.clone(), negated); - let predicate = large_string_pruning_predicate( + let predicate = large_in_list_pruning_predicate( logical2physical(&expr, &schema), Arc::clone(&schema), )?; @@ -4071,11 +4506,17 @@ mod tests { Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, true)])); let mut integer_values = (0..21).map(lit).collect::>(); integer_values.push(lit(ScalarValue::Int32(None))); - let integer_predicate = large_string_pruning_predicate( + let integer_predicate = large_in_list_pruning_predicate( logical2physical(&col("c1").in_list(integer_values, false), &integer_schema), integer_schema, )?; - assert!(integer_predicate.can_be_inverted_for_full_match()); + assert!( + integer_predicate + .predicate_expr() + .to_string() + .contains("IN_SET_INTERSECTS") + ); + assert!(!integer_predicate.can_be_inverted_for_full_match()); // The compact false predicate has the same filter result as the original // NOT IN expression, which returns UNKNOWN for values outside the list. @@ -4088,6 +4529,66 @@ mod tests { Ok(()) } + #[test] + fn compact_binary_and_primitive_null_lists_are_always_false() -> Result<()> { + let cases = [ + ( + DataType::Binary, + (0..21) + .map(|value| ScalarValue::Binary(Some(vec![value]))) + .collect::>(), + ScalarValue::Binary(None), + ), + ( + DataType::Int32, + (0..21) + .map(|value| ScalarValue::Int32(Some(value))) + .collect(), + ScalarValue::Int32(None), + ), + ]; + + for (data_type, values, null) in cases { + let schema = + Arc::new(Schema::new(vec![Field::new("c1", data_type.clone(), true)])); + let column = Arc::new(phys_expr::Column::new("c1", 0)); + + let mut with_null = values; + with_null.push(null.clone()); + let not_in = phys_expr::in_list( + Arc::clone(&column) as PhysicalExprRef, + with_null + .into_iter() + .map(|value| { + Arc::new(phys_expr::Literal::new(value)) as PhysicalExprRef + }) + .collect(), + &true, + &schema, + )?; + let predicate = large_in_list_pruning_predicate(not_in, Arc::clone(&schema))?; + assert!(is_always_false(predicate.predicate_expr())); + assert!(!predicate.can_be_inverted_for_full_match()); + assert!(predicate.required_columns.columns.is_empty()); + + let all_null = phys_expr::in_list( + column, + std::iter::repeat_n(null, 21) + .map(|value| { + Arc::new(phys_expr::Literal::new(value)) as PhysicalExprRef + }) + .collect(), + &false, + &schema, + )?; + let predicate = large_in_list_pruning_predicate(all_null, schema)?; + assert!(is_always_false(predicate.predicate_expr())); + assert!(!predicate.can_be_inverted_for_full_match()); + assert!(predicate.required_columns.columns.is_empty()); + } + Ok(()) + } + /// Statistics shared by the compact `NOT IN` tests, one row per case. fn not_in_container_stats() -> TestStatistics { TestStatistics::new().with( @@ -4158,7 +4659,7 @@ mod tests { &true, &schema, )?; - let predicate = large_string_pruning_predicate(expr, schema)?; + let predicate = large_in_list_pruning_predicate(expr, schema)?; assert_eq!( predicate.predicate_expr().to_string(), "c1_null_count@3 != row_count@2 AND NOT_IN_SET_MAY_MATCH(c1_min@0, c1_max@1, 21 values)", @@ -4208,7 +4709,7 @@ mod tests { let not_in = col("c1").in_list((0..21).map(|i| lit(format!("a{i:03}"))).collect(), true); let other = col("c2").in_list((0..21).map(|i| lit(i * 10)).collect(), false); - let predicate = large_string_pruning_predicate( + let predicate = large_in_list_pruning_predicate( logical2physical(¬_in.or(other), &schema), Arc::clone(&schema), )?; @@ -4248,13 +4749,13 @@ mod tests { fn large_string_not_in_list_inverts_without_false_full_match() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)])); let values = (0..21).map(|i| lit(format!("a{i:03}"))).collect::>(); - let forward = large_string_pruning_predicate( + let forward = large_in_list_pruning_predicate( logical2physical(&col("c1").in_list(values.clone(), true), &schema), Arc::clone(&schema), )?; // NOT(c1 NOT IN (...)) OR c1 IS NULL, the shape row_group_filter builds // once PhysicalExprSimplifier turns NOT(NOT IN) back into IN. - let inverted = large_string_pruning_predicate( + let inverted = large_in_list_pruning_predicate( logical2physical( &col("c1").in_list(values, false).or(col("c1").is_null()), &schema, @@ -4300,7 +4801,7 @@ mod tests { values.iter().map(|value| lit(value.clone())).collect(), true, ); - let compact = large_string_pruning_predicate( + let compact = large_in_list_pruning_predicate( logical2physical(¬_in, &schema), Arc::clone(&schema), )?; @@ -4312,7 +4813,7 @@ mod tests { .map(|value| col("c1").not_eq(lit(value.clone()))) .reduce(Expr::and) .unwrap(); - let per_value = large_string_pruning_predicate( + let per_value = large_in_list_pruning_predicate( logical2physical(&chain, &schema), Arc::clone(&schema), )?; diff --git a/datafusion/pruning/src/string_in_list.rs b/datafusion/pruning/src/string_in_list.rs index b9d9f68b30511..da1793f5d03a5 100644 --- a/datafusion/pruning/src/string_in_list.rs +++ b/datafusion/pruning/src/string_in_list.rs @@ -15,7 +15,6 @@ // specific language governing permissions and limitations // under the License. -use std::cmp::Ordering; use std::fmt::{self, Display, Formatter}; use std::hash::{Hash, Hasher}; use std::sync::Arc; @@ -28,115 +27,7 @@ use datafusion_common::{Result, assert_eq_or_internal_err}; use datafusion_physical_expr::{PhysicalExpr, PhysicalExprRef}; use datafusion_physical_plan::ColumnarValue; -/// Which `IN` form a sorted string domain is pruning for. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub(crate) enum SetMembership { - /// `col IN (...)`. A row matches only where the domain intersects the - /// interval, so a disjoint interval excludes every row. - In, - /// `col NOT IN (...)`. Overlap proves nothing here: values outside the - /// domain still satisfy the predicate. An interval excludes every row only - /// when it holds a single value that the domain contains. - NotIn, -} - -impl SetMembership { - fn compare_values(self, left: &[u8], right: &[u8]) -> Ordering { - match self { - // IN uses this order for interval searches. - Self::In => left.cmp(right), - // NOT IN only needs exact membership. Reject impossible lengths - // before comparing bytes that may have a long common prefix. - Self::NotIn => left.len().cmp(&right.len()).then_with(|| left.cmp(right)), - } - } -} - -/// Tests an inclusive statistics interval against a sorted string domain. -/// -/// [`PhysicalExpr::evaluate`] returns one nullable Boolean per min/max interval: -/// * `true`: matching rows may exist, so the container must be read. -/// * `false`: the available bounds prove no row can match. -/// * `NULL`: incomplete, invalid, or unusable bounds prevent a safe decision. -/// -/// [`SetMembership`] selects the test. For [`SetMembership::In`] a single known -/// bound can still prove disjointness. For [`SetMembership::NotIn`], one known -/// bound outside the domain proves the container may match, while two equal -/// bounds in the domain prove it cannot. This is the same reach as the per-value -/// `min != v OR v != max` chain it replaces. Otherwise, unknown results keep the -/// container eligible for reading. -/// -/// This expression is used only for pruning; the original IN remains the row filter. -#[derive(Debug, Eq)] -pub(crate) struct StringInListPruningExpr { - membership: SetMembership, - min: PhysicalExprRef, - max: PhysicalExprRef, - values: Arc<[String]>, -} - -impl StringInListPruningExpr { - pub(crate) fn new( - membership: SetMembership, - min: PhysicalExprRef, - max: PhysicalExprRef, - mut values: Vec, - ) -> Self { - values.sort_unstable_by(|left, right| { - membership.compare_values(left.as_bytes(), right.as_bytes()) - }); - values.dedup(); - Self { - membership, - min, - max, - values: values.into(), - } - } - - /// Does the sorted, deduplicated domain hold `value`? - fn contains(&self, value: &[u8]) -> bool { - self.values - .binary_search_by(|candidate| { - self.membership.compare_values(candidate.as_bytes(), value) - }) - .is_ok() - } -} - -impl PartialEq for StringInListPruningExpr { - fn eq(&self, other: &Self) -> bool { - self.membership == other.membership - && self.min.eq(&other.min) - && self.max.eq(&other.max) - && self.values == other.values - } -} - -impl Hash for StringInListPruningExpr { - fn hash(&self, state: &mut H) { - self.membership.hash(state); - self.min.hash(state); - self.max.hash(state); - self.values.hash(state); - } -} - -impl Display for StringInListPruningExpr { - fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { - let name = match self.membership { - SetMembership::In => "IN_SET_INTERSECTS", - SetMembership::NotIn => "NOT_IN_SET_MAY_MATCH", - }; - write!( - f, - "{name}({}, {}, {} values)", - self.min, - self.max, - self.values.len() - ) - } -} +use crate::in_list::{SetMembership, not_in_may_match}; fn has_oversized_string_buffer(array: &dyn Array, limit: usize) -> bool { match array.data_type() { @@ -150,166 +41,317 @@ fn has_oversized_string_buffer(array: &dyn Array, limit: usize) -> bool { } } -impl PhysicalExpr for StringInListPruningExpr { - fn data_type(&self, _input_schema: &Schema) -> Result { - Ok(DataType::Boolean) +fn has_oversized_binary_buffer(array: &dyn Array, limit: usize) -> bool { + match array.data_type() { + DataType::Binary => array.as_binary::().values().len() >= limit, + DataType::LargeBinary => array.as_binary::().values().len() >= limit, + DataType::Dictionary(_, _) => has_oversized_binary_buffer( + array.as_any_dictionary().values().as_ref(), + limit, + ), + _ => false, } +} - fn nullable(&self, _input_schema: &Schema) -> Result { - Ok(true) - } +macro_rules! string_bytes { + ($value:expr) => { + ($value).as_bytes() + }; +} + +macro_rules! binary_bytes { + ($value:expr) => { + &($value)[..] + }; +} - fn evaluate(&self, batch: &RecordBatch) -> Result { - // Normalize Utf8, LargeUtf8, Utf8View, and dictionary-encoded statistics. - let min = self.min.evaluate(batch)?.into_array(batch.num_rows())?; - let max = self.max.evaluate(batch)?.into_array(batch.num_rows())?; - // A short string slice can retain a buffer too large for Utf8View's - // u32 offsets. Avoid a panic in the cast and keep pruning conservative. - if has_oversized_string_buffer(min.as_ref(), u32::MAX as usize) - || has_oversized_string_buffer(max.as_ref(), u32::MAX as usize) - { - return Ok(ColumnarValue::Array(Arc::new(BooleanArray::new_null( - batch.num_rows(), - )))); +/// Defines a byte-domain pruning expression for one variable-length value type. +macro_rules! define_byte_in_list_expr { + ( + $name:ident, + $value_type:ty, + $view_type:expr, + $as_view:ident, + $bytes:ident, + $has_oversized_buffer:path + ) => { + /// Tests an inclusive statistics interval against a sorted byte domain. + /// + /// [`PhysicalExpr::evaluate`] returns one nullable Boolean per min/max + /// interval: `true` means matching rows may exist, `false` means the + /// available bounds prove no row can match, and `NULL` means the bounds + /// cannot support a safe decision. + /// + /// [`SetMembership::In`] can use one known bound to prove disjointness. + /// For [`SetMembership::NotIn`], one known bound outside the domain + /// proves the container may match, while equal bounds in the domain + /// prove it cannot. Otherwise, the result keeps the container eligible. + /// The original `IN` expression remains the row filter. + #[derive(Debug, Eq)] + pub(crate) struct $name { + membership: SetMembership, + min: PhysicalExprRef, + max: PhysicalExprRef, + values: Arc<[$value_type]>, } - // Dictionary values can be NULL behind valid keys. Preserve their - // validity even if the view cast only carries the key nulls. - // TODO: Revisit this workaround once the Arrow dependency includes - // https://github.com/apache/arrow-rs/pull/10510. - let min_nulls = min.logical_nulls(); - let max_nulls = max.logical_nulls(); - let min = cast(&min, &DataType::Utf8View)?; - let max = cast(&max, &DataType::Utf8View)?; - let min = min.as_string_view(); - let max = max.as_string_view(); - let matches: BooleanArray = (0..batch.num_rows()) - .map(|i| { - let min = (min.is_valid(i) - && min_nulls.as_ref().is_none_or(|nulls| nulls.is_valid(i))) - .then(|| min.value(i).as_bytes()); - let max = (max.is_valid(i) - && max_nulls.as_ref().is_none_or(|nulls| nulls.is_valid(i))) - .then(|| max.value(i).as_bytes()); - if self.membership == SetMembership::NotIn { - return match (min, max) { - // Check membership before comparing the bounds. NOT IN - // ordering rejects different byte lengths first, avoiding - // scans of long common prefixes. - (Some(min), Some(max)) if self.contains(min) => { - // A wider interval can hold a value outside the domain, - // which satisfies NOT IN. Only an interval pinned to one - // domain value rules out every row. Truncated Parquet - // bounds cannot fake that: min truncates downward and max - // upward, so equal bounds mean the true values were equal. - Some(min != max) - } - (Some(_), Some(_)) => Some(true), - // One known bound outside the domain is enough to preserve - // the true result that lets an enclosing OR short-circuit. - (Some(bound), None) | (None, Some(bound)) - if !self.contains(bound) => - { - Some(true) - } - _ => None, - }; + + impl $name { + pub(crate) fn new( + membership: SetMembership, + min: PhysicalExprRef, + max: PhysicalExprRef, + mut values: Vec<$value_type>, + ) -> Self { + values.sort_unstable_by(|left, right| { + membership.compare_bytes($bytes!(left), $bytes!(right)) + }); + values.dedup(); + Self { + membership, + min, + max, + values: values.into(), } - match (min, max) { - (Some(min), Some(max)) => { - if min > max { - return None; - } - // Rust string ordering and these byte comparisons both use - // unsigned lexicographic UTF-8 order, as required by the - // PruningStatistics min/max contract. Parquet adapters mask - // bounds with unusable ordering; PartitionPruningStatistics - // uses actual Arrow partition values. PrunableStatistics - // trusts file providers' bounds: there is no ordering gate - // for arbitrary statistics providers here. - let index = self.values.partition_point(|v| v.as_bytes() < min); - Some(self.values.get(index).is_some_and(|v| v.as_bytes() <= max)) - } - // A missing bound makes that end of the interval unbounded. - // Exclude only when the whole domain lies beyond the known bound; - // gaps within the domain and equality cannot prove disjointness. - (Some(min), None) - if self.values.last().is_some_and(|v| v.as_bytes() < min) => - { - Some(false) - } - (None, Some(max)) - if self.values.first().is_some_and(|v| v.as_bytes() > max) => - { - Some(false) - } - _ => None, + } + + /// Does the sorted, deduplicated domain hold `value`? + fn contains(&self, value: &[u8]) -> bool { + self.values + .binary_search_by(|candidate| { + self.membership.compare_bytes($bytes!(candidate), value) + }) + .is_ok() + } + } + + impl PartialEq for $name { + fn eq(&self, other: &Self) -> bool { + self.membership == other.membership + && self.min.eq(&other.min) + && self.max.eq(&other.max) + && self.values == other.values + } + } + + impl Hash for $name { + fn hash(&self, state: &mut H) { + self.membership.hash(state); + self.min.hash(state); + self.max.hash(state); + self.values.hash(state); + } + } + + impl Display for $name { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + let name = self.membership.display_name(); + write!( + f, + "{name}({}, {}, {} values)", + self.min, + self.max, + self.values.len() + ) + } + } + + impl PhysicalExpr for $name { + fn data_type(&self, _input_schema: &Schema) -> Result { + Ok(DataType::Boolean) + } + + fn nullable(&self, _input_schema: &Schema) -> Result { + Ok(true) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + // Normalize byte arrays and dictionary-encoded statistics to a view. + let min = self.min.evaluate(batch)?.into_array(batch.num_rows())?; + let max = self.max.evaluate(batch)?.into_array(batch.num_rows())?; + // A short slice can retain a buffer too large for view arrays' u32 + // offsets. Avoid a panic in the cast and keep pruning conservative. + if $has_oversized_buffer(min.as_ref(), u32::MAX as usize) + || $has_oversized_buffer(max.as_ref(), u32::MAX as usize) + { + return Ok(ColumnarValue::Array(Arc::new(BooleanArray::new_null( + batch.num_rows(), + )))); } - }) - .collect(); - Ok(ColumnarValue::Array(Arc::new(matches))) - } + // Unlike primitive casts, view casts can drop NULLs stored behind + // valid dictionary keys. Preserve that logical validity explicitly. + // TODO: Revisit this workaround once the Arrow dependency includes + // https://github.com/apache/arrow-rs/pull/10510. + let min_nulls = min.logical_nulls(); + let max_nulls = max.logical_nulls(); + let min = cast(&min, &$view_type)?; + let max = cast(&max, &$view_type)?; + let min = min.$as_view(); + let max = max.$as_view(); + let matches: BooleanArray = (0..batch.num_rows()) + .map(|i| { + let min: Option<&[u8]> = (min.is_valid(i) + && min_nulls.as_ref().is_none_or(|nulls| nulls.is_valid(i))) + .then(|| $bytes!(min.value(i))); + let max: Option<&[u8]> = (max.is_valid(i) + && max_nulls.as_ref().is_none_or(|nulls| nulls.is_valid(i))) + .then(|| $bytes!(max.value(i))); + if self.membership == SetMembership::NotIn { + // Membership is checked before bound equality. NOT IN ordering + // rejects different byte lengths first, avoiding scans of long + // common prefixes. + return not_in_may_match(min, max, |value| { + self.contains(value) + }); + } + match (min, max) { + (Some(min), Some(max)) => { + if min > max { + return None; + } + // String and binary values use unsigned lexicographic byte + // order, as required by the PruningStatistics min/max + // contract. Parquet adapters mask bounds with unusable + // ordering; PartitionPruningStatistics uses actual Arrow + // partition values. PrunableStatistics trusts file + // providers' bounds: there is no ordering gate for arbitrary + // statistics providers here. + let index = + self.values.partition_point(|v| $bytes!(v) < min); + Some( + self.values + .get(index) + .is_some_and(|v| $bytes!(v) <= max), + ) + } + // A missing bound makes that end of the interval unbounded. + // Exclude only when the whole domain lies beyond the known bound; + // gaps within the domain and equality cannot prove disjointness. + (Some(min), None) + if self + .values + .last() + .is_some_and(|v| $bytes!(v) < min) => + { + Some(false) + } + (None, Some(max)) + if self + .values + .first() + .is_some_and(|v| $bytes!(v) > max) => + { + Some(false) + } + _ => None, + } + }) + .collect(); + Ok(ColumnarValue::Array(Arc::new(matches))) + } - fn children(&self) -> Vec<&PhysicalExprRef> { - vec![&self.min, &self.max] - } + fn children(&self) -> Vec<&PhysicalExprRef> { + vec![&self.min, &self.max] + } - fn with_new_children( - self: Arc, - children: Vec, - ) -> Result { - assert_eq_or_internal_err!(children.len(), 2); - Ok(Arc::new(Self { - membership: self.membership, - min: Arc::clone(&children[0]), - max: Arc::clone(&children[1]), - values: Arc::clone(&self.values), - })) - } + fn with_new_children( + self: Arc, + children: Vec, + ) -> Result { + assert_eq_or_internal_err!(children.len(), 2); + Ok(Arc::new(Self { + membership: self.membership, + min: Arc::clone(&children[0]), + max: Arc::clone(&children[1]), + values: Arc::clone(&self.values), + })) + } - fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result { - write!(f, "{self}") - } + fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "{self}") + } + } + }; } +define_byte_in_list_expr!( + StringInListPruningExpr, + String, + DataType::Utf8View, + as_string_view, + string_bytes, + has_oversized_string_buffer +); + +define_byte_in_list_expr!( + BinaryInListPruningExpr, + Box<[u8]>, + DataType::BinaryView, + as_binary_view, + binary_bytes, + has_oversized_binary_buffer +); + #[cfg(test)] mod tests { use super::*; - use arrow::array::{ArrayRef, StringArray}; + use arrow::array::{ArrayRef, BinaryArray, StringArray}; - #[test] - fn oversized_buffers_check_retained_data_not_visible_offsets() -> Result<()> { - // Exercise the size boundary without allocating a 4 GiB buffer. - let limit = 32; - let padding = "p".repeat(limit - 1); - let array: ArrayRef = Arc::new(StringArray::from(vec!["a", padding.as_str()])); - - for value_type in [DataType::Utf8, DataType::LargeUtf8] { + fn assert_oversized_buffers( + array: &ArrayRef, + value_types: impl IntoIterator, + view_type: DataType, + limit: usize, + has_oversized_buffer: fn(&dyn Array, usize) -> bool, + ) -> Result<()> { + for value_type in value_types { for data_type in [ value_type.clone(), DataType::Dictionary( Box::new(DataType::Int32), Box::new(value_type.clone()), ), - DataType::Dictionary( - Box::new(DataType::UInt64), - Box::new(value_type.clone()), - ), + DataType::Dictionary(Box::new(DataType::UInt64), Box::new(value_type)), ] { - let slice = cast(&array, &data_type)?.slice(0, 1); - assert!(has_oversized_string_buffer(slice.as_ref(), limit - 1)); - assert!(has_oversized_string_buffer(slice.as_ref(), limit)); - assert!(!has_oversized_string_buffer(slice.as_ref(), limit + 1)); + let slice = cast(array, &data_type)?.slice(0, 1); + assert!(has_oversized_buffer(slice.as_ref(), limit - 1)); + assert!(has_oversized_buffer(slice.as_ref(), limit)); + assert!(!has_oversized_buffer(slice.as_ref(), limit + 1)); } } - // Already-normalized views do not have the byte-array cast limitation. for data_type in [ - DataType::Utf8View, - DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8View)), + view_type.clone(), + DataType::Dictionary(Box::new(DataType::Int32), Box::new(view_type)), ] { - let slice = cast(&array, &data_type)?.slice(0, 1); - assert!(!has_oversized_string_buffer(slice.as_ref(), limit)); + let slice = cast(array, &data_type)?.slice(0, 1); + assert!(!has_oversized_buffer(slice.as_ref(), limit)); } Ok(()) } + + #[test] + fn oversized_buffers_check_retained_data_not_visible_offsets() -> Result<()> { + // Exercise the size boundary without allocating a 4 GiB buffer. + let limit = 32; + let padding = "p".repeat(limit - 1); + let array: ArrayRef = Arc::new(StringArray::from(vec!["a", padding.as_str()])); + assert_oversized_buffers( + &array, + [DataType::Utf8, DataType::LargeUtf8], + DataType::Utf8View, + limit, + has_oversized_string_buffer, + )?; + + let padding = vec![0; limit - 1]; + let array: ArrayRef = Arc::new(BinaryArray::from(vec![&[1][..], &padding])); + assert_oversized_buffers( + &array, + [DataType::Binary, DataType::LargeBinary], + DataType::BinaryView, + limit, + has_oversized_binary_buffer, + )?; + Ok(()) + } } diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index f2a587ed72ae0..4ef57e4ae2662 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -415,7 +415,7 @@ datafusion.execution.parquet.dictionary_page_size_limit 1048576 (writing) Sets b datafusion.execution.parquet.enable_page_index true (reading) If true, reads the Parquet data page level metadata (the Page Index), if present, to reduce the I/O and number of rows decoded. datafusion.execution.parquet.encoding NULL (writing) Sets default encoding for any column. Valid values are: plain, plain_dictionary, rle, bit_packed, delta_binary_packed, delta_length_byte_array, delta_byte_array, rle_dictionary, and byte_stream_split. These values are not case sensitive. If NULL, uses default parquet writer setting datafusion.execution.parquet.force_filter_selections false (reading) Force the use of RowSelections for filter results, when pushdown_filters is enabled. If false, the reader will automatically choose between a RowSelection and a Bitmap based on the number and pattern of selected rows. -datafusion.execution.parquet.max_in_list_size 20 Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger literal string lists on a string column use a compact representation, for both `IN` and `NOT IN`, including lists with NULL members. `NOT IN` with NULL and all-NULL `IN` lists cannot match any rows. Other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. +datafusion.execution.parquet.max_in_list_size 20 Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger literal lists use a compact representation when the column type is string, variable-length binary, integer, decimal, date, time, timestamp, or duration. This applies to both `IN` and `NOT IN`, including lists with NULL members. `NOT IN` with NULL and all-NULL `IN` lists cannot match any rows. Compact lists containing NULL do not use the fully-matched-row-group optimization. Floating-point and other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. datafusion.execution.parquet.max_predicate_cache_size NULL (reading) The maximum predicate cache size, in bytes. When `pushdown_filters` is enabled, sets the maximum memory used to cache the results of predicate evaluation between filter evaluation and output generation. Decreasing this value will reduce memory usage, but may increase IO and CPU usage. None means use the default parquet reader setting. 0 means no caching. datafusion.execution.parquet.max_row_group_bytes NULL (writing) Target maximum size of each row group in bytes. When set, the writer flushes whenever either this limit or `max_row_group_size` is reached, whichever comes first. Useful for bounding writer memory on wide schemas where a row-count limit can map to very different byte sizes. Matches the behavior of `parquet.block.size` in parquet-mr. If `None` (the default), only the row-count limit applies. Currently only honored when `allow_single_file_parallelism` is `false`; by default the parallel file writer ignores this limit. datafusion.execution.parquet.max_row_group_size 1048576 (writing) Target maximum number of rows in each row group (defaults to 1M rows). Writing larger row groups requires more memory to write, but can get better compression and be faster to read. When `max_row_group_bytes` is also set, the writer flushes a row group when either limit is reached, whichever comes first. diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index ab344028cdab4..f7e47ade4f635 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -93,7 +93,7 @@ The following configuration settings are available: | datafusion.execution.parquet.coerce_int96_tz | NULL | (reading) Optional timezone applied to INT96 columns when `coerce_int96` is set. When `Some`, INT96 columns coerce to `Timestamp(, Some())` instead of the default `Timestamp(, None)`. Spark and other systems write INT96 values as UTC-adjusted instants, so callers that need the resulting Arrow type to be timezone-aware (e.g. for Spark `TimestampType` semantics) should set this to `"UTC"`. No effect when `coerce_int96` is `None`. | | datafusion.execution.parquet.bloom_filter_on_read | true | (reading) Use any available bloom filters when reading parquet files | | datafusion.execution.parquet.max_predicate_cache_size | NULL | (reading) The maximum predicate cache size, in bytes. When `pushdown_filters` is enabled, sets the maximum memory used to cache the results of predicate evaluation between filter evaluation and output generation. Decreasing this value will reduce memory usage, but may increase IO and CPU usage. None means use the default parquet reader setting. 0 means no caching. | -| datafusion.execution.parquet.max_in_list_size | 20 | Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger literal string lists on a string column use a compact representation, for both `IN` and `NOT IN`, including lists with NULL members. `NOT IN` with NULL and all-NULL `IN` lists cannot match any rows. Other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. | +| datafusion.execution.parquet.max_in_list_size | 20 | Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger literal lists use a compact representation when the column type is string, variable-length binary, integer, decimal, date, time, timestamp, or duration. This applies to both `IN` and `NOT IN`, including lists with NULL members. `NOT IN` with NULL and all-NULL `IN` lists cannot match any rows. Compact lists containing NULL do not use the fully-matched-row-group optimization. Floating-point and other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. | | datafusion.execution.parquet.data_pagesize_limit | 1048576 | (writing) Sets best effort maximum size of data page in bytes | | datafusion.execution.parquet.write_batch_size | 1024 | (writing) Sets write_batch_size in rows | | datafusion.execution.parquet.writer_version | 1.0 | (writing) Sets parquet writer version valid values are "1.0" and "2.0" |