diff --git a/datafusion/pruning/benches/string_in_list_pruning.rs b/datafusion/pruning/benches/string_in_list_pruning.rs index a55d3dad6af12..16ce752ca3417 100644 --- a/datafusion/pruning/benches/string_in_list_pruning.rs +++ b/datafusion/pruning/benches/string_in_list_pruning.rs @@ -47,12 +47,10 @@ use datafusion_common::{Column, ScalarValue}; use datafusion_expr_common::operator::Operator; use datafusion_physical_expr::PhysicalExprRef; use datafusion_physical_expr::expressions::{BinaryExpr, col, in_list, lit}; -use datafusion_pruning::{ - MAX_IN_LIST_SIZE, PruningPredicate, PruningPredicateBuilder, PruningStatistics, -}; +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]; +const CONTAINER_COUNTS: [usize; 4] = [1, 16, 256, 4096]; const BASELINE_CONTAINER_COUNT: usize = 4096; fn value(index: usize) -> String { @@ -332,7 +330,12 @@ fn assert_equivalent_results(case: &BenchmarkCase, statistics: &IntervalStatisti .not_in_list_with_null_predicate .prune(statistics) .unwrap(); - if case.size + 1 > MAX_IN_LIST_SIZE { + if case + .in_list_with_null_predicate + .predicate_expr() + .to_string() + .contains("IN_SET_INTERSECTS") + { // Compact pruning applies filter semantics: NULL does not change which // rows can satisfy IN, while NOT IN (..., NULL) can never be true. assert_eq!(in_with_null, expected); @@ -435,11 +438,11 @@ fn criterion_benchmark(criterion: &mut Criterion) { criterion.benchmark_group("string_in_list_pruning/not_in_distributions"); distributions.throughput(Throughput::Elements(BASELINE_CONTAINER_COUNT as u64)); let singleton = IntervalStatistics::uniform_singleton(); - for case in cases.iter().filter(|case| case.size > 20) { - let compact = case.not_in_list_predicate.prune(&singleton).unwrap(); - assert!(compact.iter().all(|keep| !keep)); + for case in &cases { + let actual = case.not_in_list_predicate.prune(&singleton).unwrap(); + assert!(actual.iter().all(|keep| !keep)); assert_eq!( - compact, + actual, case.expanded_and_predicate.prune(&singleton).unwrap() ); for (name, predicate) in [ diff --git a/datafusion/pruning/src/pruning_predicate.rs b/datafusion/pruning/src/pruning_predicate.rs index b49b72058e0cd..efd1c85f2741b 100644 --- a/datafusion/pruning/src/pruning_predicate.rs +++ b/datafusion/pruning/src/pruning_predicate.rs @@ -1636,14 +1636,18 @@ fn build_compact_in_list_expr( } /// Default maximum number of entries in an `IN (...)` list eligible for -/// statistics pruning. Eligible literal lists of supported ordered types above -/// this threshold use a compact sorted domain instead of per-value min/max -/// checks, for both `IN` and `NOT IN`. +/// statistics pruning. See [`PruningPredicateBuilder::with_max_in_list_size`] +/// for the representation used within this cap. /// Callers can raise the cap via [`PredicateRewriter::with_max_in_list_size`], and /// query engines can wire it from the /// `datafusion.execution.parquet.max_in_list_size` config option. pub const MAX_IN_LIST_SIZE: usize = 20; +// Keep the representation threshold independent of the configurable pruning cap. +// Small lists can be faster with vectorized per-value comparisons than with +// compact per-container searches, particularly for large statistics batches. +const MIN_COMPACT_IN_LIST_SIZE: usize = 21; + /// Rewrite a predicate expression in terms of statistics (min/max/null_counts) /// for use as a [`PruningPredicate`]. pub struct PredicateRewriter { @@ -1770,11 +1774,8 @@ fn build_predicate_expression( } } if let Some(in_list) = expr.downcast_ref::() { - // Keep the existing expression shape for lists of at most 20 values. - // This lower bound is a scope/compatibility choice, not a measured - // performance threshold; compact pruning is opt-in via a raised cap. // The compact form covers both `IN` and `NOT IN`. - if in_list.list().len() > MAX_IN_LIST_SIZE + if in_list.list().len() >= MIN_COMPACT_IN_LIST_SIZE && in_list.list().len() <= max_in_list_size && let Some(pruning_expr) = build_compact_in_list_expr(in_list, schema, required_columns, properties) @@ -4247,7 +4248,7 @@ mod tests { predicate.required_columns().single_column().unwrap().name(), "c1" ); - if count > MAX_IN_LIST_SIZE { + if count >= MIN_COMPACT_IN_LIST_SIZE { // The expression and statistics schema do not grow with the domain. assert_eq!(predicate.required_columns.columns.len(), 4); let mut nodes = 0; @@ -4385,6 +4386,41 @@ mod tests { Ok(()) } + #[test] + fn string_in_list_representation_respects_cap_and_threshold() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)])); + for count in [1, 2, 4, 8, 20, 21] { + let values = (0..count) + .map(|i| lit(format!("a{i:03}"))) + .collect::>(); + for (negated, marker) in + [(false, "IN_SET_INTERSECTS"), (true, "NOT_IN_SET_MAY_MATCH")] + { + let expr = logical2physical( + &col("c1").in_list(values.clone(), negated), + &schema, + ); + for limit in [0, count - 1, count, 32] { + 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), + count >= 21 && count <= limit, + "count={count}, limit={limit}, negated={negated}" + ); + assert_eq!( + is_always_true(predicate.predicate_expr()), + count > limit, + "count={count}, limit={limit}, negated={negated}" + ); + } + } + } + Ok(()) + } + #[test] fn large_string_in_list_respects_configured_limit() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)]));