feat: add lambda support #183

Slimsammylim · 2026-01-20T18:32:50Z

closes #187

This PR adds lambda support to substrait-go.

Note:

ExprFromProto does not include the same validation process as .Build(). ExprFromProto resolves types but does not validate field references. I will be adding this in a future PR.
LambdaInvocation is not included in this PR and will be implemented separately.

codecov · 2026-01-21T19:06:16Z

Codecov Report

❌ Patch coverage is 78.46890% with 45 lines in your changes missing coverage. Please review.
✅ Project coverage is 69.09%. Comparing base (579861e) to head (e1a53f0).

Files with missing lines	Patch %	Lines
types/types.go	67.69%	20 Missing and 1 partial ⚠️
expr/builder.go	85.93%	8 Missing and 1 partial ⚠️
expr/expression.go	55.55%	4 Missing and 4 partials ⚠️
expr/field_reference.go	75.00%	2 Missing and 2 partials ⚠️
expr/lambda.go	93.47%	2 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #183      +/-   ##
==========================================
+ Coverage   68.79%   69.09%   +0.29%     
==========================================
  Files          45       46       +1     
  Lines       10367    10576     +209     
==========================================
+ Hits         7132     7307     +175     
- Misses       2886     2912      +26     
- Partials      349      357       +8

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Slimsammylim · 2026-01-21T21:02:17Z

expr/lambda.go

+// Recurses into all descendants (e.g., Cast(FieldRef), ScalarFunction(Cast(FieldRef))).
+func resolveLambdaParamTypes(body Expression, params *types.StructType, outerParams []*types.StructType) Expression {
+	// Handle if body itself is a nested lambda - must check BEFORE Visit
+	// because Lambda.Visit would bypass our nested lambda context handling


Lambda.Visit(callback) internally calls callback(lambda.Body), which skips the current Lambda wrapper entirely. This loses the context we need to resolve stepsOut>0 references. Therefore, we need to handle if the body is a lambda first.

Slimsammylim · 2026-01-21T21:13:25Z

expr/lambda.go

+	}
+
+	// Resolve types for FieldReferences with LambdaParameterReference roots
+	resolvedBody := resolveLambdaParamTypes(body, b.parameters, outerParams)


After constructing the lambda, we walk the expression tree and explicitly resolve types by looking up params.Types[fieldIndex] for each lambda parameter reference. For nested lambdas, we pass down the outer lambda's parameters so stepsOut > 0 references can also be resolved.

benbellick

Wasn't sure where exactly to say this, but consider reducing the number of comments across all of the added code. In my view, the ideal code is self-documenting, i.e. the code should be optimized for readability. If the intent is clear, then no comments necessary. However, for anything tricky or for docstrings, comments are absolutely essential.

I haven't done a full review yet but here are a couple of things to consider. Great progress! 🚀

expr/field_reference.go

expr/lambda_test.go

expr/lambda.go

Co-authored-by: Ben Bellick <[email protected]>

… parameters

…FieldRefs

benbellick · 2026-01-29T22:11:14Z

expr/builder.go

+	targetParams := lpb.b.lambdaContext[len(lpb.b.lambdaContext)-1-int(lpb.stepsOut)]
+
+	// Validate that the parameter exists by checking field index for StructFieldRef
+	if structRef, ok := lpb.ref.(*StructFieldRef); ok {


We discussed offline, but I think this ref must be a StructFieldRef, though the spec doesn't do a great job of making that clear in the proto messages. Investigating right now...

substrait-io/substrait#955

expr/builder.go

expr/testdata/lambda/lambda-field-ref.json

types/types.go

benbellick · 2026-01-29T23:22:46Z

types/types.go

+	out := *f
+	out.Nullability = n
+	return &out
+}


I think that this is not a deep copy, so the same ParameterTypes is used in the implicit parameter and the returned parameter. Is that safe to do in this context? It might be, I genuinely am unsure.

StructType also does a shallow copy. Maybe @vbarua can offer context on if this is the intended semantics of the WithNullability function? I don't think you wrote this code, but maybe you have historical context 😅

The comment just says:

