refactor: separate compact IN-list pruning threshold from the default cap

Author: goutamadwant

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 [

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::<phys_expr::InListExpr>() {
-        // 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::<Vec<_>>();
+            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)]));