Skip to content

feat: add scoped physical expression traversal#21386

Open
rluvaton wants to merge 17 commits intoapache:mainfrom
rluvaton:add-support-for-scoped-traversal
Open

feat: add scoped physical expression traversal#21386
rluvaton wants to merge 17 commits intoapache:mainfrom
rluvaton:add-support-for-scoped-traversal

Conversation

@rluvaton
Copy link
Copy Markdown
Member

@rluvaton rluvaton commented Apr 5, 2026

Mostly tests

Which issue does this PR close?

Relates to #21231

Rationale for this change

CaseWhen (and other expression utilities) discover input columns by traversing the expression tree via children(). This breaks when a child expression introduces a new naming scope — for example, a lambda expression like array_transform(list, item -> item + external_col) where item is bound to a different schema than the outer expression.

Today, the only workaround is to not return scoped children from children(), but that breaks other legitimate use cases (e.g. finding all expressions in a tree, wrapping every expression with debug/timing, etc.). We need a way to distinguish "same-scope children" from "all children".

What changes are included in this PR?

New ScopedTreeNode trait (datafusion/common/src/tree_node.rs):

  • Extends TreeNode with the concept of scopes — a subtree boundary that traversals should not cross unless intended
  • Provides scoped variants of all standard traversal methods: apply_in_scope, visit_in_scope, rewrite_in_scope, transform_down_in_scope, transform_up_in_scope, transform_down_up_in_scope, exists_in_scope
  • Includes DynScopedTreeNode for Arc<dyn T> blanket implementations

New methods on PhysicalExpr trait:

  • children_in_scope() — returns only children that share the same input schema as the current expression. Defaults to children(), so no breaking change for existing implementations
  • with_new_children_in_scope() — reconstructs the expression with new scoped children. Defaults to with_new_children()

Updated expression utilities to use scoped traversal:

  • CaseWhen column projection now uses apply_in_scope / transform_down_in_scope so it no longer traverses into lambda bodies or other new-scope children
  • collect_columns and related helpers in physical-expr/src/utils use scoped traversal where appropriate

No changes to non-scoped traversalsTreeNode::visit, display/formatting, and other tree walks that intentionally need the full tree continue to use children().

Are these changes tested?

Yes:

  • Unit tests for ScopedTreeNode traversal on TestTreeNode with mixed scoped/non-scoped children
  • Tests verifying CaseWhen does not incorrectly project columns that belong to a different scope (e.g. lambda variables)
  • Existing tests continue to pass

Are there any user-facing changes?

Yes, but no breaking changes:

  • New public trait ScopedTreeNode in datafusion-common
  • New default methods children_in_scope() and with_new_children_in_scope() on PhysicalExpr
  • Existing implementations that don't override these methods get the current behavior automatically

@github-actions github-actions bot added physical-expr Changes to the physical-expr crates common Related to common crate physical-plan Changes to the physical-plan crate labels Apr 5, 2026
@rluvaton rluvaton force-pushed the add-support-for-scoped-traversal branch from 9f1d7c6 to c8d2648 Compare April 5, 2026 14:35
@rluvaton
Copy link
Copy Markdown
Member Author

rluvaton commented Apr 5, 2026

@gstvg how would that affect your lambda support, would that make it easier or harder to implement column capture?

note that in column capture you preferably should only keep the columns that are used to avoid a lot of copies

@rluvaton rluvaton marked this pull request as ready for review April 5, 2026 16:28
/// Get a list of child PhysicalExpr that provide the input for this expr.
fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>>;

/// Get a list of child PhysicalExpr that provide the input for this expr that are in the same scope as this expression.
Copy link
Copy Markdown
Contributor

@alamb alamb Apr 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this API is confusing because "scope" isn't something that is defined

I think a beter description of what this API is doing is "children in the outermost scope"

All built in expressions have the (implicit) property that their children are evaluated in the same input schema. I don't see how CASE is any different in theory (though I understand it is what is causing you current problems)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I also think if we add thsi API it is easy to misue and you'll be chasing down corner case bugs with other expression types that have the same problem

@gstvg
Copy link
Copy Markdown
Contributor

gstvg commented Apr 6, 2026

@gstvg how would that affect your lambda support, would that make it easier or harder to implement column capture?

I believe it wouldn't make much difference, that's okay

note that in column capture you preferably should only keep the columns that are used to avoid a lot of copies

Thanks, I will make sure to implement that

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

common Related to common crate physical-expr Changes to the physical-expr crates physical-plan Changes to the physical-plan crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants