feat: add lambda expression support

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

@@ -448,6 +448,19 @@ public O visit(Expression.IfThen expr, C context) throws E {
     return visitFallback(expr, context);
   }
 
+  /**
+   * Visits a Lambda expression.
+   *
+   * @param expr the Lambda expression
+   * @param context the visitation context
+   * @return the visit result
+   * @throws E if visitation fails
+   */
+  @Override
+  public O visit(Expression.Lambda expr, C context) throws E {
+    return visitFallback(expr, context);
+  }
+
   /**
    * Visits a scalar function invocation.
    *

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

@@ -707,6 +707,31 @@ public <R, C extends VisitationContext, E extends Throwable> R accept(
     }
   }
 
+  @Value.Immutable
+  abstract class Lambda implements Expression {
+    public abstract Type.Struct parameters();
+
+    public abstract Expression body();
+
+    @Override
+    public Type getType() {
+      List<Type> paramTypes = parameters().fields();
+      Type returnType = body().getType();
+
+      return Type.withNullability(false).func(paramTypes, returnType);
+    }
+
+    public static ImmutableExpression.Lambda.Builder builder() {
+      return ImmutableExpression.Lambda.builder();
+    }
+
+    @Override
+    public <R, C extends VisitationContext, E extends Throwable> R accept(
+        ExpressionVisitor<R, C, E> visitor, C context) throws E {
+      return visitor.visit(this, context);
+    }
+  }
+
   /**
    * Base interface for user-defined literals.
    *

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

@@ -311,6 +311,16 @@ public interface ExpressionVisitor<R, C extends VisitationContext, E extends Thr
    */
   R visit(Expression.NestedStruct expr, C context) throws E;
 
+  /**
+   * Visit a Lambda expression.
+   *
+   * @param expr the Lambda expression
+   * @param context visitation context
+   * @return visit result
+   * @throws E on visit failure
+   */
+  R visit(Expression.Lambda expr, C context) throws E;
+
   /**
    * Visit a user-defined any literal.
    *

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

@@ -20,6 +20,8 @@ public abstract class FieldReference implements Expression {
 
   public abstract Optional<Integer> outerReferenceStepsOut();
 
+  public abstract Optional<Integer> lambdaParameterReferenceStepsOut();
+
   @Override
   public Type getType() {
     return type();
@@ -38,13 +40,18 @@ public <R, C extends VisitationContext, E extends Throwable> R accept(
   public boolean isSimpleRootReference() {
     return segments().size() == 1
         && !inputExpression().isPresent()
-        && !outerReferenceStepsOut().isPresent();
+        && !outerReferenceStepsOut().isPresent()
+        && !lambdaParameterReferenceStepsOut().isPresent();
   }
 
   public boolean isOuterReference() {
     return outerReferenceStepsOut().orElse(0) > 0;
   }
 
+  public boolean isLambdaParameterReference() {
+    return lambdaParameterReferenceStepsOut().isPresent();
+  }
+
   public FieldReference dereferenceStruct(int index) {
     Type newType = StructFieldFinder.getReferencedType(type(), index);
     return dereference(newType, StructField.of(index));
@@ -134,6 +141,15 @@ public static FieldReference newInputRelReference(int index, List<Rel> rels) {
             index, currentOffset));
   }
 
+  public static FieldReference newLambdaParameterReference(
+      int paramIndex, Type.Struct lambdaParamsType, int stepsOut) {
+    return ImmutableFieldReference.builder()
+        .addSegments(StructField.of(paramIndex))
+        .type(lambdaParamsType.fields().get(paramIndex))
+        .lambdaParameterReferenceStepsOut(stepsOut)
+        .build();
+  }
+
   public interface ReferenceSegment {
     FieldReference apply(FieldReference reference);
 

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

@@ -373,6 +373,18 @@ public Expression visit(
         });
   }
 
+  @Override
+  public Expression visit(
+      io.substrait.expression.Expression.Lambda expr, EmptyVisitationContext context)
+      throws RuntimeException {
+    return io.substrait.proto.Expression.newBuilder()
+        .setLambda(
+            io.substrait.proto.Expression.Lambda.newBuilder()
+                .setParameters(typeProtoConverter.toProto(expr.parameters()).getStruct())
+                .setBody(expr.body().accept(this, context)))
+        .build();
+  }
+
   @Override
   public Expression visit(
       io.substrait.expression.Expression.UserDefinedAnyLiteral expr,
@@ -603,6 +615,10 @@ public Expression visit(FieldReference expr, EmptyVisitationContext context) {
       out.setOuterReference(
           io.substrait.proto.Expression.FieldReference.OuterReference.newBuilder()
               .setStepsOut(expr.outerReferenceStepsOut().get()));
+    } else if (expr.lambdaParameterReferenceStepsOut().isPresent()) {
+      out.setLambdaParameterReference(
+          io.substrait.proto.Expression.FieldReference.LambdaParameterReference.newBuilder()
+              .setStepsOut(expr.lambdaParameterReferenceStepsOut().get()));
     } else {
       out.setRootReference(Expression.FieldReference.RootReference.getDefaultInstance());
     }

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

@@ -37,6 +37,7 @@ public class ProtoExpressionConverter {
   private final Type.Struct rootType;
   private final ProtoTypeConverter protoTypeConverter;
   private final ProtoRelConverter protoRelConverter;
+  private final List<Type.Struct> lambdaParameterStack = new ArrayList<>();
 
   public ProtoExpressionConverter(
       ExtensionLookup lookup,
@@ -75,6 +76,26 @@ public FieldReference from(io.substrait.proto.Expression.FieldReference referenc
             reference.getDirectReference().getStructField().getField(),
             rootType,
             reference.getOuterReference().getStepsOut());
+      case LAMBDA_PARAMETER_REFERENCE:
+        {
+          io.substrait.proto.Expression.FieldReference.LambdaParameterReference lambdaParamRef =
+              reference.getLambdaParameterReference();
+
+          int stepsOut = lambdaParamRef.getStepsOut();
+          int lambdaIndex = lambdaParameterStack.size() - 1 - stepsOut;
+          if (lambdaIndex < 0 || lambdaIndex >= lambdaParameterStack.size()) {
+            throw new IllegalArgumentException(
+                String.format(
+                    "Lambda parameter reference with stepsOut=%d is invalid (current depth: %d)",
+                    stepsOut, lambdaParameterStack.size()));
+          }
+
+          Type.Struct lambdaParameters = lambdaParameterStack.get(lambdaIndex);
+          return FieldReference.newLambdaParameterReference(
+              reference.getDirectReference().getStructField().getField(),
+              lambdaParameters,
+              stepsOut);
+        }
       case ROOTTYPE_NOT_SET:
       default:
         throw new IllegalArgumentException("Unhandled type: " + reference.getRootTypeCase());
@@ -260,6 +281,27 @@ public Type visit(Type.Struct type) throws RuntimeException {
           }
         }
 
+      case LAMBDA:
+        {
+          io.substrait.proto.Expression.Lambda protoLambda = expr.getLambda();
+          Type.Struct parameters =
+              (Type.Struct)
+                  protoTypeConverter.from(
+                      io.substrait.proto.Type.newBuilder()
+                          .setStruct(protoLambda.getParameters())
+                          .build());
+
+          lambdaParameterStack.add(parameters);
+
+          Expression body;
+          try {
+            body = from(protoLambda.getBody());
+          } finally {
+            lambdaParameterStack.remove(lambdaParameterStack.size() - 1);
+          }
+
+          return Expression.Lambda.builder().parameters(parameters).body(body).build();
+        }
       // TODO enum.
       case ENUM:
         throw new UnsupportedOperationException("Unsupported type: " + expr.getRexTypeCase());

core/src/main/java/io/substrait/extension/DefaultExtensionCatalog.java

@@ -50,6 +50,9 @@ public class DefaultExtensionCatalog {
   /** Extension identifier for set functions. */
   public static final String FUNCTIONS_SET = "extension:io.substrait:functions_set";
 
