feat: add lambda expression support

expr/builder.go

@@ -1,8 +1,15 @@
 // SPDX-License-Identifier: Apache-2.0
 
+// Package expr provides types and builders for constructing Substrait expressions.
+//
+// IMPORTANT: Only ExprBuilder methods guarantee construction of valid expressions. Manual
+// construction using struct literals bypasses validation and may create invalid expressions. It is highly recommended to only construct expressions via the builders
 package expr
 
 import (
+	"fmt"
+
+	substraitgo "github.com/substrait-io/substrait-go/v7"
 	"github.com/substrait-io/substrait-go/v7/extensions"
 	"github.com/substrait-io/substrait-go/v7/types"
 )
@@ -18,9 +25,9 @@ type Builder interface {
 
 // ExprBuilder is the parent context for all the expression builders.
 // It maintains a pointer to an extension registry and, optionally,
-// a pointer to a base input schema. This allows less verbose expression
-// building as it isn't necessary to pass these to every `New*` function
-// to construct the expressions.
+// a pointer to a base input schema.
+// This allows less verbose expression building as it isn't necessary
+// to pass these to every `New*` function to construct the expressions.
 //
 // This is intended to be used like:
 //
@@ -36,8 +43,9 @@ type Builder interface {
 //
 // See the unit tests for additional examples / constructs.
 type ExprBuilder struct {
-	Reg        ExtensionRegistry
-	BaseSchema *types.RecordType
+	Reg           ExtensionRegistry
+	BaseSchema    *types.RecordType
+	lambdaContext []*types.StructType // keeps track of the lambda context for nested lambdas
 }
 
 // Literal returns a wrapped literal that can be passed as an argument
@@ -146,6 +154,21 @@ func (e *ExprBuilder) RootRef(ref Reference) *fieldRefBuilder {
 	return e.Ref(RootReference, ref)
 }
 
+// LambdaParamRef constructs a field reference to a lambda parameter.
+//
+// The stepsOut parameter specifies how many lambda scopes to traverse outward from
+// the current lambda to find the target parameter.
+//
+// The Build method validates the stepsOut value and field index against the
+// available lambda context, and resolves the reference type automatically.
+func (e *ExprBuilder) LambdaParamRef(ref Reference, stepsOut uint32) *lambdaParamRefBuilder {
+	return &lambdaParamRefBuilder{
+		b:        e,
+		ref:      ref,
+		stepsOut: stepsOut,
+	}
+}
+
 // Cast returns a builder for constructing a Cast expression. The failure
 // behavior can be specified by calling FailBehavior before calling Build.
 func (e *ExprBuilder) Cast(from Builder, to types.Type) *castBuilder {
@@ -154,6 +177,83 @@ func (e *ExprBuilder) Cast(from Builder, to types.Type) *castBuilder {
 	}
 }
 
+// Lambda returns a builder for constructing a Lambda expression with the
+// given parameters.
+//
+// When building nested lambdas (e.g., a function that takes a lambda argument
+// which itself references outer lambda parameters), the ExprBuilder maintains
+// a context stack that allows inner lambdas to validate stepsOut references
+// against outer lambda parameters.
+func (e *ExprBuilder) Lambda(params *types.StructType, body Builder) *lambdaBuilder {
+	return &lambdaBuilder{
+		b:      e,
+		params: params,
+		body:   body,
+	}
+}
+
+// pushLambdaContext pushes a lambda parameters struct onto the expression builder's context stack.
+func (e *ExprBuilder) pushLambdaContext(params *types.StructType) {
+	e.lambdaContext = append(e.lambdaContext, params)
+}
+
+// popLambdaContext pops a lambda parameters struct off the expression builder's context stack.
+func (e *ExprBuilder) popLambdaContext() {
+	e.lambdaContext = e.lambdaContext[:len(e.lambdaContext)-1]
+}
+
+type lambdaBuilder struct {
+	b      *ExprBuilder
+	params *types.StructType
+	body   Builder
+}
+
+// Build constructs and validates the Lambda expression.
+func (lb *lambdaBuilder) Build() (*Lambda, error) {
+	if lb.params == nil {
+		return nil, fmt.Errorf("%w: lambda must have parameters struct", substraitgo.ErrInvalidExpr)
+	}
+	if lb.params.Nullability != types.NullabilityRequired {
+		return nil, fmt.Errorf("%w: lambda parameters struct must have NULLABILITY_REQUIRED", substraitgo.ErrInvalidExpr)
+	}
+	for i, paramType := range lb.params.Types {
+		if paramType == nil {
+			return nil, fmt.Errorf("%w: lambda parameter %d has nil type", substraitgo.ErrInvalidExpr, i)
+		}
+	}
+
+	if lb.body == nil {
+		return nil, fmt.Errorf("%w: lambda must have a body", substraitgo.ErrInvalidExpr)
+	}
+
+	// Push this lambda's params onto context stack before building body.
+	// This allows nested lambdas to validate stepsOut references against
+	// outer lambda parameters.
+	lb.b.pushLambdaContext(lb.params)
+
+	bodyExpr, err := lb.body.BuildExpr()
+
+	// Pop our params from context stack (always, even on error)
+	lb.b.popLambdaContext()
+
+	if err != nil {
+		return nil, err
+	}
+
+	return &Lambda{Parameters: lb.params, Body: bodyExpr}, nil
+}
+
+// BuildExpr implements the Builder interface.
+func (lb *lambdaBuilder) BuildExpr() (Expression, error) {
+	return lb.Build()
+}
+
+// BuildFuncArg implements the FuncArgBuilder interface, allowing lambdas
+// to be passed directly as arguments to function builders.
+func (lb *lambdaBuilder) BuildFuncArg() (types.FuncArg, error) {
+	return lb.Build()
+}
+
 type exprWrapper struct {
 	expression Expression
 	err        error
@@ -407,3 +507,63 @@ func (rb *fieldRefBuilder) BuildFuncArg() (types.FuncArg, error) {
 func (rb *fieldRefBuilder) BuildExpr() (Expression, error) {
 	return rb.Build()
 }
+
+type lambdaParamRefBuilder struct {
+	b        *ExprBuilder
+	ref      Reference
+	stepsOut uint32
+}
+
+func (lpb *lambdaParamRefBuilder) Build() (*FieldReference, error) {
+	// Get the target lambda parameters
+	resolvedType, err := lpb.getLambdaParamType()
+	if err != nil {
+		return nil, err
+	}
+
+	return &FieldReference{
+		Root:      LambdaParameterReference{StepsOut: lpb.stepsOut},
+		Reference: lpb.ref,
+		knownType: resolvedType,
+	}, nil
+}
+
+// getLambdaParamType gets the target lambda parameters struct from the context stack and validates that the parameter exists.
+// Returns an error if the parameter does not exist.
+func (lpb *lambdaParamRefBuilder) getLambdaParamType() (types.Type, error) {
+	// Lambda parameter references must always be StructFieldRef because lambda parameters are always a StructType
+	structRef, ok := lpb.ref.(*StructFieldRef)
+	if !ok {
+		return nil, fmt.Errorf("%w: lambda parameter reference must be StructFieldRef, got %T",
+			substraitgo.ErrInvalidExpr, lpb.ref)
+	}
+
+	// Validate stepsOut - check we have enough outer lambdas
+	if int(lpb.stepsOut) >= len(lpb.b.lambdaContext) {
+		return nil, fmt.Errorf("%w: stepsOut %d references non-existent outer lambda (only %d outer lambdas available)",
+			substraitgo.ErrInvalidExpr, lpb.stepsOut, len(lpb.b.lambdaContext))
+	}
+	targetParams := lpb.b.lambdaContext[len(lpb.b.lambdaContext)-1-int(lpb.stepsOut)]
+
+	// Validate that the parameter field index exists
+	if int(structRef.Field) >= len(targetParams.Types) {
+		return nil, fmt.Errorf("%w: lambda body references parameter %d but lambda only has %d parameters",
+			substraitgo.ErrInvalidExpr, structRef.Field, len(targetParams.Types))
+	}
+
+	// Resolve type using GetType method
+	resolvedType, err := structRef.GetType(targetParams)
+	if err != nil {
+		return nil, fmt.Errorf("%w: cannot resolve lambda parameter reference type: %w",
+			substraitgo.ErrInvalidExpr, err)
+	}
+	return resolvedType, nil
+}
+
+func (lpb *lambdaParamRefBuilder) BuildExpr() (Expression, error) {
+	return lpb.Build()
+}
+
+func (lpb *lambdaParamRefBuilder) BuildFuncArg() (types.FuncArg, error) {
+	return lpb.Build()
+}

expr/expression.go

@@ -323,8 +323,36 @@ func ExprFromProto(e *proto.Expression, baseSchema *types.RecordType, reg Extens
 			return nil, fmt.Errorf("%w: subquery expressions require a subquery converter to be configured", substraitgo.ErrNotImplemented)
 		}
 		return reg.SubqueryFromProto(et.Subquery, baseSchema, reg)
-	}
+	case *proto.Expression_Lambda_:
+		if et.Lambda.Parameters == nil {
+			return nil, fmt.Errorf("%w: lambda parameters cannot be nil", substraitgo.ErrInvalidExpr)
+		}
+		if et.Lambda.Body == nil {
+			return nil, fmt.Errorf("%w: lambda body cannot be nil", substraitgo.ErrInvalidExpr)
+		}
 
+		paramTypes := make([]types.Type, len(et.Lambda.Parameters.Types))
+		for i, pt := range et.Lambda.Parameters.Types {
+			paramTypes[i] = types.TypeFromProto(pt)
+		}
+		params := &types.StructType{
+			Types:            paramTypes,
+			TypeVariationRef: et.Lambda.Parameters.TypeVariationReference,
+			Nullability:      et.Lambda.Parameters.Nullability,
+		}
+
+		if params.Nullability != types.NullabilityRequired {
+			return nil, fmt.Errorf("%w: lambda parameters struct must have NULLABILITY_REQUIRED", substraitgo.ErrInvalidExpr)
+		}
+
+		body, err := ExprFromProto(et.Lambda.Body, baseSchema, reg)
+		if err != nil {
+			return nil, err
+		}
+
+		// TODO (#189): add validation and type resolution for lambda parameter references
+		return &Lambda{Parameters: params, Body: body}, nil
+	}
 	return nil, fmt.Errorf("%w: ExprFromProto: %s", substraitgo.ErrNotImplemented, e)
 }
 

expr/field_reference.go

@@ -29,6 +29,15 @@ type OuterReference uint32
 
 func (OuterReference) isRootRef() {}
 
+// LambdaParameterReference is a root reference type for accessing lambda parameters
+// within a lambda body expression. StepsOut indicates how many lambda boundaries
+// to traverse up (0 = current lambda, 1 = outer lambda, etc.)
+type LambdaParameterReference struct {
+	StepsOut uint32
+}
+
+func (LambdaParameterReference) isRootRef() {}
+
 type ReferenceSegment interface {
 	Reference
 	fmt.Stringer
@@ -619,6 +628,8 @@ func (f *FieldReference) String() string {
 		b.WriteString(")]")
 	} else if outerRef, ok := f.Root.(OuterReference); ok {
 		fmt.Fprintf(&b, "[outerRef:%d]", outerRef)
+	} else if lambdaRef, ok := f.Root.(LambdaParameterReference); ok {
+		fmt.Fprintf(&b, "[lambdaParamRef:%d]", lambdaRef.StepsOut)
 	}
 
 	var typ string
@@ -657,6 +668,12 @@ func (f *FieldReference) ToProtoFieldRef() *proto.Expression_FieldReference {
 					StepsOut: uint32(r),
 				},
 			}
+		case LambdaParameterReference:
+			ret.RootType = &proto.Expression_FieldReference_LambdaParameterReference_{
+				LambdaParameterReference: &proto.Expression_FieldReference_LambdaParameterReference{
+					StepsOut: r.StepsOut,
+				},
+			}
 		}
 	} else {
 		ret.RootType = &proto.Expression_FieldReference_RootReference_{
@@ -696,6 +713,14 @@ func (f *FieldReference) Equals(rhs Expression) bool {
 			if !root.Equals(rhsExpr) {
 				return false
 			}
+		case LambdaParameterReference:
+			rhsRoot, ok := rhs.Root.(LambdaParameterReference)
+			if !ok {
+				return false
+			}
+			if rhsRoot.StepsOut != root.StepsOut {
+				return false
+			}
 		default:
 			if rhs.Root != RootReference {
 				return false
@@ -750,6 +775,8 @@ func FieldReferenceFromProto(p *proto.Expression_FieldReference, baseSchema *typ
 		root = OuterReference(rt.OuterReference.StepsOut)
 	case *proto.Expression_FieldReference_RootReference_:
 		root = RootReference
+	case *proto.Expression_FieldReference_LambdaParameterReference_:
+		root = LambdaParameterReference{StepsOut: rt.LambdaParameterReference.StepsOut}
 	}
 
 	switch rt := p.ReferenceType.(type) {

expr/lambda.go

@@ -0,0 +1,86 @@
+package expr
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/substrait-io/substrait-go/v7/types"
+	proto "github.com/substrait-io/substrait-protobuf/go/substraitpb"
+)
+
+// Lambda represents a lambda expression with parameters and a body.
+type Lambda struct {
+	Parameters *types.StructType // The formal lambda parameters, required to have NULLABILITY_REQUIRED
+	Body       Expression
+}
+
+func (l *Lambda) String() string {
+	var b strings.Builder
+	b.WriteString("(")
+	for i, t := range l.Parameters.Types {
+		if i > 0 {
+			b.WriteString(", ")
+		}
+		fmt.Fprintf(&b, "$%d: %s", i, t)
+	}
+	b.WriteString(") -> ")
+	b.WriteString(l.Body.String())
+	return b.String()
+}
+
+func (l *Lambda) isRootRef() {}
+
+func (l *Lambda) IsScalar() bool {
+	return false
+}
+
+func (l *Lambda) GetType() types.Type {
+	return &types.FuncType{
+		Nullability:    types.NullabilityRequired,
+		ParameterTypes: l.Parameters.Types,
+		ReturnType:     l.Body.GetType(),
+	}
+}
+
+func (l *Lambda) Equals(other Expression) bool {
+	rhs, ok := other.(*Lambda)
+	if !ok {
+		return false
+	}
+	return l.Parameters.Equals(rhs.Parameters) && l.Body.Equals(rhs.Body)
+}
+
+func (l *Lambda) ToProto() *proto.Expression {
+	children := make([]*proto.Type, len(l.Parameters.Types))
+	for i, c := range l.Parameters.Types {
+		children[i] = types.TypeToProto(c)
+	}
+	paramsProto := &proto.Type_Struct{
+		Types:                  children,
+		TypeVariationReference: l.Parameters.TypeVariationRef,
+		Nullability:            l.Parameters.Nullability,
+	}
+
+	return &proto.Expression{
+		RexType: &proto.Expression_Lambda_{
+			Lambda: &proto.Expression_Lambda{
+				Parameters: paramsProto,
+				Body:       l.Body.ToProto(),
+			},
+		},
+	}
+}
+
+func (l *Lambda) ToProtoFuncArg() *proto.FunctionArgument {
+	return &proto.FunctionArgument{
+		ArgType: &proto.FunctionArgument_Value{Value: l.ToProto()},
+	}
+}
+
+func (l *Lambda) Visit(visit VisitFunc) Expression {
+	newBody := visit(l.Body)
+	if newBody == l.Body {
+		return l
+	}
+	return &Lambda{Parameters: l.Parameters, Body: newBody}
+}