Add lambda substrait support by gstvg · Pull Request #21193 · apache/datafusion

gstvg · 2026-03-27T06:52:00Z

Which issue does this PR close?

Part of #21172

Rationale for this change

Substrait support wasn't implemented in the core lambda support to reduce PR size

What changes are included in this PR?

Substrait consuming and producing of higher-order functions, lambdas and lambda variables

Are these changes tested?

Unit tests added to datafusion/substrait/tests/cases/roundtrip_logical_plan.rs

Are there any user-facing changes?

None

benbellick · 2026-04-28T15:22:47Z

👋 Hello from @substrait-io. Great to see the core lambda PR has gotten through! Once this PR is in a ready to review state and is rebased off of main, I will be more than happy to help review it 🙂

gstvg · 2026-04-29T03:55:07Z

Thanks @benbellick, I will open this tonight. Besides rebasing, I believe it misses some tests (I tested with sqllogictests only)

gstvg · 2026-05-08T13:23:26Z

+    /// Returns a new instance of this consumer which includes the given `lambda_parameters` and the names they got assigned
+    ///
+    /// Note for custom implementations it's possible to embed a [DefaultSubstraitLambdaConsumer] and forward this method to it
+    fn with_lambda_parameters(
+        &self,
+        lambda_parameters: &[Type],
+        input_schema: &DFSchema,
+    ) -> datafusion::common::Result<(Vec<String>, Self)>;


@benbellick I usually follows the pattern of existing methods, like push/pop_outer_schema from #20439. This is my first time dealing with substrait, so I may be wrong, but I didn't followed this pattern (using push/pop_lambda_parameter) because it modifies &self via a RwLock and I'm note sure this couldn't lead to conflicts if the same consumer is used to consume two different plans at the same time in different threads. If that's not the case I can change this to use push/pop_lambda_parameter as well.

Ah, that is an interesting point... Is it expected/supported for the same SubstraitConsumer instance to be used concurrently?

If not, then I think it would be simpler and consistent to model lambda scope the same way as outer schemas.

If yes, then the scoped-consumer approach here makes sense, but it seems like the existing push_outer_schema / pop_outer_schema stack may have the same interleaving issue and should probably be addressed separately.

I'll confirm with the maintainer who ends up reviewing this, but SubstraitConsumer is both Send + Sync and all it's methods take &self. Since the default consumer is cheap to create, and I expect most/all custom ones to be cheap as well, I guess it's mostly due to async and to easily embed it into other structures which should also implement Send + Sync than to allow efficient concurrent usage.
I won't expect any consumer to be used concurrently for performance, but, since it can be used, I think it's possible that it's/will be used concurrently incidentally as the easier/natural way within a given codebase

If yes, then the scoped-consumer approach here makes sense, but it seems like the existing push_outer_schema / pop_outer_schema stack may have the same interleaving issue and should probably be addressed separately.

@benbellick is right on point here, I'd aim for consistency in this PR, and address separately the multiple threads dealing with the same consumer in a different PR.

that makes sense, done at 44384dd thanks

benbellick

Few more comments but on the whole this looks good to me! Thanks

Co-authored-by: Ben Bellick <36523439+benbellick@users.noreply.github.com>

benbellick

LGTM, great work! Excited to get this in 🚀

Feel free to summon the maintainers now.

… lambda_substrait

gstvg · 2026-05-14T11:32:13Z

LGTM, great work! Excited to get this in 🚀

Thanks for the review @benbellick
cc @rluvaton @comphead this is ready for the final review

rluvaton · 2026-05-14T12:33:31Z

I'm unfamiliar with substrait except for the idea, so I don't know if I can have meaningful review

rluvaton · 2026-05-14T12:36:03Z

@gstvg according to the breaking change detector there are no breaking change api wise in this pr, is your commemt still correct?

Are there any user-facing changes?

Yes, there are breaking changes, new methods without default implementation have been added to SubstraitConsumer and SubstraitProducer

gstvg · 2026-05-15T00:18:48Z