+  /** Extension identifier for list functions. */
+  public static final String FUNCTIONS_LIST = "extension:io.substrait:functions_list";
+
   /** Extension identifier for string functions. */
   public static final String FUNCTIONS_STRING = "extension:io.substrait:functions_string";
 

core/src/main/java/io/substrait/function/ExtendedTypeCreator.java

@@ -16,4 +16,6 @@ public interface ExtendedTypeCreator<T, I> {
   T listE(T type);
 
   T mapE(T key, T value);
+
+  T funcE(Iterable<? extends T> parameterTypes, T returnType);
 }

core/src/main/java/io/substrait/function/ParameterizedType.java

@@ -200,6 +200,23 @@ <R, E extends Throwable> R accept(final ParameterizedTypeVisitor<R, E> parameter
     }
   }
 
+  @Value.Immutable
+  abstract class Func extends BaseParameterizedType implements NullableType {
+    public abstract java.util.List<ParameterizedType> parameterTypes();
+
+    public abstract ParameterizedType returnType();
+
+    public static ImmutableParameterizedType.Func.Builder builder() {
+      return ImmutableParameterizedType.Func.builder();
+    }
+
+    @Override
+    <R, E extends Throwable> R accept(final ParameterizedTypeVisitor<R, E> parameterizedTypeVisitor)
+        throws E {
+      return parameterizedTypeVisitor.visit(this);
+    }
+  }
+
   @Value.Immutable
   abstract class ListType extends BaseParameterizedType implements NullableType {
     public abstract ParameterizedType name();

core/src/main/java/io/substrait/function/ParameterizedTypeCreator.java

@@ -96,6 +96,16 @@ public ParameterizedType listE(ParameterizedType type) {
     return ParameterizedType.ListType.builder().nullable(nullable).name(type).build();
   }
 
+  @Override
+  public ParameterizedType funcE(
+      Iterable<? extends ParameterizedType> parameterTypes, ParameterizedType returnType) {
+    return ParameterizedType.Func.builder()
+        .nullable(nullable)
+        .addAllParameterTypes(parameterTypes)
+        .returnType(returnType)
+        .build();
+  }
+
   @Override
   public ParameterizedType mapE(ParameterizedType key, ParameterizedType value) {
     return ParameterizedType.Map.builder().nullable(nullable).key(key).value(value).build();

core/src/main/java/io/substrait/function/ParameterizedTypeVisitor.java

@@ -29,6 +29,8 @@ public interface ParameterizedTypeVisitor<R, E extends Throwable> extends TypeVi
 
   R visit(ParameterizedType.StringLiteral stringLiteral) throws E;
 
+  R visit(ParameterizedType.Func expr) throws E;
+
   abstract class ParameterizedTypeThrowsVisitor<R, E extends Throwable>
       extends TypeVisitor.TypeThrowsVisitor<R, E> implements ParameterizedTypeVisitor<R, E> {
 
@@ -100,5 +102,10 @@ public R visit(ParameterizedType.Map expr) throws E {
     public R visit(ParameterizedType.StringLiteral stringLiteral) throws E {
       throw t();
     }
+
+    @Override
+    public R visit(ParameterizedType.Func expr) throws E {
+      throw t();
+    }
   }
 }

core/src/main/java/io/substrait/function/TypeExpression.java

@@ -191,6 +191,22 @@ <R, E extends Throwable> R acceptE(final TypeExpressionVisitor<R, E> visitor) th
     }
   }
 
