You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/19 01:37:31 UTC
(datafusion) 01/01: fix: preserve lambda bindings in nested-loop join projection pushdown (#26399)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-26399-fe612e2da212016905d600eb0dd689910c0899b7
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit ecdc9290ea7dec0242899c9b2c140c3eb29afba2
Author: xudong.w <wx...@gmail.com>
AuthorDate: Mon Oct 19 01:37:02 2026 +0000
fix: preserve lambda bindings in nested-loop join projection pushdown (#26399)
## Which issue does this PR close?
- Closes #26396.
## Rationale for this change
A nested-loop join whose lambda captures a column from the other input
can fail during physical optimization with `unable to unwrap lambda`,
even though the query executes correctly with physical optimization
disabled.
These queries should work with the default optimizer while retaining
projection pushdown for expressions that can be evaluated on one join
input.
## What changes are included in this PR?
- Keep each lambda and its body in the enclosing high-order function's
scope when extracting join-filter expressions. Ordinary sibling
expressions, other function arguments, and complete high-order
expressions that depend on only one input remain eligible for pushdown.
- Remap captured columns and rebase lambda parameter indices when moving
an expression to an input projection, extending the intermediate filter
schema, and pruning unused filter columns. Nested lambdas retain their
relative parameter offsets and bindings.
- Add `nested_loop_join_lambda.slt` with result and physical-plan
regressions.
This change preserves lambda bodies within their enclosing function; it
does not introduce hoisting of individual computations from inside a
lambda body.
## What is the testing strategy for this PR?
The new SQL test covers the original reproducer with either join-input
order, mixed predicates, nested lambdas, parameter shadowing, empty/NULL
arrays, and extraction of complete high-order functions onto either
input. Different input and filter schema widths exercise parameter
rebasing. Three `EXPLAIN` checks verify that ordinary expressions and
complete high-order functions are still pushed down.
## Are there any user-facing changes?
Nested-loop joins with lambda predicates now plan and execute with the
default physical optimizer. There are no public API or configuration
changes.
---
.../physical-optimizer/src/projection_pushdown.rs | 124 ++++++++-----
.../test_files/nested_loop_join_lambda.slt | 192 +++++++++++++++++++++
2 files changed, 275 insertions(+), 41 deletions(-)
diff --git a/datafusion/physical-optimizer/src/projection_pushdown.rs b/datafusion/physical-optimizer/src/projection_pushdown.rs
index f38c93a4c1..3280cd43f7 100644
--- a/datafusion/physical-optimizer/src/projection_pushdown.rs
+++ b/datafusion/physical-optimizer/src/projection_pushdown.rs
@@ -30,8 +30,8 @@ use datafusion_common::config::ConfigOptions;
use datafusion_common::tree_node::{
Transformed, TransformedResult, TreeNode, TreeNodeRecursion,
};
-use datafusion_common::{JoinSide, JoinType, Result};
-use datafusion_physical_expr::expressions::Column;
+use datafusion_common::{JoinSide, JoinType, Result, internal_datafusion_err};
+use datafusion_physical_expr::expressions::{Column, LambdaExpr, LambdaVariable};
use datafusion_physical_expr_common::physical_expr::{PhysicalExpr, is_volatile};
use datafusion_physical_plan::ExecutionPlan;
use datafusion_physical_plan::joins::NestedLoopJoinExec;
@@ -130,7 +130,7 @@ fn try_push_down_join_filter(
rhs_rewrite.data.1.column_indices(),
lhs_rewrite.data.0.schema().as_ref(),
rhs_rewrite.data.0.schema().as_ref(),
- );
+ )?;
let new_lhs_length = lhs_rewrite.data.0.schema().fields.len();
let projections = match projections.as_ref() {
@@ -219,8 +219,14 @@ fn try_push_down_projection(
})
.collect::<Fields>();
- let join_filter = JoinFilter::new(
+ let expression = remap_filter_expr(
new_expr.data,
+ join_filter.schema().fields().len(),
+ rewriter.intermediate_column_indices.len(),
+ |column| Ok(column.clone()),
+ )?;
+ let join_filter = JoinFilter::new(
+ expression,
rewriter.intermediate_column_indices,
Arc::new(Schema::new(intermediate_schema)),
);
@@ -240,15 +246,14 @@ fn minimize_join_filter(
old_column_indices: &[ColumnIndex],
lhs_schema: &Schema,
rhs_schema: &Schema,
-) -> JoinFilter {
+) -> Result<JoinFilter> {
let mut used_columns = HashSet::with_capacity(old_column_indices.len());
expr.apply(|expr| {
if let Some(col) = expr.downcast_ref::<Column>() {
used_columns.insert(col.index());
}
Ok(TreeNodeRecursion::Continue)
- })
- .expect("Closure cannot fail");
+ })?;
let new_column_indices = old_column_indices
.iter()
@@ -265,27 +270,24 @@ fn minimize_join_filter(
})
.collect::<Fields>();
- let final_expr = expr
- .transform_up(|expr| match expr.downcast_ref::<Column>() {
- None => Ok(Transformed::no(expr)),
- Some(column) => {
- let new_idx = used_columns
- .iter()
- .filter(|idx| **idx < column.index())
- .count();
- let new_column = Column::new(column.name(), new_idx);
- Ok(Transformed::yes(
- Arc::new(new_column) as Arc<dyn PhysicalExpr>
- ))
- }
- })
- .expect("Closure cannot fail");
+ let final_expr = remap_filter_expr(
+ expr,
+ old_column_indices.len(),
+ new_column_indices.len(),
+ |column| {
+ let new_idx = used_columns
+ .iter()
+ .filter(|idx| **idx < column.index())
+ .count();
+ Ok(Column::new(column.name(), new_idx))
+ },
+ )?;
- JoinFilter::new(
- final_expr.data,
+ Ok(JoinFilter::new(
+ final_expr,
new_column_indices,
Arc::new(Schema::new(fields)),
- )
+ ))
}
/// Implements the push-down machinery.
@@ -300,6 +302,7 @@ struct JoinFilterRewriter<'a> {
join_side_schema: &'a Schema,
join_side_projections: Vec<(Arc<dyn PhysicalExpr>, String)>,
intermediate_column_indices: Vec<ColumnIndex>,
+ original_filter_column_count: usize,
alias_generator: &'a AliasGenerator,
}
@@ -327,6 +330,7 @@ impl<'a> JoinFilterRewriter<'a> {
join_side,
join_side_schema,
join_side_projections: projections,
+ original_filter_column_count: column_indices.len(),
intermediate_column_indices: column_indices,
alias_generator,
}
@@ -339,6 +343,11 @@ impl<'a> JoinFilterRewriter<'a> {
&mut self,
expr: Arc<dyn PhysicalExpr>,
) -> Result<Transformed<Arc<dyn PhysicalExpr>>> {
+ // A lambda is evaluated by its enclosing high-order function. Keep its
+ // body in that scope; a complete high-order function may still be moved.
+ if expr.is::<LambdaExpr>() {
+ return Ok(Transformed::no(expr));
+ }
let depends_on_this_side = self.depends_on_join_side(&expr, self.join_side)?;
// We don't push down things that do not depend on this side (other side or no side).
if !depends_on_this_side {
@@ -378,22 +387,21 @@ impl<'a> JoinFilterRewriter<'a> {
// First, add a new projection. The expression must be rewritten, as it is no longer
// executed against the filter schema.
let new_idx = self.join_side_projections.len();
- let rewritten_expr = expr.transform_up(|expr| {
- Ok(match expr.downcast_ref::<Column>() {
- None => Transformed::no(expr),
- Some(column) => {
- let intermediate_column =
- &self.intermediate_column_indices[column.index()];
- assert_eq!(intermediate_column.side, self.join_side);
-
- let join_side_index = intermediate_column.index;
- let field = self.join_side_schema.field(join_side_index);
- let new_column = Column::new(field.name(), join_side_index);
- Transformed::yes(Arc::new(new_column) as Arc<dyn PhysicalExpr>)
- }
- })
- })?;
- self.join_side_projections.push((rewritten_expr.data, name));
+ let rewritten_expr = remap_filter_expr(
+ expr,
+ self.original_filter_column_count,
+ self.join_side_schema.fields().len(),
+ |column| {
+ let intermediate_column =
+ &self.intermediate_column_indices[column.index()];
+ assert_eq!(intermediate_column.side, self.join_side);
+
+ let join_side_index = intermediate_column.index;
+ let field = self.join_side_schema.field(join_side_index);
+ Ok(Column::new(field.name(), join_side_index))
+ },
+ )?;
+ self.join_side_projections.push((rewritten_expr, name));
// Then, update the column indices
let new_intermediate_idx = self.intermediate_column_indices.len();
@@ -420,6 +428,40 @@ impl<'a> JoinFilterRewriter<'a> {
}
}
+/// Remaps captured columns and rebases lambda parameters when the input changes.
+/// Parameters are appended after the input columns, including at nested lambda
+/// scopes. Moving a complete high-order expression preserves their relative
+/// offsets, so every parameter shifts by the change in the root input width.
+fn remap_filter_expr(
+ expr: Arc<dyn PhysicalExpr>,
+ old_input_len: usize,
+ new_input_len: usize,
+ map_column: impl Fn(&Column) -> Result<Column>,
+) -> Result<Arc<dyn PhysicalExpr>> {
+ expr.transform_up(|expr| {
+ if let Some(column) = expr.downcast_ref::<Column>() {
+ let new_column = map_column(column)?;
+ if new_column != *column {
+ return Ok(Transformed::yes(Arc::new(new_column) as _));
+ }
+ } else if let Some(variable) = expr.downcast_ref::<LambdaVariable>()
+ && old_input_len != new_input_len
+ {
+ let offset = variable.index().checked_sub(old_input_len).ok_or_else(|| {
+ internal_datafusion_err!(
+ "Lambda parameter {variable} precedes the {old_input_len} input columns"
+ )
+ })?;
+ return Ok(Transformed::yes(Arc::new(LambdaVariable::new(
+ new_input_len + offset,
+ Arc::clone(variable.field()),
+ )) as _));
+ }
+ Ok(Transformed::no(expr))
+ })
+ .data()
+}
+
#[cfg(test)]
mod test {
use super::*;
diff --git a/datafusion/sqllogictest/test_files/nested_loop_join_lambda.slt b/datafusion/sqllogictest/test_files/nested_loop_join_lambda.slt
new file mode 100644
index 0000000000..273d9db801
--- /dev/null
+++ b/datafusion/sqllogictest/test_files/nested_loop_join_lambda.slt
@@ -0,0 +1,192 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# Regression for https://github.com/apache/datafusion/issues/26396.
+statement ok
+SET datafusion.sql_parser.dialect = 'duckdb'
+
+statement ok
+SET datafusion.execution.target_partitions = 1
+
+statement ok
+SET datafusion.optimizer.join_reordering = false
+
+statement ok
+CREATE TABLE lambda_lists AS
+SELECT * FROM (VALUES
+ ([1, 30], 'a', 7, 2),
+ ([40], 'b', 8, 4),
+ ([], 'empty', 9, 2),
+ (NULL, 'null', 10, 2)
+) AS v(items, k, n, threshold)
+
+statement ok
+CREATE TABLE lambda_cutoffs AS
+SELECT * FROM (VALUES (25, -3), (5, -1)) AS v(cutoff, score)
+
+# The lambda captures a column from the opposite join input.
+query TI
+SELECT s.k, s.n
+FROM (SELECT 25 AS cutoff) c
+JOIN lambda_lists s ON array_any_match(s.items, x -> x < c.cutoff)
+----
+a 7
+
+query TI
+SELECT s.k, s.n
+FROM lambda_lists s
+JOIN (SELECT 25 AS cutoff) c ON array_any_match(s.items, x -> x < c.cutoff)
+----
+a 7
+
+# Extract ordinary expressions from both inputs while preserving the lambda.
+# Growing and then pruning the filter schema must also rebase its parameters.
+query ITI rowsort
+SELECT c.cutoff, s.k, s.n
+FROM lambda_cutoffs c JOIN lambda_lists s
+ON array_any_match(s.items, x -> x < c.cutoff)
+ AND abs(c.score) > abs(s.threshold)
+----
+25 a 7
+
+# The outer high-order function spans both inputs. Its array argument is a
+# complete single-input high-order function that can still be projected out.
+query IT rowsort
+SELECT c.cutoff, s.k
+FROM lambda_cutoffs c JOIN lambda_lists s
+ON array_any_match(array_transform(s.items, y -> y + 1), x -> x < c.cutoff)
+ AND abs(c.score) > abs(s.threshold)
+----
+25 a
+
+# Nested lambdas capture both a join column and the outer lambda's parameter.
+query IT rowsort
+SELECT c.cutoff, s.k
+FROM lambda_cutoffs c JOIN lambda_lists s
+ON array_any_match(s.items,
+ x -> array_any_match(s.items, y -> y <= x AND x < c.cutoff))
+ AND abs(c.score) > abs(s.threshold)
+----
+25 a
+
+# Shadowed names and an unused outer parameter must keep their own bindings.
+query IT rowsort
+SELECT c.cutoff, s.k
+FROM lambda_cutoffs c JOIN lambda_lists s
+ON array_any_match(s.items,
+ x -> array_any_match(s.items, x -> x < c.cutoff))
+ AND abs(c.score) > abs(s.threshold)
+----
+25 a
+
+# Projecting an entire high-order function is valid when its inputs and
+# captures all belong to one side. Selected payload columns widen that input
+# beyond the filter's schema, exercising rebasing into the input projection.
+query ITII rowsort
+SELECT c.cutoff, s.k, s.n, s.threshold
+FROM lambda_cutoffs c JOIN lambda_lists s
+ON CAST(array_any_match(s.items, x -> x < s.n) AS BIGINT) < abs(c.score)
+----
+25 a 7 2
+25 b 8 4
+25 empty 9 2
+5 b 8 4
+5 empty 9 2
+
+# The same single-side extraction with nested lambdas and an outer capture.
+query ITII rowsort
+SELECT c.cutoff, s.k, s.n, s.threshold
+FROM lambda_cutoffs c JOIN lambda_lists s
+ON CAST(array_any_match(s.items,
+ x -> array_any_match(s.items, y -> y <= x AND x < s.n)) AS BIGINT) < abs(c.score)
+----
+25 a 7 2
+25 b 8 4
+25 empty 9 2
+5 b 8 4
+5 empty 9 2
+
+# Extract the complete function on the left input as well.
+query ITII rowsort
+SELECT c.cutoff, s.k, s.n, s.threshold
+FROM lambda_lists s JOIN lambda_cutoffs c
+ON CAST(array_any_match(s.items, x -> x < s.n) AS BIGINT) < abs(c.score)
+----
+25 a 7 2
+25 b 8 4
+25 empty 9 2
+5 b 8 4
+5 empty 9 2
+
+statement ok
+SET datafusion.explain.physical_plan_only = true
+
+# Both abs expressions must still be pushed below the join.
+query TT
+EXPLAIN SELECT c.cutoff, s.k, s.n
+FROM lambda_cutoffs c JOIN lambda_lists s
+ON array_any_match(s.items, x -> x < c.cutoff)
+ AND abs(c.score) > abs(s.threshold)
+----
+physical_plan
+01)ProjectionExec: expr=[cutoff@0 as cutoff, k@3 as k, n@4 as n]
+02)--NestedLoopJoinExec: join_type=Inner, filter=array_any_match(items@1, (x) -> x@4 < cutoff@0) AND join_proj_push_down_1@2 > join_proj_push_down_2@3, projection=[cutoff@0, score@1, items@3, k@4, n@5, threshold@6]
+03)----ProjectionExec: expr=[cutoff@0 as cutoff, score@1 as score, abs(score@1) as join_proj_push_down_1]
+04)------DataSourceExec: partitions=1, partition_sizes=[1]
+05)----ProjectionExec: expr=[items@0 as items, k@1 as k, n@2 as n, threshold@3 as threshold, abs(threshold@3) as join_proj_push_down_2]
+06)------DataSourceExec: partitions=1, partition_sizes=[1]
+
+# A whole high-order function is still eligible for projection pushdown.
+query TT
+EXPLAIN SELECT c.cutoff, s.k, s.n, s.threshold
+FROM lambda_cutoffs c JOIN lambda_lists s
+ON CAST(array_any_match(s.items, x -> x < s.n) AS BIGINT) < abs(c.score)
+----
+physical_plan
+01)ProjectionExec: expr=[cutoff@0 as cutoff, k@3 as k, n@4 as n, threshold@5 as threshold]
+02)--NestedLoopJoinExec: join_type=Inner, filter=join_proj_push_down_2@1 < join_proj_push_down_1@0, projection=[cutoff@0, score@1, items@3, k@4, n@5, threshold@6]
+03)----ProjectionExec: expr=[cutoff@0 as cutoff, score@1 as score, abs(score@1) as join_proj_push_down_1]
+04)------DataSourceExec: partitions=1, partition_sizes=[1]
+05)----ProjectionExec: expr=[items@0 as items, k@1 as k, n@2 as n, threshold@3 as threshold, CAST(array_any_match(items@0, (x) -> x@4 < n@2) AS Int64) as join_proj_push_down_2]
+06)------DataSourceExec: partitions=1, partition_sizes=[1]
+
+# The array argument may be extracted without extracting either lambda body.
+query TT
+EXPLAIN SELECT c.cutoff, s.k
+FROM lambda_cutoffs c JOIN lambda_lists s
+ON array_any_match(array_transform(s.items, y -> y + 1), x -> x < c.cutoff)
+ AND abs(c.score) > abs(s.threshold)
+----
+physical_plan
+01)ProjectionExec: expr=[cutoff@0 as cutoff, k@3 as k]
+02)--NestedLoopJoinExec: join_type=Inner, filter=array_any_match(join_proj_push_down_2@2, (x) -> x@4 < cutoff@0) AND join_proj_push_down_1@1 > join_proj_push_down_3@3, projection=[cutoff@0, score@1, items@3, k@4, threshold@5]
+03)----ProjectionExec: expr=[cutoff@0 as cutoff, score@1 as score, abs(score@1) as join_proj_push_down_1]
+04)------DataSourceExec: partitions=1, partition_sizes=[1]
+05)----ProjectionExec: expr=[items@0 as items, k@1 as k, threshold@2 as threshold, array_transform(items@0, (y) -> y@3 + 1) as join_proj_push_down_2, abs(threshold@2) as join_proj_push_down_3]
+06)------DataSourceExec: partitions=1, partition_sizes=[1]
+
+statement ok
+RESET datafusion.explain.physical_plan_only
+
+statement ok
+SET datafusion.execution.target_partitions = 4
+
+statement ok
+RESET datafusion.optimizer.join_reordering
+
+statement ok
+RESET datafusion.sql_parser.dialect
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org