@rluvaton, indeed, after applying @benbellick suggestions there's no breaking changes anymore, description updated, thanks

gstvg · 2026-05-18T19:39:32Z

Hey @gabotechs, I see you worked on some substrait PRs, bringing you here in case you have interest in this. thanks

gabotechs

Looks good! gave a shallow pass and just made one suggestion, but I trust @benbellick in his review (he's a work colleague). Thanks @gstvg for your work with lambdas 🙏

gstvg · 2026-05-19T06:43:36Z

Thanks @gabotechs. The derive(Clone) is removed. There's only two comments from @benbellick that were outside substrait translation and weren't fully addressed, #21193 (review) and #21193 (comment)

gabotechs · 2026-05-19T07:28:16Z

There's only two comments from @benbellick that were outside substrait translation and weren't fully addressed, #21193 (review) and #21193 (comment)

Answered there

gstvg · 2026-05-19T10:00:59Z

        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, (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)))
+        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)))


@gabotechs that's expected since now there's one global next_lambda_parameter_name, before it was one per lambda scope (applies to the other test as well)

gstvg mentioned this pull request Mar 27, 2026

[EPIC] Full lambda support #21172

Open

27 tasks

add producer test

6cdfc05

github-actions Bot added the auto detected api change Auto detected API change label May 8, 2026

gstvg commented May 8, 2026

View reviewed changes

Comment thread datafusion/substrait/tests/cases/roundtrip_logical_plan.rs Outdated

add default impl to new lambda methods

2c23439

github-actions Bot removed the auto detected api change Auto detected API change label May 9, 2026

benbellick reviewed May 11, 2026

View reviewed changes

gstvg and others added 5 commits May 11, 2026 16:05

Apply suggestion from @benbellick

3d83fb4

Co-authored-by: Ben Bellick <36523439+benbellick@users.noreply.github.com>

improve serialize test readability

83e261f

add lambda field reference tests

c019769

add lambda tests

93f17e4

check invalid lambda

347bf12

benbellick approved these changes May 13, 2026

View reviewed changes

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

Comment thread datafusion/substrait/tests/cases/serialize.rs

gstvg added 2 commits May 14, 2026 08:23

add comment to collect_lambda_ref test helper

59fb308

Merge branch 'main' of https://github.com/gstvg/arrow-datafusion into…

e9a9441

… lambda_substrait

fix for new sqlparser version

898bb21

gabotechs approved these changes May 19, 2026

View reviewed changes

Comment thread datafusion/substrait/src/extensions.rs Outdated

gstvg added 2 commits May 19, 2026 03:29

remove derive(Clone) from Extensions

d78d49c

Merge branch 'main' into lambda_substrait

ad22039

use push/pop_lambda_parameters instead of with_lambda_parameters

44384dd

gstvg commented May 19, 2026

View reviewed changes

gstvg added 2 commits May 19, 2026 08:19

fix doc comment

bee313a

revert submodule bump

555a49f

Conversation

gstvg commented Mar 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

benbellick commented Apr 28, 2026

Uh oh!

gstvg commented Apr 29, 2026

Uh oh!

gstvg May 8, 2026

Choose a reason for hiding this comment

Uh oh!

benbellick May 8, 2026

Choose a reason for hiding this comment

Uh oh!

gstvg May 9, 2026

Choose a reason for hiding this comment

Uh oh!

gabotechs May 19, 2026

Choose a reason for hiding this comment

Uh oh!

gstvg May 19, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

benbellick left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

benbellick left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

gstvg commented May 14, 2026

Uh oh!

rluvaton commented May 14, 2026

Uh oh!

rluvaton commented May 14, 2026

Are there any user-facing changes?

Uh oh!

gstvg commented May 15, 2026

Uh oh!

gstvg commented May 18, 2026

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!

gstvg commented May 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gabotechs commented May 19, 2026

Uh oh!

gstvg May 19, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

gstvg commented Mar 27, 2026 •

edited

Loading

gabotechs left a comment •

edited

Loading

gstvg commented May 19, 2026 •

edited

Loading