feat: add lambda expression support by limameml · Pull Request #710 · substrait-io/substrait-java

limameml · 2026-02-24T14:21:17Z

This PR is to add support for the Lambda Expression

Summary

Add lambda expression support to substrait-java core and isthmus
Implement Type.Func for function types
Add Expression.Lambda and lambda parameter references
Full proto round-trip support
Calcite (isthmus) integration (single-level lambdas only)

Test plan

Added LambdaExpressionRoundtripTest in core
Added LambdaExpressionTest in isthmus

closes #687

CLAassistant · 2026-02-24T14:21:25Z

All committers have signed the CLA.

benbellick

Only just started looking at this but flushing a few comments just to get the ball rolling. On the whole this is looking good though, nice work!

core/src/main/java/io/substrait/expression/Expression.java

benbellick · 2026-02-24T17:07:58Z

core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java

  private final Type.Struct rootType;
  private final ProtoTypeConverter protoTypeConverter;
  private final ProtoRelConverter protoRelConverter;
+  private final List<Type.Struct> lambdaParameterStack = new ArrayList<>();


Can we use a stack here?

I need to access lambdaIndex = lambdaParameterStack.size() - 1 - stepsOut here: lambdaParameters = lambdaParameterStack.get(lambdaIndex); when building the Lambda parametre ref that's why I didn't use a stack

That's fair. Java's stack does extend Vector though so you still get the get method. I'm not so particular on this one.

The only thing I will say is I do think that the whole stack parameter stuff can be confusing for people without a little bit of PL experience, which is why I pushed @Slimsammylim to encapsulate some of that logic and put it in docstring comments. I could see a case for encapsulating this logic in a local private class, though it probably makes more sense to do this only if this logic comes up again.

I expected to see something like this logic elsewhere, as you need to do it to do validation when building lambdas. Though I didn't find anything like this. Does the builder for lambdas in this PR actually validate that the lambda is semantically correct?

I agree that it would be nice for clarity to encapsulate this quite specific behaviour in a class (single responsibility principle). It doesn't need to extend one of the full collection interfaces, just provide the methods you need:

void push(Type.Struct value)

Type.Struct pop() (throw runtime NoSuchElementException if empty)

Type.Struct get(int stepsOut) (throw runtime IndexOutOfBoundsException on invalid stepsOut)

Or make it generic if you will reuse it for other types. The caller doesn't need to worry about clever indexing logic or what collection implementation is used internally; just make simple method calls that clearly describe the intent.

Hey @benbellick, for your last point I don't quite understand what validation I should add, doesn't the validation happen when building the LAMBDA_PARAMETER_REFERENCE?

core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java

Slimsammylim · 2026-02-26T21:33:39Z

core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java

