diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 7c7604e04136..d2bd0ca23322 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1373,10 +1373,10 @@ config_namespace! { /// rewrite; other predicates and Bloom-filter pruning remain available. /// /// Within the cap, nonempty lists of at most 20 values use the existing - /// per-value rewrite. Larger positive, non-null literal string lists - /// on a string column use a compact sorted domain. Other lists retain - /// the existing per-value rewrite, so raising the cap can make those - /// predicates expensive to build and evaluate. + /// per-value rewrite. Larger non-null literal string lists on a string + /// column use a compact sorted domain, for both `IN` and `NOT IN`. + /// Other lists retain the existing per-value rewrite, so raising the cap + /// can make those predicates expensive to build and evaluate. /// /// Defaults to 20. pub max_in_list_size: usize, default = 20 diff --git a/datafusion/core/tests/parquet/string_in_list_pruning.rs b/datafusion/core/tests/parquet/string_in_list_pruning.rs index e1005daaa1a3..ad588a9325e3 100644 --- a/datafusion/core/tests/parquet/string_in_list_pruning.rs +++ b/datafusion/core/tests/parquet/string_in_list_pruning.rs @@ -15,8 +15,8 @@ // specific language governing permissions and limitations // under the License. -//! End-to-end coverage for compact, large string IN-list pruning. The positive -//! IN-list cases disable the row and Bloom filters to isolate min/max pruning. +//! End-to-end coverage for compact, large string IN-list pruning. The `IN` and +//! `NOT IN` cases disable the row and Bloom filters to isolate min/max pruning. use std::sync::Arc; @@ -47,10 +47,29 @@ const UNITS: usize = 4; const TOTAL_ROWS: usize = ROWS_PER_UNIT * UNITS; const MATCHING_ROWS: usize = ROWS_PER_UNIT * 2; -/// Write either four row groups or four pages in one row group. The second -/// unit lies in a gap between two members of every test IN list; an enclosing -/// min/max range for the list cannot prune it. +/// Write either four row groups or four pages in one row group. Each unit holds +/// a single repeated value, so `NOT IN` can exclude the two units whose value is +/// a list member. The second unit lies in a gap between two members of every +/// test IN list; an enclosing min/max range for the list cannot prune it. fn make_file(page_pruning: bool) -> NamedTempFile { + let values = ["v000000", "v000001", "v000010", "v999999"] + .into_iter() + .flat_map(|value| std::iter::repeat_n(value, ROWS_PER_UNIT)) + .collect::>(); + write_file(page_pruning, values) +} + +/// Write one mixed unit and one single-valued unit. The mixed unit contains +/// both a list member and a value that satisfies `NOT IN`. +fn make_mixed_not_in_file(page_pruning: bool) -> NamedTempFile { + let values = std::iter::repeat_n("v000000", ROWS_PER_UNIT / 2) + .chain(std::iter::repeat_n("v000001", ROWS_PER_UNIT / 2)) + .chain(std::iter::repeat_n("v000010", ROWS_PER_UNIT)) + .collect::>(); + write_file(page_pruning, values) +} + +fn write_file(page_pruning: bool, values: Vec<&str>) -> NamedTempFile { let mut file = tempfile::Builder::new() .prefix("string_in_list_pruning") .suffix(".parquet") @@ -61,17 +80,15 @@ fn make_file(page_pruning: bool) -> NamedTempFile { DataType::Utf8, false, )])); - let values = ["v000000", "v000001", "v000010", "v999999"] - .into_iter() - .flat_map(|value| std::iter::repeat_n(value, ROWS_PER_UNIT)) - .collect::>(); + let total_rows = values.len(); + assert_eq!(total_rows % ROWS_PER_UNIT, 0); let batch = RecordBatch::try_new( Arc::clone(&schema), vec![Arc::new(StringArray::from(values))], ) .unwrap(); let rows_per_group = if page_pruning { - TOTAL_ROWS + total_rows } else { ROWS_PER_UNIT }; @@ -86,7 +103,7 @@ fn make_file(page_pruning: bool) -> NamedTempFile { let mut writer = ArrowWriter::try_new(&mut file, schema, Some(properties)).unwrap(); writer.write(&batch).unwrap(); let metadata = writer.close().unwrap(); - assert_eq!(metadata.num_row_groups(), TOTAL_ROWS / rows_per_group); + assert_eq!(metadata.num_row_groups(), total_rows / rows_per_group); let offsets = metadata.offset_index().unwrap(); for row_group in offsets { assert_eq!( @@ -151,6 +168,27 @@ impl ScanOutput { ], &self.batches ); + self.assert_no_filter_interference(); + } + + /// The complement of [`Self::assert_results`]: the two units whose value is + /// not a list member. + fn assert_negated_results(&self) { + assert_batches_eq!( + [ + "+---------+----+", + "| value | n |", + "+---------+----+", + "| v000001 | 16 |", + "| v999999 | 16 |", + "+---------+----+", + ], + &self.batches + ); + self.assert_no_filter_interference(); + } + + fn assert_no_filter_interference(&self) { assert_eq!(self.counter("predicate_evaluation_errors"), 0); assert_eq!(self.counter("pushdown_rows_pruned"), 0); assert_eq!(self.pruned("row_groups_pruned_bloom_filter"), 0); @@ -162,6 +200,7 @@ async fn scan( list_size: usize, max_in_list_size: Option, page_pruning: bool, + negated: bool, ) -> ScanOutput { let mut config = SessionConfig::new() .with_target_partitions(1) @@ -183,9 +222,10 @@ async fn scan( .map(|index| format!("'v{:06}'", index * 10)) .collect::>() .join(", "); + let op = if negated { "NOT IN" } else { "IN" }; let sql = format!( "SELECT value, count(*) AS n FROM t \ - WHERE value IN ({values}) GROUP BY value ORDER BY value" + WHERE value {op} ({values}) GROUP BY value ORDER BY value" ); let plan = ctx .sql(&sql) @@ -209,14 +249,14 @@ async fn check_string_in_list_pruning(page_pruning: bool) { for list_size in [20, 21, 256, 1024] { // A zero cap provides a result-equivalence control that cannot use // min/max IN-list pruning at either granularity. - let unpruned = scan(&file, list_size, Some(0), page_pruning).await; + let unpruned = scan(&file, list_size, Some(0), page_pruning, false).await; unpruned.assert_results(); assert!(!unpruned.plan.contains("IN_SET_INTERSECTS")); assert_eq!(unpruned.pruned("row_groups_pruned_statistics"), 0); assert_eq!(unpruned.pruned("page_index_rows_pruned"), 0); assert_eq!(unpruned.counter("output_rows"), TOTAL_ROWS); - let output = scan(&file, list_size, Some(list_size), page_pruning).await; + let output = scan(&file, list_size, Some(list_size), page_pruning, false).await; output.assert_results(); assert_eq!( pretty_format_batches(&output.batches).unwrap().to_string(), @@ -247,7 +287,7 @@ async fn check_string_in_list_pruning(page_pruning: bool) { // The default remains 20: enabling the compact representation must not // silently change the public cap's meaning. - let default = scan(&file, 21, None, page_pruning).await; + let default = scan(&file, 21, None, page_pruning, false).await; default.assert_results(); assert!(!default.plan.contains("IN_SET_INTERSECTS")); assert_eq!(default.pruned("row_groups_pruned_statistics"), 0); @@ -255,6 +295,96 @@ async fn check_string_in_list_pruning(page_pruning: bool) { assert_eq!(default.counter("output_rows"), TOTAL_ROWS); } +/// The compact `NOT IN` form must prune exactly the units whose single repeated +/// value is a list member, and nothing else. Overlapping a list member is not +/// enough: units 1 and 3 each sit inside the list's enclosing range. +async fn check_string_not_in_list_pruning(page_pruning: bool) { + let file = make_file(page_pruning); + for list_size in [20, 21, 256, 1024] { + let unpruned = scan(&file, list_size, Some(0), page_pruning, true).await; + unpruned.assert_negated_results(); + assert!(!unpruned.plan.contains("NOT_IN_SET_MAY_MATCH")); + assert_eq!(unpruned.pruned("row_groups_pruned_statistics"), 0); + assert_eq!(unpruned.pruned("page_index_rows_pruned"), 0); + assert_eq!(unpruned.counter("output_rows"), TOTAL_ROWS); + + let output = scan(&file, list_size, Some(list_size), page_pruning, true).await; + output.assert_negated_results(); + assert_eq!( + pretty_format_batches(&output.batches).unwrap().to_string(), + pretty_format_batches(&unpruned.batches) + .unwrap() + .to_string() + ); + assert_eq!( + output.plan.contains("NOT_IN_SET_MAY_MATCH"), + list_size > 20, + "list_size={list_size}, plan={}", + output.plan + ); + // The compact form and the per-value AND chain it replaces prune the + // same units, so the counts do not depend on which one ran. + assert_eq!( + output.pruned("row_groups_pruned_statistics"), + if page_pruning { 0 } else { 2 }, + "list_size={list_size}, metrics={}", + output.metrics + ); + assert_eq!( + output.pruned("page_index_rows_pruned"), + if page_pruning { MATCHING_ROWS } else { 0 }, + "list_size={list_size}, metrics={}", + output.metrics + ); + assert_eq!(output.counter("output_rows"), MATCHING_ROWS); + } + + let default = scan(&file, 21, None, page_pruning, true).await; + default.assert_negated_results(); + assert!(!default.plan.contains("NOT_IN_SET_MAY_MATCH")); + assert_eq!(default.pruned("row_groups_pruned_statistics"), 0); + assert_eq!(default.pruned("page_index_rows_pruned"), 0); + assert_eq!(default.counter("output_rows"), TOTAL_ROWS); + + // A mixed interval that overlaps a list member can still contain matching + // rows. Only the adjacent single-valued unit can be excluded. + let mixed_file = make_mixed_not_in_file(page_pruning); + let unpruned = scan(&mixed_file, 21, Some(0), page_pruning, true).await; + assert_batches_eq!( + [ + "+---------+---+", + "| value | n |", + "+---------+---+", + "| v000001 | 8 |", + "+---------+---+", + ], + &unpruned.batches + ); + unpruned.assert_no_filter_interference(); + assert_eq!(unpruned.pruned("row_groups_pruned_statistics"), 0); + assert_eq!(unpruned.pruned("page_index_rows_pruned"), 0); + assert_eq!(unpruned.counter("output_rows"), ROWS_PER_UNIT * 2); + + let output = scan(&mixed_file, 21, Some(21), page_pruning, true).await; + assert_eq!( + pretty_format_batches(&output.batches).unwrap().to_string(), + pretty_format_batches(&unpruned.batches) + .unwrap() + .to_string() + ); + output.assert_no_filter_interference(); + assert!(output.plan.contains("NOT_IN_SET_MAY_MATCH")); + assert_eq!( + output.pruned("row_groups_pruned_statistics"), + usize::from(!page_pruning) + ); + assert_eq!( + output.pruned("page_index_rows_pruned"), + if page_pruning { ROWS_PER_UNIT } else { 0 } + ); + assert_eq!(output.counter("output_rows"), ROWS_PER_UNIT); +} + #[tokio::test] async fn string_in_list_row_group_pruning() { check_string_in_list_pruning(false).await; @@ -265,6 +395,16 @@ async fn string_in_list_page_pruning() { check_string_in_list_pruning(true).await; } +#[tokio::test] +async fn string_not_in_list_row_group_pruning() { + check_string_not_in_list_pruning(false).await; +} + +#[tokio::test] +async fn string_not_in_list_page_pruning() { + check_string_not_in_list_pruning(true).await; +} + #[tokio::test] async fn string_not_in_list_with_null_does_not_bypass_row_filter() { let mut file = tempfile::Builder::new() diff --git a/datafusion/pruning/benches/string_in_list_pruning.rs b/datafusion/pruning/benches/string_in_list_pruning.rs index d26112611bf3..c542e9be327f 100644 --- a/datafusion/pruning/benches/string_in_list_pruning.rs +++ b/datafusion/pruning/benches/string_in_list_pruning.rs @@ -17,12 +17,16 @@ //! Compare compact string IN-list pruning with per-value min/max expansion. //! -//! Both cases raise `max_in_list_size` to the domain size. On a baseline -//! without compact pruning, `in_list` measures the ordinary raised-cap path. -//! The explicit `expanded_or` is a balanced tree of equalities, which produces -//! the same per-value statistics checks without making the baseline depend on -//! a deeply nested expression. Half of the statistics intervals hit a domain -//! member and half fall in a sparse gap. Bloom filters are not involved. +//! Every case raises `max_in_list_size` to the domain size. On a baseline +//! without compact pruning, `in_list` and `not_in_list` measure the ordinary +//! raised-cap path. The explicit `expanded_or` and `expanded_and` are balanced +//! trees of (in)equalities, which produce the same per-value statistics checks +//! without making the baseline depend on a deeply nested expression. +//! +//! Half of the statistics intervals are pinned to a domain member and half span +//! a sparse gap, so `IN` keeps the first half and `NOT IN` keeps the second. +//! Both directions therefore prune real containers rather than measuring an +//! always-true fallback. Bloom filters are not involved. //! //! Run with `cargo bench -p datafusion-pruning --bench string_in_list_pruning`. //! The construction benchmarks reuse their input physical expressions; the @@ -48,18 +52,35 @@ fn value(index: usize) -> String { format!("key{index:08}") } -fn balanced_or(expressions: &[PhysicalExprRef]) -> PhysicalExprRef { +fn balanced(expressions: &[PhysicalExprRef], op: Operator) -> PhysicalExprRef { if expressions.len() == 1 { return Arc::clone(&expressions[0]); } let middle = expressions.len() / 2; Arc::new(BinaryExpr::new( - balanced_or(&expressions[..middle]), - Operator::Or, - balanced_or(&expressions[middle..]), + balanced(&expressions[..middle], op), + op, + balanced(&expressions[middle..], op), )) } +/// A balanced tree of `column value`, one branch per domain member. +fn expanded( + column: &PhysicalExprRef, + values: &[PhysicalExprRef], + op: Operator, + combine: Operator, +) -> PhysicalExprRef { + let comparisons = values + .iter() + .map(|value| { + Arc::new(BinaryExpr::new(Arc::clone(column), op, Arc::clone(value))) + as PhysicalExprRef + }) + .collect::>(); + balanced(&comparisons, combine) +} + fn build_predicate( expression: &PhysicalExprRef, schema: &SchemaRef, @@ -133,8 +154,12 @@ struct BenchmarkCase { schema: SchemaRef, in_list: PhysicalExprRef, expanded_or: PhysicalExprRef, + not_in_list: PhysicalExprRef, + expanded_and: PhysicalExprRef, in_list_predicate: PruningPredicate, expanded_or_predicate: PruningPredicate, + not_in_list_predicate: PruningPredicate, + expanded_and_predicate: PruningPredicate, statistics: IntervalStatistics, } @@ -149,42 +174,57 @@ impl BenchmarkCase { let values = (0..size) .map(|index| lit(ScalarValue::new_utf8view(value(index * 10)))) .collect::>(); + let not_in_list = + in_list(Arc::clone(&column), values.clone(), &true, &schema).unwrap(); let in_list = in_list(Arc::clone(&column), values.clone(), &false, &schema).unwrap(); - let equalities = values - .into_iter() - .map(|value| { - Arc::new(BinaryExpr::new(Arc::clone(&column), Operator::Eq, value)) - as PhysicalExprRef - }) - .collect::>(); - let expanded_or = balanced_or(&equalities); + let expanded_or = expanded(&column, &values, Operator::Eq, Operator::Or); + let expanded_and = expanded(&column, &values, Operator::NotEq, Operator::And); let in_list_predicate = build_predicate(&in_list, &schema, size); let expanded_or_predicate = build_predicate(&expanded_or, &schema, size); + let not_in_list_predicate = build_predicate(¬_in_list, &schema, size); + let expanded_and_predicate = build_predicate(&expanded_and, &schema, size); eprintln!( - "string_in_list_pruning: {size} values, compact={}", + "string_in_list_pruning: {size} values, compact in={}, compact not in={}", in_list_predicate .predicate_expr() .to_string() - .contains("IN_SET_INTERSECTS") + .contains("IN_SET_INTERSECTS"), + not_in_list_predicate + .predicate_expr() + .to_string() + .contains("NOT_IN_SET_MAY_MATCH") ); let statistics = IntervalStatistics::new(size); - // Check that both benchmark paths do the same useful work, rather + // Check that every benchmark path does the same useful work, rather // than comparing compact pruning with an always-true fallback. - let expected = (0..CONTAINERS) + let kept = (0..CONTAINERS) .map(|index| index % 2 == 0) .collect::>(); - assert_eq!(in_list_predicate.prune(&statistics).unwrap(), expected); - assert_eq!(expanded_or_predicate.prune(&statistics).unwrap(), expected); + let negated_kept = kept.iter().map(|keep| !keep).collect::>(); + assert_eq!(in_list_predicate.prune(&statistics).unwrap(), kept); + assert_eq!(expanded_or_predicate.prune(&statistics).unwrap(), kept); + assert_eq!( + not_in_list_predicate.prune(&statistics).unwrap(), + negated_kept + ); + assert_eq!( + expanded_and_predicate.prune(&statistics).unwrap(), + negated_kept + ); Self { size, schema, in_list, expanded_or, + not_in_list, + expanded_and, in_list_predicate, expanded_or_predicate, + not_in_list_predicate, + expanded_and_predicate, statistics, } } @@ -198,6 +238,8 @@ fn criterion_benchmark(criterion: &mut Criterion) { for (name, expression) in [ ("in_list", &case.in_list), ("expanded_or", &case.expanded_or), + ("not_in_list", &case.not_in_list), + ("expanded_and", &case.expanded_and), ] { construction.bench_with_input( BenchmarkId::new(name, case.size), @@ -222,6 +264,8 @@ fn criterion_benchmark(criterion: &mut Criterion) { for (name, predicate) in [ ("in_list", &case.in_list_predicate), ("expanded_or", &case.expanded_or_predicate), + ("not_in_list", &case.not_in_list_predicate), + ("expanded_and", &case.expanded_and_predicate), ] { evaluation.bench_with_input( BenchmarkId::new(name, case.size), diff --git a/datafusion/pruning/src/pruning_predicate.rs b/datafusion/pruning/src/pruning_predicate.rs index 4985cadac72b..7e1873e6eee6 100644 --- a/datafusion/pruning/src/pruning_predicate.rs +++ b/datafusion/pruning/src/pruning_predicate.rs @@ -22,7 +22,7 @@ use std::collections::HashSet; use std::sync::Arc; -use crate::string_in_list::StringInListPruningExpr; +use crate::string_in_list::{SetMembership, StringInListPruningExpr}; use arrow::array::AsArray; use arrow::{ @@ -450,7 +450,7 @@ impl<'a> PruningPredicateBuilder<'a> { /// | Condition | Pruning representation | /// | --- | --- | /// | `N <= min(20, C)` | Existing per-value rewrite | - /// | `20 < N <= C`, positive, non-null literal strings on a string column | Compact sorted domain | + /// | `20 < N <= C`, non-null literal strings on a string column | Compact sorted domain | /// | `20 < N <= C`, other lists | Existing per-value rewrite | /// | `N > C` | Unhandled-predicate hook, normally "keep the container" | /// @@ -459,8 +459,9 @@ impl<'a> PruningPredicateBuilder<'a> { /// pruning (such as Bloom filters). The default cap is [`MAX_IN_LIST_SIZE`] /// (20), so the compact path requires an explicitly raised cap. /// - /// Raising the cap can still build large comparison trees for non-string - /// lists, `NOT IN`, or lists containing NULL; their handling is unchanged. + /// The compact form covers `IN` and `NOT IN` alike. Raising the cap can + /// still build large comparison trees for non-string lists or lists + /// containing NULL; their handling is unchanged. /// /// Query engines typically pass /// `datafusion.execution.parquet.max_in_list_size` here. @@ -1470,15 +1471,27 @@ fn build_is_null_column_expr( } } -/// Keep large literal string domains compact instead of building an OR tree. +/// Keep large literal string domains compact instead of building a per-value +/// tree: an OR tree for `IN`, an AND chain for `NOT IN`. +/// +/// `IN` excludes a container whose interval is disjoint from the domain. That +/// matches the per-value OR tree except for inverted bounds, which the compact +/// form reports as UNKNOWN rather than excluding. +/// +/// `NOT IN` excludes a container only where `min` and `max` agree on a single +/// value the domain holds, which is exactly what makes the per-value +/// `min != v OR v != max` chain false. Its decisions match that chain +/// everywhere, including absent and inverted bounds. fn build_string_in_list_expr( in_list: &phys_expr::InListExpr, schema: &Schema, required_columns: &mut RequiredColumns, ) -> Option> { - if in_list.negated() { - return None; - } + let membership = if in_list.negated() { + SetMembership::NotIn + } else { + SetMembership::In + }; let column = in_list.expr().downcast_ref::()?; let field = schema.fields().get(column.index())?; let data_type = match field.data_type() { @@ -1490,6 +1503,7 @@ fn build_string_in_list_expr( } // NULLs must remain unhandled: the inverse predicate is also used to prove // that every row matches, and IN (..., NULL) can evaluate to UNKNOWN. + // Tracked in https://github.com/apache/datafusion/issues/24711. let values = in_list .list() .iter() @@ -1503,17 +1517,18 @@ fn build_string_in_list_expr( .ok()?; let non_null = build_is_null_column_expr(in_list.expr(), schema, required_columns, true)?; - let intersects = Arc::new(StringInListPruningExpr::new(min, max, values)); + let may_match = Arc::new(StringInListPruningExpr::new(membership, min, max, values)); Some(Arc::new(phys_expr::BinaryExpr::new( non_null, Operator::And, - intersects, + may_match, ))) } /// Default maximum number of entries in an `IN (...)` list eligible for -/// statistics pruning. Eligible positive literal string lists above this -/// threshold use a compact sorted domain instead of per-value min/max checks. +/// statistics pruning. Eligible literal string lists above this threshold use a +/// compact sorted domain instead of per-value min/max checks, for both `IN` and +/// `NOT IN`. /// Callers can raise the cap via [`PredicateRewriter::with_max_in_list_size`], and /// query engines can wire it from the /// `datafusion.execution.parquet.max_in_list_size` config option. @@ -1599,9 +1614,9 @@ impl PredicateRewriter { /// Returns the pruning predicate as an [`PhysicalExpr`] /// /// `max_in_list_size` is the largest `IN (...)` list eligible for statistics -/// pruning. Large positive literal string lists use a compact sorted domain; -/// other eligible lists use per-value checks. Longer lists fall back to -/// `unhandled_hook`. +/// pruning. Large literal string lists use a compact sorted domain, for both +/// `IN` and `NOT IN`; other eligible lists use per-value checks. Longer lists +/// fall back to `unhandled_hook`. fn build_predicate_expression( expr: &Arc, schema: &SchemaRef, @@ -1645,6 +1660,7 @@ fn build_predicate_expression( // Keep the existing expression shape for lists of at most 20 values. // This lower bound is a scope/compatibility choice, not a measured // performance threshold; compact pruning is opt-in via a raised cap. + // The compact form covers both `IN` and `NOT IN`. if in_list.list().len() > MAX_IN_LIST_SIZE && in_list.list().len() <= max_in_list_size && let Some(pruning_expr) = @@ -3883,33 +3899,43 @@ mod tests { fn large_string_in_list_respects_configured_limit() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)])); let values = (0..21).map(|i| lit(format!("a{i:03}"))).collect::>(); - let expr = logical2physical(&col("c1").in_list(values, false), &schema); - for (limit, compact) in [(0, false), (20, false), (21, true), (32, true)] { - let predicate = PruningPredicateBuilder::new() + for (negated, marker) in + [(false, "IN_SET_INTERSECTS"), (true, "NOT_IN_SET_MAY_MATCH")] + { + let expr = + logical2physical(&col("c1").in_list(values.clone(), negated), &schema); + + for (limit, compact) in [(0, false), (20, false), (21, true), (32, true)] { + let predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .with_max_in_list_size(limit) + .try_build(Arc::clone(&expr))?; + assert_eq!( + predicate.predicate_expr().to_string().contains(marker), + compact, + "negated={negated}, limit={limit}" + ); + assert_eq!( + is_always_true(predicate.predicate_expr()), + !compact, + "negated={negated}, limit={limit}" + ); + } + + let default = PruningPredicateBuilder::new() .with_file_schema(Arc::clone(&schema)) - .with_max_in_list_size(limit) - .try_build(Arc::clone(&expr))?; - assert_eq!( - predicate - .predicate_expr() - .to_string() - .contains("IN_SET_INTERSECTS"), - compact, - "limit={limit}" + .try_build(expr)?; + assert!( + is_always_true(default.predicate_expr()), + "negated={negated}" ); - assert_eq!(is_always_true(predicate.predicate_expr()), !compact); } - - let default = PruningPredicateBuilder::new() - .with_file_schema(schema) - .try_build(expr)?; - assert!(is_always_true(default.predicate_expr())); Ok(()) } #[test] - fn large_string_in_list_keeps_null_and_not_in_semantics() -> Result<()> { + fn large_string_in_list_keeps_null_semantics() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)])); let values = (0..21).map(|i| lit(format!("a{i:03}"))).collect::>(); let stats = TestStatistics::new().with( @@ -3928,10 +3954,9 @@ mod tests { )?; assert_eq!(predicate.prune(&stats)?, [true, false]); - let mut with_null = values.clone(); + let mut with_null = values; with_null.push(lit(ScalarValue::Utf8(None))); for expr in [ - col("c1").in_list(values, true), col("c1").in_list(with_null.clone(), false), col("c1").in_list(with_null.clone(), true), ] { @@ -3942,16 +3967,15 @@ mod tests { assert!(is_always_true(default.predicate_expr()), "{expr}"); assert_eq!(default.prune(&stats)?, [true, true]); - // Raising the cap retains the existing per-value rewrite for - // NOT IN and lists containing NULL; neither uses the new path. + // Raising the cap retains the existing per-value rewrite for lists + // containing NULL, in either direction. Dropping the NULL literal + // would turn UNKNOWN into FALSE, which the inverse-predicate proof + // in identify_fully_matched_row_groups cannot absorb. + // See https://github.com/apache/datafusion/issues/24711. let raised = large_string_pruning_predicate(physical, Arc::clone(&schema))?; - assert!( - !raised - .predicate_expr() - .to_string() - .contains("IN_SET_INTERSECTS"), - "{expr}" - ); + let raised = raised.predicate_expr().to_string(); + assert!(!raised.contains("IN_SET_INTERSECTS"), "{expr}"); + assert!(!raised.contains("NOT_IN_SET_MAY_MATCH"), "{expr}"); } // Inverting NOT IN (..., NULL) must not prove a full match and bypass @@ -3965,6 +3989,221 @@ mod tests { Ok(()) } + /// Statistics shared by the compact `NOT IN` tests, one row per case. + fn not_in_container_stats() -> TestStatistics { + TestStatistics::new().with( + "c1", + ContainerStats::new_utf8( + [ + Some("a005"), // single value inside the domain + Some("a005"), // same, with some NULL rows + Some("zzz"), // single value outside the domain + Some("a005"), // two adjacent domain values + Some("a000"), // spans the whole domain + None, // every row NULL + None, // no bounds at all + Some("z"), // inverted bounds + Some("a005"), // unbounded above + None, // unbounded below + ], + [ + Some("a005"), + Some("a005"), + Some("zzz"), + Some("a006"), + Some("a020"), + None, + None, + Some("m"), + None, + Some("a005"), + ], + ) + .with_null_counts([ + Some(0), + Some(3), + Some(0), + Some(0), + Some(0), + Some(10), + Some(0), + Some(0), + Some(0), + Some(0), + ]) + .with_row_counts([Some(10); 10]), + ) + } + + #[test] + fn large_string_not_in_list_prunes_only_single_valued_containers() -> Result<()> { + let types = [ + DataType::Utf8, + DataType::LargeUtf8, + DataType::Utf8View, + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + ]; + for data_type in &types { + let schema = + Arc::new(Schema::new(vec![Field::new("c1", data_type.clone(), true)])); + let values = (0..21) + .map(|i| { + let value = + ScalarValue::from(format!("a{i:03}")).cast_to(data_type)?; + Ok(Arc::new(phys_expr::Literal::new(value)) as PhysicalExprRef) + }) + .collect::>>()?; + let expr = phys_expr::in_list( + Arc::new(phys_expr::Column::new("c1", 0)), + values, + &true, + &schema, + )?; + let predicate = large_string_pruning_predicate(expr, schema)?; + assert_eq!( + predicate.predicate_expr().to_string(), + "c1_null_count@3 != row_count@2 AND NOT_IN_SET_MAY_MATCH(c1_min@0, c1_max@1, 21 values)", + "type={data_type:?}" + ); + + // Overlap never excludes a NOT IN container: only an interval pinned + // to one value the domain holds, or a container that is entirely NULL. + assert_eq!( + predicate.prune(¬_in_container_stats())?, + [ + false, false, true, true, true, false, true, true, true, true + ], + "type={data_type:?}" + ); + + // The expression and statistics schema do not grow with the domain. + assert_eq!(predicate.required_columns.columns.len(), 4); + let mut nodes = 0; + predicate.predicate_expr().apply(|_| { + nodes += 1; + Ok(TreeNodeRecursion::Continue) + })?; + assert_eq!(nodes, 7, "type={data_type:?}"); + } + Ok(()) + } + + /// `identify_fully_matched_row_groups` proves "every row matches" by showing + /// the pruning predicate for `NOT P OR IsNull(col)` excludes the container. + /// That inference needs `P` to be two-valued, so the compact `NOT IN` form + /// and the compact `IN` it inverts to must stay exact. Both sides use a + /// raised cap here, which is what wiring the configured cap into the + /// inverted builder, or removing the lower bound, would produce. + #[test] + fn large_string_not_in_list_inverts_without_false_full_match() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)])); + let values = (0..21).map(|i| lit(format!("a{i:03}"))).collect::>(); + let forward = large_string_pruning_predicate( + logical2physical(&col("c1").in_list(values.clone(), true), &schema), + Arc::clone(&schema), + )?; + // NOT(c1 NOT IN (...)) OR c1 IS NULL, the shape row_group_filter builds + // once PhysicalExprSimplifier turns NOT(NOT IN) back into IN. + let inverted = large_string_pruning_predicate( + logical2physical( + &col("c1").in_list(values, false).or(col("c1").is_null()), + &schema, + ), + schema, + )?; + let inverted_expr = inverted.predicate_expr().to_string(); + assert!( + inverted_expr.contains("IN_SET_INTERSECTS"), + "{inverted_expr}" + ); + + let stats = not_in_container_stats(); + let kept = forward.prune(&stats)?; + let fully_matched = inverted + .prune(&stats)? + .into_iter() + .map(|keep| !keep) + .collect::>(); + + // Only container 2 qualifies: every row holds the single value "zzz", + // which is outside the domain, so every row satisfies NOT IN. Container + // 0 repeats a domain member, and container 5 is entirely NULL, so + // neither may be claimed even though both are single-valued. + assert_eq!( + fully_matched, + [ + false, false, true, false, false, false, false, false, false, false + ] + ); + for (index, (keep, matched)) in kept.iter().zip(&fully_matched).enumerate() { + assert!(*keep || !matched, "container {index} claimed after pruning"); + } + Ok(()) + } + + #[test] + fn large_string_not_in_list_matches_per_value_chain() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)])); + let values = (0..21).map(|i| format!("a{i:03}")).collect::>(); + + let not_in = col("c1").in_list( + values.iter().map(|value| lit(value.clone())).collect(), + true, + ); + let compact = large_string_pruning_predicate( + logical2physical(¬_in, &schema), + Arc::clone(&schema), + )?; + + // The per-value rewrite this replaces, written out as the AND chain of + // `col != value` that build_predicate_expression would produce. + let chain = values + .iter() + .map(|value| col("c1").not_eq(lit(value.clone()))) + .reduce(Expr::and) + .unwrap(); + let per_value = large_string_pruning_predicate( + logical2physical(&chain, &schema), + Arc::clone(&schema), + )?; + let per_value_expr = per_value.predicate_expr().to_string(); + assert!(!per_value_expr.contains("NOT_IN_SET_MAY_MATCH")); + assert!(per_value_expr.contains("c1_min@0 != a000 OR a000 != c1_max@1")); + + // Containers 7 and 9 carry inverted and absent bounds, where the two + // forms could diverge without the compact side matching the chain. + let stats = not_in_container_stats(); + assert_eq!(compact.prune(&stats)?, per_value.prune(&stats)?); + + // The chain above must be what the per-value NOT IN rewrite actually + // builds, or the comparison checks the wrong baseline. Twenty values + // still take that path under the default cap, so both agree there. + let short = (0..20).map(|i| format!("a{i:03}")).collect::>(); + let rewritten = |expr| -> Result { + Ok(PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(logical2physical(expr, &schema))? + .predicate_expr() + .to_string()) + }; + assert_eq!( + rewritten( + &col("c1").in_list( + short.iter().map(|value| lit(value.clone())).collect(), + true + ) + )?, + rewritten( + &short + .iter() + .map(|value| col("c1").not_eq(lit(value.clone()))) + .reduce(Expr::and) + .unwrap() + )? + ); + Ok(()) + } + #[test] fn row_group_predicate_cast_int_int() -> Result<()> { let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]); diff --git a/datafusion/pruning/src/string_in_list.rs b/datafusion/pruning/src/string_in_list.rs index bc5c8fa908e7..eadebf216569 100644 --- a/datafusion/pruning/src/string_in_list.rs +++ b/datafusion/pruning/src/string_in_list.rs @@ -27,19 +27,35 @@ use datafusion_common::{Result, assert_eq_or_internal_err}; use datafusion_physical_expr::{PhysicalExpr, PhysicalExprRef}; use datafusion_physical_plan::ColumnarValue; -/// Tests whether a sorted string domain intersects an inclusive statistics interval. +/// Which `IN` form a sorted string domain is pruning for. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub(crate) enum SetMembership { + /// `col IN (...)`. A row matches only where the domain intersects the + /// interval, so a disjoint interval excludes every row. + In, + /// `col NOT IN (...)`. Overlap proves nothing here: values outside the + /// domain still satisfy the predicate. An interval excludes every row only + /// when it holds a single value that the domain contains. + NotIn, +} + +/// Tests an inclusive statistics interval against a sorted string domain. /// /// [`PhysicalExpr::evaluate`] returns one nullable Boolean per min/max interval: -/// * `true`: the interval intersects the domain, so matching rows may exist. -/// * `false`: the available bounds prove the interval disjoint from the domain. +/// * `true`: matching rows may exist, so the container must be read. +/// * `false`: the available bounds prove no row can match. /// * `NULL`: incomplete, invalid, or unusable bounds prevent a safe decision. /// -/// A single known bound can still prove disjointness. Otherwise, unknown results -/// keep the container eligible for reading. +/// [`SetMembership`] selects the test. For [`SetMembership::In`] a single known +/// bound can still prove disjointness. [`SetMembership::NotIn`] needs both +/// bounds and excludes only a single-valued interval, which is the same reach as +/// the per-value `min != v OR v != max` chain it replaces. Otherwise, unknown +/// results keep the container eligible for reading. /// /// This expression is used only for pruning; the original IN remains the row filter. #[derive(Debug, Eq)] pub(crate) struct StringInListPruningExpr { + membership: SetMembership, min: PhysicalExprRef, max: PhysicalExprRef, values: Arc<[String]>, @@ -47,6 +63,7 @@ pub(crate) struct StringInListPruningExpr { impl StringInListPruningExpr { pub(crate) fn new( + membership: SetMembership, min: PhysicalExprRef, max: PhysicalExprRef, mut values: Vec, @@ -54,21 +71,33 @@ impl StringInListPruningExpr { values.sort_unstable(); 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| candidate.as_bytes().cmp(value)) + .is_ok() + } } impl PartialEq for StringInListPruningExpr { fn eq(&self, other: &Self) -> bool { - self.min.eq(&other.min) && self.max.eq(&other.max) && self.values == other.values + self.membership == other.membership + && self.min.eq(&other.min) + && self.max.eq(&other.max) + && self.values == other.values } } impl Hash for StringInListPruningExpr { fn hash(&self, state: &mut H) { + self.membership.hash(state); self.min.hash(state); self.max.hash(state); self.values.hash(state); @@ -77,9 +106,13 @@ impl Hash for StringInListPruningExpr { impl Display for StringInListPruningExpr { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + let name = match self.membership { + SetMembership::In => "IN_SET_INTERSECTS", + SetMembership::NotIn => "NOT_IN_SET_MAY_MATCH", + }; write!( f, - "IN_SET_INTERSECTS({}, {}, {} values)", + "{name}({}, {}, {} values)", self.min, self.max, self.values.len() @@ -151,19 +184,46 @@ impl PhysicalExpr for StringInListPruningExpr { // 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)) + match self.membership { + SetMembership::In => { + 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 wider interval can always 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 too. This arm also + // compares for equality rather than order, so it does + // not rely on the bound ordering the IN arm needs. + SetMembership::NotIn => { + Some(min != max || !self.contains(min)) + } + } } // 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. + // An unbounded interval is never single-valued, so NOT IN takes + // neither arm. (Some(min), None) - if self.values.last().is_some_and(|v| v.as_bytes() < min) => + if self.membership == SetMembership::In + && 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) => + if self.membership == SetMembership::In + && self + .values + .first() + .is_some_and(|v| v.as_bytes() > max) => { Some(false) } @@ -184,6 +244,7 @@ impl PhysicalExpr for StringInListPruningExpr { ) -> Result { assert_eq_or_internal_err!(children.len(), 2); Ok(Arc::new(Self { + membership: self.membership, min: Arc::clone(&children[0]), max: Arc::clone(&children[1]), values: Arc::clone(&self.values), diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 757ddc9ffdf5..e895b6b8be1c 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -413,7 +413,7 @@ datafusion.execution.parquet.dictionary_page_size_limit 1048576 (writing) Sets b datafusion.execution.parquet.enable_page_index true (reading) If true, reads the Parquet data page level metadata (the Page Index), if present, to reduce the I/O and number of rows decoded. datafusion.execution.parquet.encoding NULL (writing) Sets default encoding for any column. Valid values are: plain, plain_dictionary, rle, bit_packed, delta_binary_packed, delta_length_byte_array, delta_byte_array, rle_dictionary, and byte_stream_split. These values are not case sensitive. If NULL, uses default parquet writer setting datafusion.execution.parquet.force_filter_selections false (reading) Force the use of RowSelections for filter results, when pushdown_filters is enabled. If false, the reader will automatically choose between a RowSelection and a Bitmap based on the number and pattern of selected rows. -datafusion.execution.parquet.max_in_list_size 20 Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger positive, non-null literal string lists on a string column use a compact sorted domain. Other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. +datafusion.execution.parquet.max_in_list_size 20 Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger non-null literal string lists on a string column use a compact sorted domain, for both `IN` and `NOT IN`. Other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. datafusion.execution.parquet.max_predicate_cache_size NULL (reading) The maximum predicate cache size, in bytes. When `pushdown_filters` is enabled, sets the maximum memory used to cache the results of predicate evaluation between filter evaluation and output generation. Decreasing this value will reduce memory usage, but may increase IO and CPU usage. None means use the default parquet reader setting. 0 means no caching. datafusion.execution.parquet.max_row_group_bytes NULL (writing) Target maximum size of each row group in bytes. When set, the writer flushes whenever either this limit or `max_row_group_size` is reached, whichever comes first. Useful for bounding writer memory on wide schemas where a row-count limit can map to very different byte sizes. Matches the behavior of `parquet.block.size` in parquet-mr. If `None` (the default), only the row-count limit applies. Currently only honored when `allow_single_file_parallelism` is `false`; by default the parallel file writer ignores this limit. datafusion.execution.parquet.max_row_group_size 1048576 (writing) Target maximum number of rows in each row group (defaults to 1M rows). Writing larger row groups requires more memory to write, but can get better compression and be faster to read. When `max_row_group_bytes` is also set, the writer flushes a row group when either limit is reached, whichever comes first. diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 7c21103cd781..313aa985ed85 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -93,7 +93,7 @@ The following configuration settings are available: | datafusion.execution.parquet.coerce_int96_tz | NULL | (reading) Optional timezone applied to INT96 columns when `coerce_int96` is set. When `Some`, INT96 columns coerce to `Timestamp(, Some())` instead of the default `Timestamp(, None)`. Spark and other systems write INT96 values as UTC-adjusted instants, so callers that need the resulting Arrow type to be timezone-aware (e.g. for Spark `TimestampType` semantics) should set this to `"UTC"`. No effect when `coerce_int96` is `None`. | | datafusion.execution.parquet.bloom_filter_on_read | true | (reading) Use any available bloom filters when reading parquet files | | datafusion.execution.parquet.max_predicate_cache_size | NULL | (reading) The maximum predicate cache size, in bytes. When `pushdown_filters` is enabled, sets the maximum memory used to cache the results of predicate evaluation between filter evaluation and output generation. Decreasing this value will reduce memory usage, but may increase IO and CPU usage. None means use the default parquet reader setting. 0 means no caching. | -| datafusion.execution.parquet.max_in_list_size | 20 | Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger positive, non-null literal string lists on a string column use a compact sorted domain. Other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. | +| datafusion.execution.parquet.max_in_list_size | 20 | Maximum number of input values in an `IN (...)` list eligible for min/max pruning. Lists above this cap, or a cap of 0, skip this rewrite; other predicates and Bloom-filter pruning remain available. Within the cap, nonempty lists of at most 20 values use the existing per-value rewrite. Larger non-null literal string lists on a string column use a compact sorted domain, for both `IN` and `NOT IN`. Other lists retain the existing per-value rewrite, so raising the cap can make those predicates expensive to build and evaluate. Defaults to 20. | | datafusion.execution.parquet.data_pagesize_limit | 1048576 | (writing) Sets best effort maximum size of data page in bytes | | datafusion.execution.parquet.write_batch_size | 1024 | (writing) Sets write_batch_size in rows | | datafusion.execution.parquet.writer_version | 1.0 | (writing) Sets parquet writer version valid values are "1.0" and "2.0" |