perf: mirror ASOF equality-key filters to right

benchmarks/sql_benchmarks/asof_join/asof_join.suite

@@ -5,6 +5,7 @@
 # - latest right row at or before, or earliest right row at or after, each left row
 # - narrow integer or wide UTF-8 payloads
 # - selective left-side predicates that can run before matching
+# - equality-key predicates that can also prune the right input
 
 description = "ASOF join SQL benchmarks"
 
@@ -24,3 +25,7 @@ description = "Run the pre-sorted keyed ASOF join query."
 [[examples]]
 command = "cargo run --release --bin benchmark_runner -- asof_join --query 8"
 description = "Run the selective left-filter ASOF join query."
+
+[[examples]]
+command = "cargo run --release --bin benchmark_runner -- asof_join --query 9"
+description = "Run the equality-key filter ASOF join query."

benchmarks/sql_benchmarks/asof_join/benchmarks/q09.benchmark

@@ -0,0 +1,25 @@
+name Q09
+group asof_join
+
+expect_plan AsOfJoinExec
+
+run
+-- Both sides: 1M rows in 10K equality groups. The left predicate retains
+-- one group; mirroring it prunes the right input before sorting and broadcast.
+WITH left_input AS (
+    SELECT value % 10000 AS key,
+           value / 10000 + 1 AS ts,
+           value AS payload
+    FROM range(1000000)
+),
+right_input AS (
+    SELECT value % 10000 AS key,
+           value / 10000 AS ts,
+           value AS payload
+    FROM range(1000000)
+)
+SELECT l.key, l.ts, l.payload, r.payload AS right_payload
+FROM left_input l
+ASOF JOIN right_input r MATCH_CONDITION (l.ts >= r.ts)
+ON l.key = r.key
+WHERE l.key = 42;

datafusion/optimizer/src/push_down_filter.rs

@@ -40,7 +40,7 @@ use datafusion_expr::utils::{
     conjunction, expr_to_columns, split_conjunction, split_conjunction_owned,
 };
 use datafusion_expr::{
-    BinaryExpr, Distinct, Expr, Filter, Operator, Projection,
+    BinaryExpr, Distinct, Expr, ExprSchemable, Filter, Operator, Projection,
     TableProviderFilterPushDown, and, or,
 };
 
@@ -1156,6 +1156,48 @@ impl OptimizerRule for PushDownFilter {
                                 })
                         });
 
+                // A literal comparison on an equal, same-typed key has the
+                // same value for every matching pair. Mirroring it to the
+                // right can prune groups without changing the ASOF candidate.
+                let mut right_predicates = Vec::new();
+                for predicate in &push_predicates {
+                    let Expr::BinaryExpr(BinaryExpr {
+                        left,
+                        op: Operator::Eq,
+                        right,
+                    }) = predicate
+                    else {
+                        continue;
+                    };
+                    let ((Expr::Column(left_column), Expr::Literal(_, _))
+                    | (Expr::Literal(_, _), Expr::Column(left_column))) =
+                        (left.as_ref(), right.as_ref())
+                    else {
+                        continue;
+                    };
+                    for (left_key, right_key) in &join.on {
+                        let (Some(left_key_column), Some(right_key_column)) =
+                            (left_key.try_as_col(), right_key.try_as_col())
+                        else {
+                            continue;
+                        };
+                        if left_column == left_key_column
+                            && left_key.get_type(join.left.schema())?
+                                == right_key.get_type(join.right.schema())?
+                        {
+                            let replacements =
+                                HashMap::from([(left_key_column, right_key_column)]);
+                            right_predicates
+                                .push(replace_col(predicate.clone(), &replacements)?);
+                            break;
+                        }
+                    }
+                }
+                if let Some(predicate) = conjunction(right_predicates) {
+                    join.right =
+                        Arc::new(LogicalPlan::Filter(Filter::new(predicate, join.right)));
+                }
+
                 let result = if let Some(predicate) = conjunction(push_predicates) {
                     filter.predicate = predicate;
                     filter.input = join.left;

datafusion/sqllogictest/test_files/asof_join.slt

@@ -342,6 +342,77 @@ physical_plan
 07)----SortExec: expr=[grp@0 ASC, ts@1 ASC], preserve_partitioning=[false]
 08)------DataSourceExec: partitions=1, partition_sizes=[1]
 
+# An equality-key predicate on the left can also remove right-side groups
+# without changing matches or the unmatched left rows.
+query IT
+SELECT l.id, r.val
+FROM asof_left l
+ASOF JOIN asof_right r
+MATCH_CONDITION (l.ts >= r.ts)
+ON l.grp = r.grp
+WHERE l.grp = 'A'
+ORDER BY l.id;
+----
+1 NULL
+2 a4
+3 a6
+7 NULL
+
+query IT
+SELECT l.id, r.val
+FROM asof_left l
+ASOF JOIN asof_right r
+MATCH_CONDITION (l.ts >= r.ts)
+ON l.grp = r.grp
+WHERE l.grp = 'A' AND r.val IS NULL
+ORDER BY l.id;
+----
+1 NULL
+7 NULL
+
+# The disjunction also retains group B, so it cannot prune the right to A.
+query IT
+SELECT l.id, r.val
+FROM asof_left l
+ASOF JOIN asof_right r
+MATCH_CONDITION (l.ts >= r.ts)
+ON l.grp = r.grp
+WHERE l.grp = 'A' OR l.id = 4
+ORDER BY l.id;
+----
+1 NULL
+2 a4
+3 a6
+4 b1
+7 NULL
+
+query TT
+EXPLAIN SELECT l.id, r.val
+FROM asof_left l
+ASOF JOIN asof_right r
+MATCH_CONDITION (l.ts >= r.ts)
+ON l.grp = r.grp
+WHERE l.grp = 'A';
+----
+logical_plan
+01)Projection: l.id, r.val
+02)--AsOf Join: match=[l.ts >= r.ts], constraint=On, on=[l.grp = r.grp]
+03)----SubqueryAlias: l
+04)------Filter: asof_left.grp = Utf8View("A")
+05)--------TableScan: asof_left projection=[id, grp, ts]
+06)----SubqueryAlias: r
+07)------Filter: asof_right.grp = Utf8View("A")
+08)--------TableScan: asof_right projection=[grp, ts, val]
+physical_plan
+01)AsOfJoinExec: on=[(grp = grp)], match=[ts >= ts], projection=[id@0, val@5]
+02)--SortExec: expr=[ts@2 ASC], preserve_partitioning=[true]
+03)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
+04)------FilterExec: grp@1 = A
+05)--------DataSourceExec: partitions=1, partition_sizes=[1]
+06)--SortExec: expr=[ts@1 ASC], preserve_partitioning=[false]
+07)----FilterExec: grp@0 = A
+08)------DataSourceExec: partitions=1, partition_sizes=[1]
+
 # An ASOF join emits each left row exactly once, so the primary key from
 # asof_left remains unique. Because the outer LEFT JOIN does not use columns
 # from the ASOF side, it can be removed entirely.