refactor(substrait): pass immutable scope context through consumers by alexanderbianchi · Pull Request #26890 · apache/datafusion

alexanderbianchi · 2026-10-29T16:01:34Z

Rationale for this change

Lambda and outer-query scopes currently live in mutable stacks on the Substrait consumer. Passing an explicit context makes each conversion's scope visible and removes the need to push and pop state around nested conversions.

This PR keeps input consumption sequential. Parallel consumption and its performance evaluation can be handled separately.

What changes are included in this PR?

Add SubstraitConsumerContext and pass it through relation and expression callbacks and conversion helpers.
Derive child contexts when entering lambda or outer-query scopes; the parent stays unchanged on success, error, or cancellation.
Support custom scoped values through with_extension / get_extension, using DataFusion’s typed Extensions container. Child scopes inherit values and can override them without changing parents or siblings.
Remove scope stacks and their push/pop methods from the consumer, including DefaultSubstraitLambdaConsumer.
Preserve existing sequential input consumption.
Keep plan-level entrypoint signatures unchanged and document custom-consumer migration in the DataFusion 57.0.0 upgrade guide.

Testing

Existing Substrait tests cover relation conversion, correlated subqueries, and lambda behavior. Direct context tests cover scope inheritance and sibling isolation; a roundtrip execution test covers differently typed sibling lambdas and nested captures. Concurrency-only tests and their fixture have been removed.

Validation passed: Rust and Markdown formatting, workspace Clippy with all targets and features (warnings denied), 68 Substrait unit tests, 230 integration tests (6 ignored), and 3 doctests.

Are there any user-facing changes?

Breaking API change: custom relation/expression callbacks and direct conversion-helper calls must accept and forward &SubstraitConsumerContext. The upgrade guide explains how to replace the removed scope methods.

Query results and the default sequential scheduling policy are preserved. Generated lambda parameter names can differ in plan displays because independent lexical scopes may reuse names.

Custom consumers can keep their scoped values in the context using typed extensions. Values are shared through Arc; child-local changes replace a value. Interior mutable state is still shared, so opting into concurrency remains the custom consumer’s responsibility.

alexanderbianchi · 2026-11-03T17:31:26Z

@gabotechs This PR is fairly straightforward, even if the diff is scary. The main change is in datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs and the rest are updating function signatures + callers.

If we agree that structure makes sense, the 2 remaining question are

Should we do parallelization here or leave that to the user?
Do we still want the multi-input consume function from feat(substrait): add a hook for consuming relation inputs #26603?

gabotechs

Looks good! for this PR, I'd just remove the consume_rel_inputs hook, and I think it should be ready to go.

gabotechs · 2026-11-10T08:22:14Z

+    /// The default consumes inputs sequentially, stopping at the first error.
+    /// Overrides can provide a different scheduling policy, but must account for
+    /// any additional mutable traversal state in the custom consumer.
+    async fn consume_rel_inputs(


I'm still not convinced about this change. For the sake of this PR, I'd remove it, and we can discuss in a follow up

gabotechs · 2026-11-10T08:22:39Z

    pub(super) state: &'a SessionState,
-    outer_schemas: RwLock<Vec<Arc<DFSchema>>>,
-    lambda_consumer: DefaultSubstraitLambdaConsumer,
 }



gabotechs · 2026-11-10T08:24:35Z

