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,104 @@ public static ImmutableExpression.NestedList.Builder builder() {
     }
   }
 
+  /**
+   * A nested map expression with one or more key-value pairs.
+   *
+   * <p>The pairs are held as an ordered list rather than a {@code Map} because that is what the
+   * Substrait map expression is: a repeated list of key-value pairs. Two pairs may carry equal
+   * keys, and a {@code Map} would silently drop one of them.
+   *
+   * <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 List<KeyValue> keyValues();
+
+    /**
+     * 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 (keyValues().isEmpty()) {
+        throw new IllegalArgumentException(
+            "To specify an empty map, use ExpressionCreator.emptyMap");
+      }
+      if (keyValues().stream().map(keyValue -> keyValue.key().getType()).distinct().count() > 1) {
+        throw new IllegalArgumentException("All keys in a NestedMap must have the same type");
+      }
+      if (keyValues().stream().map(keyValue -> keyValue.value().getType()).distinct().count() > 1) {
+        throw new IllegalArgumentException("All values in a NestedMap must have the same type");
+      }
+    }
+
+    @Override
+    public Type getType() {
+      KeyValue first = keyValues().get(0);
+      return Type.withNullability(nullable()).map(first.key().getType(), first.value().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();
+    }
+
+    /** A single key-value pair of a {@link NestedMap}. */
+    @Value.Immutable
+    public abstract static class KeyValue {
+      /**
+       * Returns the key expression of this pair.
+       *
+       * @return the key
+       */
+      public abstract Expression key();
+
+      /**
+       * Returns the value expression of this pair.
+       *
+       * @return the value
+       */
+      public abstract Expression value();
+
+      /**
+       * Creates a key-value pair.
+       *
+       * @param key the key expression
+       * @param value the value expression
+       * @return the key-value pair
+       */
+      public static KeyValue of(Expression key, Expression value) {
+        return builder().key(key).value(value).build();
+      }
+
+      /**
+       * Creates a new builder for constructing a KeyValue.
+       *
+       * @return a new builder instance
+       */
+      public static ImmutableExpression.KeyValue.Builder builder() {
+        return ImmutableExpression.KeyValue.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 keyValues the key-value pairs in the nested map, in the order they should be serialized
+   * @return a NestedMap expression
+   */
+  public static Expression.NestedMap nestedMap(
+      boolean nullable, List<Expression.NestedMap.KeyValue> keyValues) {
+    return Expression.NestedMap.builder().nullable(nullable).addAllKeyValues(keyValues).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,23 @@ 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();
+          for (io.substrait.expression.Expression.NestedMap.KeyValue keyValue : expr.keyValues()) {
+            mapBldr.addKeyValues(
+                Expression.Nested.Map.KeyValue.newBuilder()
+                    .setKey(toProto(keyValue.key()))
+                    .setValue(toProto(keyValue.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

@@ -481,13 +481,29 @@ 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:
+        // The pairs are kept in a list, in the order the producer emitted them, so that a map
+        // repeating a key keeps both of its pairs.
+        List<Expression.NestedMap.KeyValue> keyValues =
+            nested.getMap().getKeyValuesList().stream()
+                .map(
+                    keyValue ->
+                        Expression.NestedMap.KeyValue.of(
+                            from(keyValue.getKey()), from(keyValue.getValue())))
+                .collect(Collectors.toList());
+        return ExpressionCreator.nestedMap(nested.getNullable(), keyValues);
       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,6 +10,7 @@
 import io.substrait.expression.FunctionArg;
 import io.substrait.expression.ImmutableExpression;
 import io.substrait.util.EmptyVisitationContext;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
 
@@ -402,6 +403,26 @@ 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;
+    List<Expression.NestedMap.KeyValue> keyValues = new ArrayList<>();
+    for (Expression.NestedMap.KeyValue keyValue : expr.keyValues()) {
+      Optional<Expression> key = keyValue.key().accept(this, context);
+      Optional<Expression> value = keyValue.value().accept(this, context);
+      changed |= !allEmpty(key, value);
+      keyValues.add(
+          Expression.NestedMap.KeyValue.of(
+              key.orElse(keyValue.key()), value.orElse(keyValue.value())));
+    }
+
+    if (!changed) {
+      return Optional.empty();
+    }
+    return Optional.of(Expression.NestedMap.builder().from(expr).keyValues(keyValues).build());
+  }
+
   /**
    * Visits a multi-or-list record.
    *

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

@@ -0,0 +1,137 @@
+package io.substrait.type.proto;
+
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+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 java.util.List;
+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()
+            .addKeyValues(Expression.NestedMap.KeyValue.of(sb.str("a"), literalExpression))
+            .addKeyValues(Expression.NestedMap.KeyValue.of(sb.i32(1), literalExpression));
+    assertThrows(IllegalArgumentException.class, builder::build);
+  }
+
+  @Test
+  void rejectNestedMapWithValuesOfDifferentTypes() {
+    ImmutableExpression.NestedMap.Builder builder =
+        Expression.NestedMap.builder()
+            .addKeyValues(Expression.NestedMap.KeyValue.of(sb.str("a"), literalExpression))
+            .addKeyValues(Expression.NestedMap.KeyValue.of(sb.str("b"), sb.i32(1)));
+    assertThrows(IllegalArgumentException.class, builder::build);
+  }
+
+  @Test
+  void acceptNestedMapWithKeysAndValuesOfSameType() {
+    ImmutableExpression.NestedMap.Builder builder =
+        Expression.NestedMap.builder()
+            .addKeyValues(Expression.NestedMap.KeyValue.of(sb.str("a"), nonLiteralExpression))
+            .addKeyValues(Expression.NestedMap.KeyValue.of(sb.str("b"), sb.i32(12)));
+    assertDoesNotThrow(builder::build);
+
+    verifyRoundTrip(projectOf(builder.build()));
+  }
+
+  @Test
+  void literalNestedMapTest() {
+    Expression.NestedMap literalNestedMap =
+        Expression.NestedMap.builder()
+            .addKeyValues(Expression.NestedMap.KeyValue.of(sb.str("a"), literalExpression))
+            .addKeyValues(Expression.NestedMap.KeyValue.of(sb.str("b"), literalExpression))
+            .build();
+
+    verifyRoundTrip(projectOf(literalNestedMap));
+  }
+
+  @Test
+  void literalNullableNestedMapTest() {
+    Expression.NestedMap literalNestedMap =
+        Expression.NestedMap.builder()
+            .addKeyValues(Expression.NestedMap.KeyValue.of(sb.str("a"), literalExpression))
+            .addKeyValues(Expression.NestedMap.KeyValue.of(sb.str("b"), literalExpression))
+            .nullable(true)
+            .build();
+
+    verifyRoundTrip(projectOf(literalNestedMap));
+  }
+
+  @Test
+  void nonLiteralNestedMapTest() {
+    Expression.NestedMap nonLiteralNestedMap =
+        Expression.NestedMap.builder()
+            .addKeyValues(
+                Expression.NestedMap.KeyValue.of(nonLiteralExpression, nonLiteralExpression))
+            .addKeyValues(Expression.NestedMap.KeyValue.of(sb.i32(12), sb.i32(13)))
+            .build();
+
+    verifyRoundTrip(projectOf(nonLiteralNestedMap));
+  }
+
+  @Test
+  void nestedMapOfNestedMapsTest() {
+    Expression.NestedMap inner =
+        Expression.NestedMap.builder()
+            .addKeyValues(Expression.NestedMap.KeyValue.of(sb.str("a"), sb.i32(1)))
+            .build();
+
+    Expression.NestedMap outer =
+        Expression.NestedMap.builder()
+            .addKeyValues(Expression.NestedMap.KeyValue.of(sb.str("outer"), inner))
+            .build();
+
+    verifyRoundTrip(projectOf(outer));
+  }
+
+  @Test
+  void repeatedKeysNestedMapTest() {
+    // A Substrait map expression is a repeated list of key-value pairs, so the same key may appear
+    // more than once. Both pairs have to survive a round trip.
+    Expression.NestedMap repeatedKeys =
+        Expression.NestedMap.builder()
+            .addKeyValues(Expression.NestedMap.KeyValue.of(sb.i32(1), sb.i32(10)))
+            .addKeyValues(Expression.NestedMap.KeyValue.of(sb.i32(1), sb.i32(20)))
+            .build();
+
+    assertEquals(2, repeatedKeys.keyValues().size());
+    verifyRoundTrip(projectOf(repeatedKeys));
+  }
+
+  @Test
+  void keyValueOrderIsPreservedTest() {
+    // Keys deliberately out of natural order, so that a representation which reorders the pairs
+    // would fail here.
+    List<Expression.NestedMap.KeyValue> keyValues =
+        List.of(
+            Expression.NestedMap.KeyValue.of(sb.str("zzz"), sb.i32(1)),
+            Expression.NestedMap.KeyValue.of(sb.str("aaa"), sb.i32(2)),
+            Expression.NestedMap.KeyValue.of(sb.str("mmm"), sb.i32(3)));
+
+    Expression.NestedMap nestedMap =
+        Expression.NestedMap.builder().addAllKeyValues(keyValues).build();
+
+    assertEquals(keyValues, nestedMap.keyValues());
+    verifyRoundTrip(projectOf(nestedMap));
+  }
+
+  private Project projectOf(Expression expression) {
+    return Project.builder().addExpressions(expression).input(sb.emptyVirtualTableScan()).build();
+  }
+}