// WithNullability returns a copy of this type but with // the nullability set to the passed in value

Yes, it is a copy... but unclear if it is meant to be a deep copy or not.

I think it would be very strange if you had a reason to alter the parameters of the type after it had been constructed, but it is probably worth adding a comment somewhere about this invariant if there isn't one already.

Hm, I thought this function was a bit strange too. This definition was mostly me just following how previous types have defined this function

Maybe making a deep copy is too expensive for a deeply nested struct?

expr/lambda_test.go

Co-authored-by: Ben Bellick <[email protected]>

…substrait-go into slim/adding-lambda-support

Co-authored-by: Ben Bellick <[email protected]>

…substrait-go into slim/adding-lambda-support

benbellick

Some minor cosmetic things, but the implementation itself is awesome 💯

Going through the tests now, but wanted to dump comments so that you can work on them in parallel.

benbellick · 2026-01-30T20:17:41Z

expr/builder.go

+// the current lambda to find the target parameter (0 = current lambda, 1 = outer lambda, etc.).
+//
+// The Build method validates the stepsOut value and field index against the
+// available lambda context, and resolves the reference type automatically.


After a discussion with @vbarua, I believe we may end up leaving it in the spec for now that you can use a MaskedExpression reference with lambdas, though we don't know what that means. Thus, this implementation is completely sufficient for now, we can just say we don't build that yet. But it will be good to call it out, specifically linking to an issue in upstream about clarification. I haven't made that issue, but when it exists, we should link it in these docs.

benbellick · 2026-01-30T20:27:06Z

expr/builder.go