+  @Value.Immutable
+  abstract class Func extends BaseTypeExpression implements NullableType {
+    public abstract java.util.List<TypeExpression> parameterTypes();
+
+    public abstract TypeExpression returnType();
+
+    public static ImmutableTypeExpression.Func.Builder builder() {
+      return ImmutableTypeExpression.Func.builder();
+    }
+
+    @Override
+    <R, E extends Throwable> R acceptE(final TypeExpressionVisitor<R, E> visitor) throws E {
+      return visitor.visit(this);
+    }
+  }
+
   @Value.Immutable
   abstract class BinaryOperation extends BaseTypeExpression {
     public enum OpType {

core/src/main/java/io/substrait/function/TypeExpressionCreator.java

@@ -82,6 +82,16 @@ public TypeExpression mapE(TypeExpression key, TypeExpression value) {
     return TypeExpression.Map.builder().nullable(nullable).key(key).value(value).build();
   }
 
+  @Override
+  public TypeExpression funcE(
+      Iterable<? extends TypeExpression> parameterTypes, TypeExpression returnType) {
+    return TypeExpression.Func.builder()
+        .nullable(nullable)
+        .addAllParameterTypes(parameterTypes)
+        .returnType(returnType)
+        .build();
+  }
+
   public static class Assign {
     String name;
     TypeExpression expr;

core/src/main/java/io/substrait/function/TypeExpressionVisitor.java

@@ -24,6 +24,8 @@ public interface TypeExpressionVisitor<R, E extends Throwable>
 
   R visit(TypeExpression.Map expr) throws E;
 
+  R visit(TypeExpression.Func expr) throws E;
+
   R visit(TypeExpression.BinaryOperation expr) throws E;
 
   R visit(TypeExpression.NotOperation expr) throws E;
@@ -97,6 +99,11 @@ public R visit(TypeExpression.Map expr) throws E {
       throw t();
     }
 
+    @Override
+    public R visit(TypeExpression.Func expr) throws E {
+      throw t();
+    }
+
     @Override
     public R visit(TypeExpression.BinaryOperation expr) throws E {
       throw t();

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

@@ -432,6 +432,18 @@ public Optional<Expression> visit(
             .build());
   }
 
+  @Override
+  public Optional<Expression> visit(Expression.Lambda lambda, EmptyVisitationContext context)
+      throws E {
+    Optional<Expression> newBody = lambda.body().accept(this, context);
+
+    if (allEmpty(newBody)) {
+      return Optional.empty();
+    }
+    return Optional.of(
+        Expression.Lambda.builder().from(lambda).body(newBody.orElse(lambda.body())).build());
+  }
+
   // utilities
 
   protected Optional<List<Expression>> visitExprList(

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

@@ -242,5 +242,11 @@ public Integer visit(Type.Map type) throws RuntimeException {
     public Integer visit(Type.UserDefined type) throws RuntimeException {
       return 0;
     }
+
+    @Override
+    public Integer visit(Type.Func type) throws RuntimeException {
+      return type.parameterTypes().stream().mapToInt(p -> p.accept(this)).sum()
+          + type.returnType().accept(this);
+    }
   }
 }

core/src/main/java/io/substrait/type/StringTypeVisitor.java

@@ -150,4 +150,13 @@ public String visit(Type.Map type) throws RuntimeException {
   public String visit(Type.UserDefined type) throws RuntimeException {
     return String.format("u!%s%s", type.name(), n(type));
   }
+
+  @Override
+  public String visit(Type.Func type) throws RuntimeException {
+    return String.format(
+        "func%s<%s -> %s>",
+        n(type),
+        type.parameterTypes().stream().map(t -> t.accept(this)).collect(Collectors.joining(", ")),
+        type.returnType().accept(this));
+  }
 }