-        assert!(consumer.get_outer_schema(0).is_none());
-        assert!(consumer.get_outer_schema(1).is_none());
-        assert!(consumer.get_outer_schema(2).is_none());
+    fn extensions_inherit_and_override_without_changing_other_scopes() {


🤔 these tests were pretty coupled to implementation details, and now they are getting in the way rather than helping.

No action required in this PR, but if you find yourself in the position of adding new test, rather than contributing to these ones, I'd try to just test through the public API.

gabotechs · 2026-11-10T13:59:37Z

@@ -545,7 +558,7 @@ mod tests {
        assert_snapshot!(
        plan,
        @"
-        Projection: array_transform2(make_array(make_array(data3.p1)), (p0, p2) -> array_concat(array_transform2(p0, (p3, p4) -> p3 * p2 * p4), array_transform2(p0, (p5, p6) -> p5 * p2 * p6))) AS array_transform2(make_array(make_array(data3.p1)),(v, i) -> array_concat(array_transform2(v,(v, j) -> v * i * j),array_transform2(v,(v, j) -> v * i * j)))
+        Projection: array_transform2(make_array(make_array(data3.p1)), (p0, p2) -> array_concat(array_transform2(p0, (p3, p4) -> p3 * p2 * p4), array_transform2(p0, (p3, p4) -> p3 * p2 * p4))) AS array_transform2(make_array(make_array(data3.p1)),(v, i) -> array_concat(array_transform2(v,(v, j) -> v * i * j),array_transform2(v,(v, j) -> v * i * j)))
          TableScan: data3


🤔 at first I thought there was something wrong with this snapshot change, but it does look correct, just that the argument numeration is not global to whole plan, but to the specific plan branch in which that array_transform happened to fall.

gabotechs · 2026-11-10T14:01:01Z

+nested conversions. There is no matching pop operation; the parent context stays
+unchanged on success, error, and cancellation.
+
+Input consumption remains sequential. Custom consumers


This looks like the classical LLM correction over a previous state of this PR that is not going to be shipped. I'd revisit the upgrade guide.

Also, branch-56 was already cut, and this PR is not there, so this might belong to the 57 upgrade guide instead.

I think humans also screw this up 😆

…opes

gabotechs

All good here, thanks @alexanderbianchi and @LiaCastaneda!

github-actions Bot added substrait Changes to the substrait crate auto detected api change Auto detected API change labels Oct 29, 2026

This was referenced Oct 29, 2026

feat(substrait): fork consumers for concurrent relation inputs #26896

Draft

feat(substrait): add a hook for consuming relation inputs #26603

Closed

alexanderbianchi marked this pull request as ready for review November 3, 2026 17:28

LiaCastaneda reviewed Nov 5, 2026

View reviewed changes

Comment thread datafusion/substrait/README.md Resolved

Comment thread datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs Resolved

Comment thread datafusion/substrait/src/logical_plan/consumer/rel/join_rel.rs Resolved

github-actions Bot added the documentation Improvements or additions to documentation label Nov 6, 2026

gabotechs reviewed Nov 10, 2026

View reviewed changes

alexanderbianchi added 9 commits November 10, 2026 09:35

feat(substrait): add a hook for consuming relation inputs

7e9180a

refactor(substrait): pass immutable scope context through consumers

41a1f36

test(substrait): use table_scan in consumer test fixture

0e977f0

test(substrait): simplify scope and concurrent input tests

99ae0a9

refactor(substrait): preserve sequential consumption with explicit sc…

58ae265

…opes

docs(substrait): keep migration notes in the upgrade guide

87ae625

feat(substrait): support scoped consumer context extensions

4d8d227

refactor(substrait): defer input consumption hook to a follow-up

4e24776

docs(substrait): document consumer migration for DataFusion 57

861b014

alexanderbianchi force-pushed the codex/substrait-explicit-context branch from c889abf to 861b014 Compare November 10, 2026 14:52

gabotechs approved these changes Nov 10, 2026

View reviewed changes

LiaCastaneda approved these changes Nov 10, 2026

View reviewed changes

github-actions Bot removed the auto detected api change Auto detected API change label Nov 10, 2026

gabotechs added this pull request to the merge queue Nov 10, 2026

Merged via the queue into apache:main with commit 8fbcc00 Nov 10, 2026
45 of 46 checks passed

github-actions Bot added the v57.0.0 label Nov 10, 2026

Conversation

alexanderbianchi commented Oct 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Rationale for this change

What changes are included in this PR?

Testing

Are there any user-facing changes?

Uh oh!

alexanderbianchi commented Nov 3, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gabotechs left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gabotechs left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

alexanderbianchi commented Oct 29, 2026 •

edited

Loading

gabotechs left a comment •

edited

Loading