From 85a5fe3dc1ed72b3ca9ec537c0ea2866f05e5106 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Mon, 10 Aug 2026 11:59:04 -0400 Subject: [PATCH 1/2] feat(pruning): prune containers through `CASE` predicates `build_predicate_expression` had no arm for `CaseExpr`, so any predicate whose top level (or a conjunct of it) is a `CASE` was handed to the unhandled hook and rewritten to `true`, contributing nothing to pruning. This shows up for dynamic filters pushed down from hash partitioned joins: the filter arrives at the scan wrapped in a partition switch, one arm per partition, e.g. CASE hash(l_orderkey) % 4 WHEN 0 THEN l_orderkey >= 165 AND l_orderkey <= 5999783 AND ... WHEN 1 THEN ... ELSE false END A container (row group, file, ...) may hold rows from any partition, so the container level predicate is the disjunction of the arms' `THEN` predicates. The `WHEN` values are not expressible in terms of container statistics and are dropped, which only weakens the predicate and is therefore sound. Arms that can never be `true` (`false` or `NULL` literals, and the implicit `NULL` of a missing `ELSE`) drop out of the disjunction; an arm that can not be rewritten becomes `true` via the unhandled hook and so makes the whole `CASE` `true` (no pruning), never `false`. Co-Authored-By: Claude Opus 5 --- datafusion/pruning/src/pruning_predicate.rs | 100 ++++++++++++++++++++ 1 file changed, 100 insertions(+) diff --git a/datafusion/pruning/src/pruning_predicate.rs b/datafusion/pruning/src/pruning_predicate.rs index 3a63451495e4c..37dc6c5e0060b 100644 --- a/datafusion/pruning/src/pruning_predicate.rs +++ b/datafusion/pruning/src/pruning_predicate.rs @@ -805,6 +805,19 @@ fn is_always_false(expr: &Arc) -> bool { .unwrap_or_default() } +/// Returns true if `expr` is a literal that can never evaluate to `true`, +/// i.e. either `false` or `NULL`. +/// +/// A filter only keeps rows for which the predicate evaluates to `true`, so a +/// `NULL` literal selects no rows, exactly like `false`. +fn is_never_true(expr: &Arc) -> bool { + expr.downcast_ref::() + .map(|l| { + matches!(l.value(), ScalarValue::Boolean(Some(false))) || l.value().is_null() + }) + .unwrap_or_default() +} + /// Describes which columns statistics are necessary to evaluate a /// [`PruningPredicate`]. /// @@ -1621,6 +1634,15 @@ fn build_predicate_expression( return unhandled_hook.handle(expr); } } + if let Some(case) = expr.downcast_ref::() { + return build_case_predicate_expression( + case, + schema, + required_columns, + unhandled_hook, + max_in_list_size, + ); + } let (left, op, right) = { if let Some(bin_expr) = expr.downcast_ref::() { @@ -1711,6 +1733,84 @@ fn build_predicate_expression( .unwrap_or_else(|_| unhandled_hook.handle(expr)) } +/// Rewrite a `CASE` expression used as a predicate into a container level +/// pruning predicate. +/// +/// A container (row group, file, partition, ...) holds many rows and a +/// container can only be pruned if *no* row in it can pass the predicate. Each +/// row picks exactly one arm of the `CASE`, but different rows in the same +/// container may pick different arms, so the container level predicate is the +/// **disjunction** of the arms' `THEN` predicates: +/// +/// ```text +/// CASE [expr] WHEN v1 THEN t1 WHEN v2 THEN t2 ELSE e END +/// => prune(t1) OR prune(t2) OR prune(e) +/// ``` +/// +/// This is what a dynamic filter pushed down from a hash partitioned join +/// looks like: one arm per partition, keyed on the partition number, e.g. +/// `CASE hash(a) % 4 WHEN 0 THEN a >= 1 AND a <= 9 WHEN 1 THEN ... END`. The +/// `WHEN` values are not expressible in terms of container statistics and are +/// deliberately ignored: dropping them only makes the predicate weaker (more +/// containers are kept), which is always sound. +/// +/// Notes: +/// * A missing `ELSE` is an implicit `NULL`, which never passes a filter, so it +/// contributes nothing to the disjunction (it must *not* be treated as +/// `true`). +/// * An arm that cannot be rewritten becomes `true` via `unhandled_hook`, which +/// makes the whole disjunction `true` (no pruning). That is sound. +/// * An arm that is statically `false` (empty partitions are emitted as +/// `lit(false)`) drops out of the disjunction. +/// * If every arm drops out the result is `false`: no row of any container can +/// pass, so all containers can be pruned. +fn build_case_predicate_expression( + case: &phys_expr::CaseExpr, + schema: &SchemaRef, + required_columns: &mut RequiredColumns, + unhandled_hook: &Arc, + max_in_list_size: usize, +) -> Arc { + let arms = case + .when_then_expr() + .iter() + .map(|(_when, then)| then) + .chain(case.else_expr()); + + let mut result: Option> = None; + for arm in arms { + // Arms that can never evaluate to `true` select no rows at all and so + // contribute nothing to the disjunction. + if is_never_true(arm) { + continue; + } + let arm_expr = build_predicate_expression( + arm, + schema, + required_columns, + unhandled_hook, + max_in_list_size, + ); + if is_always_true(&arm_expr) { + // `x OR true` is `true`: this container can never be pruned. + return arm_expr; + } + if is_always_false(&arm_expr) { + continue; + } + result = Some(match result { + None => arm_expr, + Some(prev) => { + Arc::new(phys_expr::BinaryExpr::new(prev, Operator::Or, arm_expr)) + } + }); + } + + result.unwrap_or_else(|| { + Arc::new(phys_expr::Literal::new(ScalarValue::Boolean(Some(false)))) + }) +} + /// Count of distinct column references in an expression. /// This is the same as [`collect_columns`] but optimized to stop counting /// once more than one distinct column is found. From db3fd52282f56d0df875f3b0bf2c3db216cc17f8 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Mon, 10 Aug 2026 11:59:18 -0400 Subject: [PATCH 2/2] test(pruning): cover `CASE` predicates in the pruning rewriter Covers both `CASE` shapes (with and without a base expression), a missing/NULL `ELSE`, `false` arms dropping out of the disjunction, an all-`false` `CASE` pruning everything, an unhandled arm degrading to `true`, and an end to end prune with the partitioned-ranges shape a dynamic filter from a hash partitioned join produces. Co-Authored-By: Claude Opus 5 --- datafusion/pruning/src/pruning_predicate.rs | 187 +++++++++++++++++++- 1 file changed, 186 insertions(+), 1 deletion(-) diff --git a/datafusion/pruning/src/pruning_predicate.rs b/datafusion/pruning/src/pruning_predicate.rs index 37dc6c5e0060b..f47eaf7b8951f 100644 --- a/datafusion/pruning/src/pruning_predicate.rs +++ b/datafusion/pruning/src/pruning_predicate.rs @@ -2281,7 +2281,7 @@ mod tests { use super::*; use datafusion_common::test_util::batches_to_string; - use datafusion_expr::{and, col, lit, or}; + use datafusion_expr::{and, case, col, lit, or, when}; use datafusion_physical_expr::utils::collect_columns; use insta::assert_snapshot; @@ -6161,4 +6161,189 @@ mod tests { "c1_null_count@2 != row_count@3 AND c1_min@0 <= a AND a <= c1_max@1"; assert_eq!(res.to_string(), expected); } + + /// Schema used by the `CASE` tests: `c1` is the column the arms filter on, + /// `c2` only appears in the `WHEN` conditions (like a partition selector) + fn case_test_schema() -> Schema { + Schema::new(vec![ + Field::new("c1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ]) + } + + /// The distinct source columns whose statistics the predicate requires + fn required_column_names(required_columns: &RequiredColumns) -> Vec { + required_columns + .iter() + .map(|(c, _t, _f)| c.name().to_string()) + .unique() + .collect() + } + + /// `CASE WHEN .. THEN .. ELSE .. END`: the container predicate is the + /// disjunction of the arms + #[test] + fn test_build_predicate_expression_case_with_base_expr() { + let schema = case_test_schema(); + // CASE c2 % 2 WHEN 0 THEN c1 < 10 WHEN 1 THEN c1 > 100 ELSE false END + let expr = case(col("c2").rem(lit(2))) + .when(lit(0), col("c1").lt(lit(10))) + .when(lit(1), col("c1").gt(lit(100))) + .otherwise(lit(ScalarValue::Boolean(Some(false)))) + .unwrap(); + let mut required_columns = RequiredColumns::new(); + let res = test_build_predicate_expression(&expr, &schema, &mut required_columns); + assert_eq!( + res.to_string(), + "c1_null_count@1 != row_count@2 AND c1_min@0 < 10 \ + OR c1_null_count@1 != row_count@2 AND c1_max@3 > 100" + ); + // the `WHEN` values are not prunable and must not require statistics + assert_eq!(required_column_names(&required_columns), vec!["c1"]); + } + + /// `CASE WHEN THEN .. END`: same, for the form without a base + /// expression + #[test] + fn test_build_predicate_expression_case_without_base_expr() { + let schema = case_test_schema(); + // CASE WHEN c2 = 0 THEN c1 < 10 WHEN c2 = 1 THEN c1 > 100 ELSE false END + let expr = when(col("c2").eq(lit(0)), col("c1").lt(lit(10))) + .when(col("c2").eq(lit(1)), col("c1").gt(lit(100))) + .otherwise(lit(ScalarValue::Boolean(Some(false)))) + .unwrap(); + let mut required_columns = RequiredColumns::new(); + let res = test_build_predicate_expression(&expr, &schema, &mut required_columns); + assert_eq!( + res.to_string(), + "c1_null_count@1 != row_count@2 AND c1_min@0 < 10 \ + OR c1_null_count@1 != row_count@2 AND c1_max@3 > 100" + ); + assert_eq!(required_column_names(&required_columns), vec!["c1"]); + } + + /// A `CASE` without an `ELSE` has an implicit `NULL` else, which selects no + /// rows and so must not weaken the predicate to `true` + #[test] + fn test_build_predicate_expression_case_without_else() { + let schema = case_test_schema(); + // CASE c2 WHEN 0 THEN c1 < 10 END + let expr = case(col("c2")) + .when(lit(0), col("c1").lt(lit(10))) + .end() + .unwrap(); + let res = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!( + res.to_string(), + "c1_null_count@1 != row_count@2 AND c1_min@0 < 10" + ); + + // an explicit NULL else behaves the same way + let expr = case(col("c2")) + .when(lit(0), col("c1").lt(lit(10))) + .otherwise(lit(ScalarValue::Boolean(None))) + .unwrap(); + let res = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!( + res.to_string(), + "c1_null_count@1 != row_count@2 AND c1_min@0 < 10" + ); + } + + /// `lit(false)` arms (e.g. an empty partition) drop out of the disjunction + #[test] + fn test_build_predicate_expression_case_with_false_arms() { + let schema = case_test_schema(); + // CASE c2 WHEN 0 THEN false WHEN 1 THEN c1 < 10 ELSE false END + let expr = case(col("c2")) + .when(lit(0), lit(ScalarValue::Boolean(Some(false)))) + .when(lit(1), col("c1").lt(lit(10))) + .otherwise(lit(ScalarValue::Boolean(Some(false)))) + .unwrap(); + let res = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!( + res.to_string(), + "c1_null_count@1 != row_count@2 AND c1_min@0 < 10" + ); + } + + /// If every arm is `false` no row can pass the predicate, so every + /// container can be pruned + #[test] + fn test_build_predicate_expression_case_all_false_arms() { + let schema = case_test_schema(); + let expr = case(col("c2")) + .when(lit(0), lit(ScalarValue::Boolean(Some(false)))) + .otherwise(lit(ScalarValue::Boolean(Some(false)))) + .unwrap(); + let mut required_columns = RequiredColumns::new(); + let res = test_build_predicate_expression(&expr, &schema, &mut required_columns); + assert_eq!(res.to_string(), "false"); + assert!(required_column_names(&required_columns).is_empty()); + } + + /// An arm that can not be rewritten degrades the whole `CASE` to `true` + /// (i.e. no pruning), never to `false` + #[test] + fn test_build_predicate_expression_case_with_unhandled_arm() { + let schema = case_test_schema(); + // CASE c2 WHEN 0 THEN c1 < 10 WHEN 1 THEN array_has([1], c1) END + let expr = case(col("c2")) + .when(lit(0), col("c1").lt(lit(10))) + .when(lit(1), array_has(make_array(vec![lit(1)]), col("c1"))) + .end() + .unwrap(); + let res = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(res.to_string(), "true"); + + // ... and the same when the unhandled arm comes first + let expr = case(col("c2")) + .when(lit(0), array_has(make_array(vec![lit(1)]), col("c1"))) + .when(lit(1), col("c1").lt(lit(10))) + .end() + .unwrap(); + let res = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(res.to_string(), "true"); + } + + /// End to end pruning with a `CASE` predicate of the shape produced by a + /// dynamic filter pushed down from a hash partitioned join + #[test] + fn prune_case_partitioned_ranges() { + let schema = Arc::new(Schema::new(vec![ + Field::new("c1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ])); + // CASE c2 % 2 + // WHEN 0 THEN c1 >= 0 AND c1 <= 10 + // WHEN 1 THEN c1 >= 100 AND c1 <= 110 + // ELSE false + // END + let expr = case(col("c2").rem(lit(2))) + .when( + lit(0), + col("c1").gt_eq(lit(0)).and(col("c1").lt_eq(lit(10))), + ) + .when( + lit(1), + col("c1").gt_eq(lit(100)).and(col("c1").lt_eq(lit(110))), + ) + .otherwise(lit(ScalarValue::Boolean(Some(false)))) + .unwrap(); + + let statistics = TestStatistics::new().with( + "c1", + ContainerStats::new_i32( + vec![Some(0), Some(50), Some(100), Some(200)], // min + vec![Some(5), Some(60), Some(105), Some(300)], // max + ), + ); + // only the container that overlaps neither range is pruned + prune_with_expr(expr, &schema, &statistics, &[true, false, true, false]); + } }