+
+	// Validate that the parameter field index exists
+	if int(lpb.ref.Field) >= len(targetParams.Types) {
+		return nil, fmt.Errorf("%w: lambda body references parameter %d but lambda only has %d parameters",


Should we mention something about the steps out here too?

Otherwise, it could be ambiguous from the error message which one we are talking about. Let's also make sure some test captures that this info is present in the error message.

I added a better error message to include the steps out value. I also updated the builder error tests to search for this info in the error message.

expr/builder.go

benbellick · 2026-01-30T20:38:34Z

expr/expression.go

-	}
+	case *proto.Expression_Lambda_:
+		if et.Lambda.Parameters == nil {
+			return nil, fmt.Errorf("%w: lambda parameters cannot be nil", substraitgo.ErrInvalidExpr)


This is totally fine, but it is a bummer that we use ErrInvalidExpr throughout this file already. IMO, there should be ErrInvalidExprBuilder and ErrInvalidExprProto or something like that. But that's for a later PR 😅

benbellick · 2026-01-30T20:47:32Z

types/func_type_test.go

Can we actually drop this file and move it to type_test.go? We can add the roundtrip tests as just additional cases in:

substrait-go/types/types_test.go

Line 61 in 579861e

func TestTypeRoundtrip(t *testing.T) {

I think there is similarly already also a test for type short names that we can add one more case too. Maybe one for the long name though I haven't looked.

Co-authored-by: Ben Bellick <[email protected]>

…substrait-go into slim/adding-lambda-support

benbellick · 2026-01-30T21:21:01Z

expr/testdata/lambda/basic-lambda.json

The library doesn't have the facilities to catch these yet, but I think all of the JSON tests here are syntactically correct, but semantically meaningless. To simplify the discussion, I am using a dummy explain format.

This plan here is basically the following thing:

Root Project expressions: [λ($0: i32) -> 42] input: Read(virtual_table [[0]])

However, I think what you intend is something like:

Root Project expressions: [transform(field_ref(0), λ($0: i32) -> 42)] input: Read( virtual_table [[[0]]], schema=[values: list<i32>] )

The difference here is that in the first one, the expression just evaluates to a lambda. We may technically allows lambdas as column values... but I'm not sure about that. But the second says to apply the lambda.

I also think these don't all need to be virtual tables, and it might simplify them to just make a read rel of some made up table:

Root Project expressions: [transform(field_ref(0), λ($0: i32) -> 42)] input: Read( namedTable="test_table", schema=[values: list<i32>] )

Also, using transform means that we are testing that functions_list.yaml is loaded correctly. However, it seems like it isn't working locally for me. This YAML file must be missing somehow.

When all that other stuff is said and done though, then I think this is a possible form of that basic plan:

{ "version": {"majorNumber": 0, "minorNumber": 79}, "extensionUrns": [ { "extensionUrnAnchor": 1, "urn": "extension:io.substrait:functions_list" } ], "extensions": [ { "extensionFunction": { "extensionUrnReference": 1, "functionAnchor": 1, "name": "transform" } } ], "relations": [{ "root": { "input": { "project": { "common": {"direct": {}}, "input": { "read": { "common": {"direct": {}}, "baseSchema": { "names": ["values"], "struct": { "nullability": "NULLABILITY_REQUIRED", "types": [{ "list": { "type": {"i32": {"nullability": "NULLABILITY_REQUIRED"}}, "nullability": "NULLABILITY_REQUIRED" } }] } }, "namedTable": { "names": ["test_table"] } } }, "expressions": [{ "scalarFunction": { "functionReference": 1, "outputType": { "list": { "type": {"i32": {"nullability": "NULLABILITY_REQUIRED"}}, "nullability": "NULLABILITY_REQUIRED" } }, "arguments": [ { "value": { "selection": { "directReference": { "structField": {"field": 0} }, "rootReference": {} } } }, { "value": { "lambda": { "parameters": { "nullability": "NULLABILITY_REQUIRED", "types": [{"i32": {"nullability": "NULLABILITY_REQUIRED"}}] }, "body": { "literal": {"i32": 42} } } } } ] } }] } }, "names": ["result"] } }] }

It would be awesome to be able to have something like substrait-explain for this to make these readable, but that is for another day and another battle.

benbellick · 2026-01-30T21:24:33Z

go.mod

@@ -16,7 +16,7 @@ require (
 	github.com/google/uuid v1.6.0
 	github.com/stretchr/testify v1.10.0
 	github.com/substrait-io/substrait v0.78.1


Suggested change

github.com/substrait-io/substrait v0.78.1

github.com/substrait-io/substrait v0.79.0

this should get us access to functions_list.yaml

Okay... this is opening up a can of worms. If you do this, we will fail parsing these YAML files. To fix this, you then need to update this line:

substrait-go/grammar/generate.go

Line 5 in 579861e

// using substrait v0.69.0

to v0.79.0

Then when you call go generate, it will update to use the new grammar which includes funcs.

Then unfortunately, you'll need to update some more code.

Unresolving this one, I think you still haven't called go generate to regenerate all of the code here

Co-authored-by: Ben Bellick <[email protected]>

Slimsammylim added 6 commits January 20, 2026 13:30

First try at lambda support

29ebf33

Changing code to use the builder pattern

86704d5

Adding an extra guard for invalid plans when we try to resolve types

bdae181

Changing the builder struct to not use interface{}

4b6ef3b

Adding clarifying comments

941b5a5

Making visitors recursive

c63aca6

Handling when a body is a lambda when resolving param types

a76e2ac

Slimsammylim commented Jan 21, 2026

View reviewed changes

Slimsammylim changed the title ~~First try at lambda support~~ feat: Add Lambda support Jan 22, 2026

Slimsammylim changed the title ~~feat: Add Lambda support~~ feat: add lambda support Jan 22, 2026

Slimsammylim changed the title ~~feat: add lambda support~~ feat: add lambda support (#187) Jan 22, 2026

Slimsammylim changed the title ~~feat: add lambda support (#187)~~ feat: add lambda support Jan 22, 2026

Slimsammylim added 3 commits January 22, 2026 10:49

Adding tests for basic functions

9b64203

More test coverage

2d8ac51

Validation Error test cases

86b2046

benbellick reviewed Jan 22, 2026

View reviewed changes

Slimsammylim and others added 12 commits January 22, 2026 14:30

Update expr/field_reference.go

49ce3d2

Co-authored-by: Ben Bellick <[email protected]>

Update expr/lambda_test.go

e77547b

Co-authored-by: Ben Bellick <[email protected]>

Update expr/lambda_test.go

1c283a7

Co-authored-by: Ben Bellick <[email protected]>

Update expr/lambda_test.go

b787bb3

Co-authored-by: Ben Bellick <[email protected]>

Update expr/lambda_test.go

6db3eb1

Co-authored-by: Ben Bellick <[email protected]>

Using URNs not URIs

9aa140d

Using the builder interface in builder.go

30b650e

Using a full plan for TestLambdaWithFunctionExprFromProto

0966f42

Removing redundant getters and unnecessary tests

b05f2e0

Editing tests and comments to make it clear that a lambda can have no…

c639da5

… parameters

Removing unnecessary comments

87590bd

Fixing inconsistent testcase

153cdae

Slimsammylim added 3 commits January 29, 2026 16:07

Adding FuncType as the type for lambda bodies

0b5805b

Adding FuncType as the type for lambda bodies

e1a53f0

Realized we can assume that all lambdaParamRefBuilder refs are Struct…

cfe657c

…FieldRefs

benbellick reviewed Jan 29, 2026

View reviewed changes

benbellick mentioned this pull request Jan 29, 2026

docs: forbid masked_reference for lambda parameters substrait-io/substrait#955

Closed

benbellick reviewed Jan 29, 2026

View reviewed changes

Slimsammylim and others added 11 commits January 30, 2026 10:00

Update expr/builder.go

5c08270

Co-authored-by: Ben Bellick <[email protected]>

Only using full plans for lambda tests

a2e003c

Merge branch 'slim/adding-lambda-support' of github.com:Slimsammylim/…

03c6ab6

…substrait-go into slim/adding-lambda-support

lambdaParamRefs only use StructFieldRef

c35ae76

Update expr/lambda_test.go

c8ee7d7

Co-authored-by: Ben Bellick <[email protected]>

Update expr/lambda_test.go

41422d1

Co-authored-by: Ben Bellick <[email protected]>

Update expr/lambda_test.go

e9c27a8

Co-authored-by: Ben Bellick <[email protected]>

Update expr/lambda_test.go

f56ba94

Co-authored-by: Ben Bellick <[email protected]>

Adding literal package to lambda tests

51dd93e

Merge branch 'slim/adding-lambda-support' of github.com:Slimsammylim/…

cd63be4

…substrait-go into slim/adding-lambda-support

Adding comments to clarify the desired lambda in tests

9fd7d68

benbellick reviewed Jan 30, 2026

View reviewed changes

Slimsammylim and others added 9 commits January 30, 2026 15:59

Better error message for getting lambda param types

7a58a30

Updating lambda tests

05e9370

Update expr/builder.go

d1e76bf

Co-authored-by: Ben Bellick <[email protected]>

Update expr/builder.go

a75b372

Co-authored-by: Ben Bellick <[email protected]>

Update expr/builder.go

3717967

Co-authored-by: Ben Bellick <[email protected]>

Update expr/builder.go

2b898a8

Co-authored-by: Ben Bellick <[email protected]>

Update expr/builder.go

6e04873

Co-authored-by: Ben Bellick <[email protected]>

Moving func type tests to types_test.go

17b813c

Merge branch 'slim/adding-lambda-support' of github.com:Slimsammylim/…

b3fe321

…substrait-go into slim/adding-lambda-support

benbellick reviewed Jan 30, 2026

View reviewed changes

Slimsammylim and others added 2 commits January 30, 2026 16:46

Update go.mod

1becf60

Co-authored-by: Ben Bellick <[email protected]>

Updating substrait version

8835174

	github.com/substrait-io/substrait v0.78.1
	github.com/substrait-io/substrait v0.79.0

feat: add lambda support #183

Are you sure you want to change the base?

feat: add lambda support #183

Uh oh!

Conversation

Slimsammylim commented Jan 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov bot commented Jan 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Slimsammylim Jan 21, 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!

benbellick left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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!

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!

Uh oh!

Uh oh!

Uh oh!

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

Uh oh!

Uh oh!

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!

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

Slimsammylim commented Jan 20, 2026 •

edited

Loading

codecov bot commented Jan 21, 2026 •

edited

Loading

Slimsammylim Jan 21, 2026 •

edited

Loading