+
+          Type.Struct lambdaParameters = lambdaParameterStack.get(lambdaIndex);
+          return FieldReference.newLambdaParameterReference(
+              reference.getDirectReference().getStructField().getField(),


Something I'm curious about is why we use reference.getDirectReference().getStructField().getField() here instead of getDirectReferenceSegments(reference.getDirectReference()) like the ROOT_REFERENCE case does.

I'm unfamiliar with substrait-java, but does this current implementation support nested field access through lambda parameters?

*This may be an unimportant issue especially if it's not common to access nested fields in lambdas

because the ofRoot method takes the struct itself while the newLambdaParameterReference takes the index, like the newRootStructOuterReference methode

Ah, I see. I am still confused as to whether this supports nested field access through lambda parameters. I messed around w some tests through claude, and I believe it is unsupported. However, this is a nit and I'm not sure it will come up often in practice, so we can come back to this later if it ends up being an issue!

I am still confused as to whether this supports nested field access through lambda parameters.

If by nested access you mean accessing nested fields of a lambda parameter inside the body, for example
($0) -> $0.a.[0].hello then I think that is broken here actually. That might be fine for now and we can treat it is a follow-up item, but we should detect if we have a nested field lookup happening by checking to to content of reference.getDirectReference().getStructField().getChild(); If it's empty, there is only a direct reference to the lambda param. If it's not empty we can blow up.

To @Slimsammylim's we point, to handle this correctly we would need to use getDirectReferenceSegments(reference.getDirectReference()) to create the list of reference segments and then re-use some of the other constructors.

ok I will add a check and go back to it in a follow up pr.

core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java

vbarua · 2026-03-06T03:39:52Z

So I was trying to test this by reading in the test plans from https://github.com/substrait-io/substrait-go/tree/main/expr/testdata/lambda, but then I realized that we were not loading in the functions_list.yaml extensions, and then I realize that we cannot load the extensions because the ANTLR parser hasn't been updated.

I made a PR to do just that, because it's very out of date and required a little extra care #728. Once that's in, I can use the test vectors to help verify this work. We also need that PR anyways, because we can't load any of the functions that use lambdas without it 😓

Note that these had to be tweaked slightly because they were not entirely valid. This will be fixed in substrait-go.

vbarua · 2026-03-10T22:54:05Z

I made a PR into this PR as part of understanding how it works, and fixing a small gap when loading Substrait plans with lambdas: limameml#1

test: add lambda plan roundtrip tests

vbarua

Started taking a more thorough look and found a couple of others things. I'll continue reviwing this tomorrow.

vbarua · 2026-03-12T05:35:19Z

core/src/main/java/io/substrait/expression/Expression.java

+      // TO DO: fix Lambda return type once this issue
+      // https://github.com/substrait-io/substrait/issues/976 is resolved


Suggested change

// TO DO: fix Lambda return type once this issue

// https://github.com/substrait-io/substrait/issues/976 is resolved

// TODO: Type.Func nullability is hardcoded to false here because the spec does not allow for

// declaring otherwise.

// See: https://github.com/substrait-io/substrait/issues/976

Being a bit more precise about what the issue is, and what the implications are in the code.

vbarua · 2026-03-12T06:01:31Z

core/src/main/java/io/substrait/relation/VirtualTableScan.java

+      return type.parameterTypes().stream().mapToInt(p -> p.accept(this)).sum()
+          + type.returnType().accept(this);


This visitor is used to determine how many names are needed for a ReadRel based on the schema.

For example, a table with row type struct<i8, i16, i32> needs 3 names. A table like struct<i8, struct<i16, i32>, i64> needs 5 names, because in Substrait we name every struct field. For full details see https://substrait.io/types/named_structs/#determining-names

Now, given something like struct<i8, func<...>, i32> we don't need to look inside the func for names. The type of the column is just func, and none of the parameters to the func can be reference by name.

All of this is to say that this can just be 0.

Suggested change

return type.parameterTypes().stream().mapToInt(p -> p.accept(this)).sum()

+ type.returnType().accept(this);

return 0;

vbarua · 2026-03-12T06:08:28Z

examples/substrait-spark/src/main/java/io/substrait/examples/util/ExpressionStringify.java

+  @Override
+  public String visit(Expression.Lambda expr, EmptyVisitationContext context)
+      throws RuntimeException {
+    return "<Lambda>";


minor: extract for consistency with the other methods

Suggested change

return "<Lambda>";

return "<Lambda >";

vbarua · 2026-03-12T06:51:14Z

core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java

+
+          Type.Struct lambdaParameters = lambdaParameterStack.get(lambdaIndex);
+          return FieldReference.newLambdaParameterReference(
+              reference.getDirectReference().getStructField().getField(),


I am still confused as to whether this supports nested field access through lambda parameters.

If by nested access you mean accessing nested fields of a lambda parameter inside the body, for example
($0) -> $0.a.[0].hello then I think that is broken here actually. That might be fine for now and we can treat it is a follow-up item, but we should detect if we have a nested field lookup happening by checking to to content of reference.getDirectReference().getStructField().getChild(); If it's empty, there is only a direct reference to the lambda param. If it's not empty we can blow up.

To @Slimsammylim's we point, to handle this correctly we would need to use getDirectReferenceSegments(reference.getDirectReference()) to create the list of reference segments and then re-use some of the other constructors.

vbarua

I did leave a couple of very minor suggestions, and I think that there is one check worth enforcing in #710 (comment).

This is a solid pass at integrating lambdas end-to-end, thanks for working on in @limameml. If you can take a look at the above I think we can get this merged this week.

cc: @bestbeforetoday is there anything else you'd like to see in this PR

benbellick · 2026-03-13T13:31:18Z

core/src/test/java/io/substrait/type/proto/LambdaExpressionRoundtripTest.java

+ */
+class LambdaExpressionRoundtripTest extends TestBase {
+
+  /** Test that lambdas with no parameters are valid. Building: () -> i32(42) : func<() -> i32> */


nit: The parenthesis make it look like a cast to me 🤷

Suggested change

/** Test that lambdas with no parameters are valid. Building: () -> i32(42) : func<() -> i32> */

/** Test that lambdas with no parameters are valid. Building: () -> 42 : func<() -> i32> */

benbellick · 2026-03-13T13:33:33Z

core/src/test/java/io/substrait/type/proto/LambdaExpressionRoundtripTest.java

+    FieldReference invalidRef = FieldReference.newLambdaParameterReference(0, params, 1);
+
+    Expression.Lambda lambda =
+        Expression.Lambda.builder().parameters(params).body(invalidRef).build();


Shouldn't this throw an exception? Perhaps I am misunderstanding what the builder does, but don't we expect the builder to validate that we built a correct lambda? Why are we able to build it, and convert it to a proto, but not convert it back?

IMO this issue will need to be resolved before merging. Sorry to be a blocker! It doesn't make sense to me that we would be able to make an invalid proto and only validate on the way back, unless I am misunderstanding something.

https://github.com/substrait-io/substrait-go/blob/48bbdb43eeccc61c7adfecc3549a1b004fd52fd8/expr/lambda_test.go#L148-L165

^ You can see an example here of how in substrait-go the builder throws an error

Good point Ben. It would make sense for these to fail, especially given that we're already claiming them as validation errors that should fail.

benbellick · 2026-03-13T13:34:04Z

core/src/test/java/io/substrait/type/proto/LambdaExpressionRoundtripTest.java

+        Expression.Lambda.builder().parameters(outerParams).body(innerLambda).build();
+
+    // Convert to proto
+    io.substrait.proto.Expression protoExpression = expressionProtoConverter.toProto(outerLambda);


Same issue here as above I think.

benbellick

Sorry to block this, I know that you have already gotten an approval. I just want to make sure that this is addressed before merging. Thanks!

benbellick · 2026-03-13T13:38:08Z

core/src/test/java/io/substrait/type/proto/LambdaExpressionRoundtripTest.java

+    FieldReference invalidRef = FieldReference.newLambdaParameterReference(0, params, 1);
+
+    Expression.Lambda lambda =
+        Expression.Lambda.builder().parameters(params).body(invalidRef).build();


IMO this issue will need to be resolved before merging. Sorry to be a blocker! It doesn't make sense to me that we would be able to make an invalid proto and only validate on the way back, unless I am misunderstanding something.

feat: add lambda expression support

5ca6ff8

benbellick self-requested a review February 24, 2026 17:00

benbellick reviewed Feb 24, 2026

View reviewed changes

limameml added 2 commits February 25, 2026 10:23

Merge branch 'main' into limame.malainine/add-lambda-support

ed61127

adresse some of @benbellick's comments

a90c609

bestbeforetoday reviewed Feb 26, 2026

View reviewed changes

core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java Outdated Show resolved Hide resolved

Slimsammylim reviewed Feb 26, 2026

View reviewed changes

tweak: encapsulate LambdaParameterStack logic in a class

4053357

benbellick reviewed Mar 2, 2026

View reviewed changes

core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java Show resolved Hide resolved

limameml added 2 commits March 2, 2026 16:25

Merge branch 'main' into limame.malainine/add-lambda-support

ed5de56

tweak: adding comments expailining LambdaParameterStack

eb7b2b2

vbarua added 6 commits March 6, 2026 11:21

Merge branch 'main' into limame.malainine/add-lambda-support

60b66be

build: ignore build related generated files

1d1b26e

feat: enable parsing of func types in extensions

2224c4b

test: copy substrait-go lambda plans

83b9e24

Note that these had to be tweaked slightly because they were not entirely valid. This will be fixed in substrait-go.

test: add LambdaRoundtripTests

d802c19

feat(isthmus): add TRANSFORM SqlFunction to handle transform:list_func

14738e3

limameml added 2 commits March 11, 2026 10:38

Merge pull request #1 from substrait-io/vbarua/lambda-testing

2b5436c

test: add lambda plan roundtrip tests

tweak: add filter in the function mapping

6fae010

vbarua reviewed Mar 12, 2026

View reviewed changes

vbarua approved these changes Mar 13, 2026

View reviewed changes

adress @vbarua's comments

1c1f8d2

benbellick reviewed Mar 13, 2026

View reviewed changes

benbellick requested changes Mar 13, 2026

View reviewed changes

		// TO DO: fix Lambda return type once this issue
		// https://github.com/substrait-io/substrait/issues/976 is resolved

-      // TO DO: fix Lambda return type once this issue
-      // https://github.com/substrait-io/substrait/issues/976 is resolved
+      // TODO: Type.Func nullability is hardcoded to false here because the spec does not allow for
+      //   declaring otherwise.
+      // See: https://github.com/substrait-io/substrait/issues/976

		return type.parameterTypes().stream().mapToInt(p -> p.accept(this)).sum()
		+ type.returnType().accept(this);

	return type.parameterTypes().stream().mapToInt(p -> p.accept(this)).sum()
	+ type.returnType().accept(this);
	return 0;

	/** Test that lambdas with no parameters are valid. Building: () -> i32(42) : func<() -> i32> */
	/** Test that lambdas with no parameters are valid. Building: () -> 42 : func<() -> i32> */

Conversation

limameml commented Feb 24, 2026

Summary

Test plan

Uh oh!

CLAassistant commented Feb 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

benbellick left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

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!

bestbeforetoday Feb 26, 2026 • 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!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Slimsammylim Feb 26, 2026 • 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!

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!

Uh oh!

vbarua commented Mar 6, 2026

Uh oh!

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

Uh oh!

Uh oh!

vbarua 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!

vbarua 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!

benbellick left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

CLAassistant commented Feb 24, 2026 •

edited

Loading

bestbeforetoday Feb 26, 2026 •

edited

Loading

Slimsammylim Feb 26, 2026 •

edited

Loading

vbarua commented Mar 10, 2026 •

edited

Loading