feat(core)!: support nested map expressions and import nested structs

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

@@ -513,6 +513,19 @@ public O visit(Expression.NestedList expr, C context) throws E {
     return visitFallback(expr, context);
   }
 
+  /**
+   * Visits a nested map expression.
+   *
+   * @param expr the nested map
+   * @param context the visitation context
+   * @return the visit result
+   * @throws E if visitation fails
+   */
+  @Override
+  public O visit(Expression.NestedMap expr, C context) throws E {
+    return visitFallback(expr, context);
+  }
+
   /**
    * Visits a field reference.
    *

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

@@ -1819,6 +1819,64 @@ public static ImmutableExpression.NestedList.Builder builder() {
     }
   }
 
+  /**
+   * A nested map expression with one or more key-value pairs.
+   *
+   * <p>Note: This class cannot be used to construct an empty map. To create an empty map, use
+   * {@link ExpressionCreator#emptyMap(boolean, Type, Type)} which returns an {@link
+   * EmptyMapLiteral}.
+   */
+  @Value.Immutable
+  abstract class NestedMap implements Nested {
+    /**
+     * Returns the key-value pairs in this nested map, in the order they were added.
+     *
+     * @return the key-value pairs
+     */
+    public abstract Map<Expression, Expression> values();
+
+    /**
+     * Validates that the nested map is not empty and that all keys, and all values, have a single
+     * common type.
+     */
+    @Value.Check
+    protected void check() {
+      if (values().isEmpty()) {
+        throw new IllegalArgumentException(
+            "To specify an empty map, use ExpressionCreator.emptyMap");
+      }
+      if (values().keySet().stream().map(Expression::getType).distinct().count() > 1) {
+        throw new IllegalArgumentException("All keys in a NestedMap must have the same type");
+      }
+      if (values().values().stream().map(Expression::getType).distinct().count() > 1) {
+        throw new IllegalArgumentException("All values in a NestedMap must have the same type");
+      }
+    }
+
+    @Override
+    public Type getType() {
+      return Type.withNullability(nullable())
+          .map(
+              values().keySet().iterator().next().getType(),
+              values().values().iterator().next().getType());
+    }
+
+    @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);
+    }
+
+    /**
+     * Creates a new builder for constructing a NestedMap.
+     *
+     * @return a new builder instance
+     */
+    public static ImmutableExpression.NestedMap.Builder builder() {
+      return ImmutableExpression.NestedMap.builder();
+    }
+  }
+
   /** Represents a single record (combination of values) in a multi-or-list expression. */
   @Value.Immutable
   abstract class MultiOrListRecord {

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

@@ -571,6 +571,22 @@ public static Expression.NestedStruct nestedStruct(boolean nullable, Expression.
     return Expression.NestedStruct.builder().nullable(nullable).addFields(fields).build();
   }
 
+  /**
+   * Creates a nested map expression with one or more key-value pairs.
+   *
+   * <p>Note: This method cannot be used to construct an empty map. To create an empty map, use
+   * {@link ExpressionCreator#emptyMap(boolean, Type, Type)} which returns an {@link
+   * Expression.EmptyMapLiteral}.
+   *
+   * @param nullable whether the map can be null
+   * @param values the key-value pairs in the nested map
+   * @return a NestedMap expression
+   */
+  public static Expression.NestedMap nestedMap(
+      boolean nullable, Map<Expression, Expression> values) {
+    return Expression.NestedMap.builder().nullable(nullable).putAllValues(values).build();
+  }
+
   /**
    * Create a UserDefinedAnyLiteral with google.protobuf.Any representation.
    *

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

@@ -401,6 +401,16 @@ public interface ExpressionVisitor<R, C extends VisitationContext, E extends Thr
    */
   R visit(Expression.NestedList expr, C context) throws E;
 
+  /**
+   * Visit a nested map.
+   *
+   * @param expr the nested map
+   * @param context visitation context
+   * @return visit result
+   * @throws E on visit failure
+   */
+  R visit(Expression.NestedMap expr, C context) throws E;
+
   /**
    * Visit a field reference.
    *

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

@@ -616,6 +616,24 @@ public Expression visit(
         .build();
   }
 
+  @Override
+  public Expression visit(
+      io.substrait.expression.Expression.NestedMap expr, EmptyVisitationContext context)
+      throws RuntimeException {
+    return nested(
+        bldr -> {
+          Expression.Nested.Map.Builder mapBldr = Expression.Nested.Map.newBuilder();
+          expr.values()
+              .forEach(
+                  (key, value) ->
+                      mapBldr.addKeyValues(
+                          Expression.Nested.Map.KeyValue.newBuilder()
+                              .setKey(toProto(key))
+                              .setValue(toProto(value))));
+          bldr.setMap(mapBldr).setNullable(expr.nullable());
+        });
+  }
+
   @Override
   public Expression visit(FieldReference expr, EmptyVisitationContext context) {
 

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

@@ -21,7 +21,9 @@
 import io.substrait.type.proto.ProtoTypeConverter;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.LinkedHashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -481,13 +483,26 @@ private WindowBound toWindowBound(io.substrait.proto.Expression.WindowFunction.B
    */
   public Expression.Nested from(io.substrait.proto.Expression.Nested nested) {
     switch (nested.getNestedTypeCase()) {
+      case STRUCT:
+        List<Expression> fields =
+            nested.getStruct().getFieldsList().stream()
+                .map(this::from)
+                .collect(Collectors.toList());
+        return ExpressionCreator.nestedStruct(nested.getNullable(), fields);
       case LIST:
         List<Expression> list =
             nested.getList().getValuesList().stream().map(this::from).collect(Collectors.toList());
         return ExpressionCreator.nestedList(nested.getNullable(), list);
+      case MAP:
+        // A LinkedHashMap keeps the key-value pairs in the order the producer emitted them.
+        Map<Expression, Expression> map = new LinkedHashMap<>();
+        for (io.substrait.proto.Expression.Nested.Map.KeyValue keyValue :
+            nested.getMap().getKeyValuesList()) {
+          map.put(from(keyValue.getKey()), from(keyValue.getValue()));
+        }
+        return ExpressionCreator.nestedMap(nested.getNullable(), map);
       default:
-        throw new UnsupportedOperationException(
-            "Unimplemented nested type: " + nested.getNestedTypeCase());
+        throw new IllegalStateException("Unexpected nested type: " + nested.getNestedTypeCase());
     }
   }
 

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

@@ -10,7 +10,9 @@
 import io.substrait.expression.FunctionArg;
 import io.substrait.expression.ImmutableExpression;
 import io.substrait.util.EmptyVisitationContext;
+import java.util.LinkedHashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 
 /**
@@ -402,6 +404,25 @@ public Optional<Expression> visit(Expression.NestedList expr, EmptyVisitationCon
             Expression.NestedList.builder().from(expr).values(expressionList).build());
   }
 
+  @Override
+  public Optional<Expression> visit(Expression.NestedMap expr, EmptyVisitationContext context)
+      throws E {
+    boolean changed = false;
+    // A LinkedHashMap keeps the key-value pairs in their original order.
+    Map<Expression, Expression> values = new LinkedHashMap<>();
+    for (Map.Entry<Expression, Expression> keyValue : expr.values().entrySet()) {
+      Optional<Expression> key = keyValue.getKey().accept(this, context);
+      Optional<Expression> value = keyValue.getValue().accept(this, context);
+      changed |= !allEmpty(key, value);
+      values.put(key.orElse(keyValue.getKey()), value.orElse(keyValue.getValue()));
+    }
+
+    if (!changed) {
+      return Optional.empty();
+    }
+    return Optional.of(Expression.NestedMap.builder().from(expr).values(values).build());
+  }
+
   /**
    * Visits a multi-or-list record.
    *

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

@@ -0,0 +1,99 @@
+package io.substrait.type.proto;
+
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import io.substrait.TestBase;
+import io.substrait.expression.Expression;
+import io.substrait.expression.ImmutableExpression;
+import io.substrait.relation.Project;
+import org.junit.jupiter.api.Test;
+
+class NestedMapExpressionTest extends TestBase {
+  Expression literalExpression = Expression.BoolLiteral.builder().value(true).build();
+  Expression.ScalarFunctionInvocation nonLiteralExpression = sb.add(sb.i32(7), sb.i32(42));
+
+  @Test
+  void rejectEmptyNestedMap() {
+    ImmutableExpression.NestedMap.Builder builder = Expression.NestedMap.builder();
+    assertThrows(IllegalArgumentException.class, builder::build);
+  }
+
+  @Test
+  void rejectNestedMapWithKeysOfDifferentTypes() {
+    ImmutableExpression.NestedMap.Builder builder =
+        Expression.NestedMap.builder()
+            .putValues(sb.str("a"), literalExpression)
+            .putValues(sb.i32(1), literalExpression);
+    assertThrows(IllegalArgumentException.class, builder::build);
+  }
+
+  @Test
+  void rejectNestedMapWithValuesOfDifferentTypes() {
+    ImmutableExpression.NestedMap.Builder builder =
+        Expression.NestedMap.builder()
+            .putValues(sb.str("a"), literalExpression)
+            .putValues(sb.str("b"), sb.i32(1));
+    assertThrows(IllegalArgumentException.class, builder::build);
+  }
+
+  @Test
+  void acceptNestedMapWithKeysAndValuesOfSameType() {
+    ImmutableExpression.NestedMap.Builder builder =
+        Expression.NestedMap.builder()
+            .putValues(sb.str("a"), nonLiteralExpression)
+            .putValues(sb.str("b"), sb.i32(12));
+    assertDoesNotThrow(builder::build);
+
+    verifyRoundTrip(projectOf(builder.build()));
+  }
+
+  @Test
+  void literalNestedMapTest() {
+    Expression.NestedMap literalNestedMap =
+        Expression.NestedMap.builder()
+            .putValues(sb.str("a"), literalExpression)
+            .putValues(sb.str("b"), literalExpression)
+            .build();
+
+    verifyRoundTrip(projectOf(literalNestedMap));
+  }
+
+  @Test
+  void literalNullableNestedMapTest() {
+    Expression.NestedMap literalNestedMap =
+        Expression.NestedMap.builder()
+            .putValues(sb.str("a"), literalExpression)
+            .putValues(sb.str("b"), literalExpression)
+            .nullable(true)
+            .build();
+
+    verifyRoundTrip(projectOf(literalNestedMap));
+  }
+
+  @Test
+  void nonLiteralNestedMapTest() {
+    Expression.NestedMap nonLiteralNestedMap =
+        Expression.NestedMap.builder()
+            .putValues(nonLiteralExpression, nonLiteralExpression)
+            .putValues(sb.i32(12), sb.i32(13))
+            .build();
+
+    verifyRoundTrip(projectOf(nonLiteralNestedMap));
+  }
+
+  @Test
+  void nestedMapOfNestedMapsTest() {
+    Expression.NestedMap inner =
+        Expression.NestedMap.builder().putValues(sb.str("a"), sb.i32(1)).build();
+
+    Expression.NestedMap outer =
+        Expression.NestedMap.builder().putValues(sb.str("outer"), inner).build();
+
+    verifyRoundTrip(projectOf(outer));
+  }
+
+  private Project projectOf(Expression expression) {
+    return Project.builder().addExpressions(expression).input(sb.emptyVirtualTableScan()).build();
+  }
+}

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

@@ -0,0 +1,66 @@
+package io.substrait.type.proto;
+
+import io.substrait.TestBase;
+import io.substrait.expression.Expression;
+import io.substrait.relation.Project;
+import org.junit.jupiter.api.Test;
+
+class NestedStructExpressionTest extends TestBase {
+  Expression literalExpression = Expression.BoolLiteral.builder().value(true).build();
+  Expression.ScalarFunctionInvocation nonLiteralExpression = sb.add(sb.i32(7), sb.i32(42));
+
+  @Test
+  void emptyNestedStructTest() {
+    verifyRoundTrip(projectOf(Expression.NestedStruct.builder().build()));
+  }
+
+  @Test
+  void literalNestedStructTest() {
+    Expression.NestedStruct literalNestedStruct =
+        Expression.NestedStruct.builder()
+            .addFields(literalExpression)
+            .addFields(sb.str("a"))
+            .build();
+
+    verifyRoundTrip(projectOf(literalNestedStruct));
+  }
+
+  @Test
+  void literalNullableNestedStructTest() {
+    Expression.NestedStruct literalNestedStruct =
+        Expression.NestedStruct.builder().addFields(literalExpression).nullable(true).build();
+
+    verifyRoundTrip(projectOf(literalNestedStruct));
+  }
+
+  @Test
+  void heterogeneouslyTypedNestedStructTest() {
+    Expression.NestedStruct nestedStruct =
+        Expression.NestedStruct.builder()
+            .addFields(nonLiteralExpression)
+            .addFields(sb.str("a"))
+            .addFields(literalExpression)
+            .build();
+
+    verifyRoundTrip(projectOf(nestedStruct));
+  }
+
+  @Test
+  void nestedStructOfNestedTypesTest() {
+    Expression.NestedStruct inner =
+        Expression.NestedStruct.builder().addFields(sb.i32(1)).nullable(true).build();
+    Expression.NestedList list =
+        Expression.NestedList.builder().addValues(sb.i32(2)).addValues(sb.i32(3)).build();
+    Expression.NestedMap map =
+        Expression.NestedMap.builder().putValues(sb.str("a"), sb.i32(4)).build();
+
+    Expression.NestedStruct outer =
+        Expression.NestedStruct.builder().addFields(inner).addFields(list).addFields(map).build();
+
+    verifyRoundTrip(projectOf(outer));
+  }
+
+  private Project projectOf(Expression expression) {
+    return Project.builder().addExpressions(expression).input(sb.emptyVirtualTableScan()).build();
+  }
+}