diff --git a/core/src/main/java/io/substrait/expression/AbstractExpressionVisitor.java b/core/src/main/java/io/substrait/expression/AbstractExpressionVisitor.java index eeb3b0954..f39a5e164 100644 --- a/core/src/main/java/io/substrait/expression/AbstractExpressionVisitor.java +++ b/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. * diff --git a/core/src/main/java/io/substrait/expression/Expression.java b/core/src/main/java/io/substrait/expression/Expression.java index e394682b0..3c22ef70d 100644 --- a/core/src/main/java/io/substrait/expression/Expression.java +++ b/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. + * + *

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. + * + *

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 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 accept( + ExpressionVisitor 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 { diff --git a/core/src/main/java/io/substrait/expression/ExpressionCreator.java b/core/src/main/java/io/substrait/expression/ExpressionCreator.java index 714b741be..bd4e819c4 100644 --- a/core/src/main/java/io/substrait/expression/ExpressionCreator.java +++ b/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. + * + *

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 keyValues) { + return Expression.NestedMap.builder().nullable(nullable).addAllKeyValues(keyValues).build(); + } + /** * Create a UserDefinedAnyLiteral with google.protobuf.Any representation. * diff --git a/core/src/main/java/io/substrait/expression/ExpressionVisitor.java b/core/src/main/java/io/substrait/expression/ExpressionVisitor.java index f9a8f62ba..e9af20c44 100644 --- a/core/src/main/java/io/substrait/expression/ExpressionVisitor.java +++ b/core/src/main/java/io/substrait/expression/ExpressionVisitor.java @@ -401,6 +401,16 @@ public interface ExpressionVisitor { + 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) { diff --git a/core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java b/core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java index 8c51027d4..d0cfbbd81 100644 --- a/core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java +++ b/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 fields = + nested.getStruct().getFieldsList().stream() + .map(this::from) + .collect(Collectors.toList()); + return ExpressionCreator.nestedStruct(nested.getNullable(), fields); case LIST: List 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 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()); } } diff --git a/core/src/main/java/io/substrait/relation/ExpressionCopyOnWriteVisitor.java b/core/src/main/java/io/substrait/relation/ExpressionCopyOnWriteVisitor.java index 1a8923083..bc4d0228a 100644 --- a/core/src/main/java/io/substrait/relation/ExpressionCopyOnWriteVisitor.java +++ b/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 visit(Expression.NestedList expr, EmptyVisitationCon Expression.NestedList.builder().from(expr).values(expressionList).build()); } + @Override + public Optional visit(Expression.NestedMap expr, EmptyVisitationContext context) + throws E { + boolean changed = false; + List keyValues = new ArrayList<>(); + for (Expression.NestedMap.KeyValue keyValue : expr.keyValues()) { + Optional key = keyValue.key().accept(this, context); + Optional 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. * diff --git a/core/src/test/java/io/substrait/type/proto/NestedMapExpressionTest.java b/core/src/test/java/io/substrait/type/proto/NestedMapExpressionTest.java new file mode 100644 index 000000000..3b2a574be --- /dev/null +++ b/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 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(); + } +} diff --git a/core/src/test/java/io/substrait/type/proto/NestedStructExpressionTest.java b/core/src/test/java/io/substrait/type/proto/NestedStructExpressionTest.java new file mode 100644 index 000000000..f64656909 --- /dev/null +++ b/core/src/test/java/io/substrait/type/proto/NestedStructExpressionTest.java @@ -0,0 +1,68 @@ +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() + .addKeyValues(Expression.NestedMap.KeyValue.of(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(); + } +} diff --git a/examples/substrait-spark/src/main/java/io/substrait/examples/util/ExpressionStringify.java b/examples/substrait-spark/src/main/java/io/substrait/examples/util/ExpressionStringify.java index 07e09b821..dcf8fc8c3 100644 --- a/examples/substrait-spark/src/main/java/io/substrait/examples/util/ExpressionStringify.java +++ b/examples/substrait-spark/src/main/java/io/substrait/examples/util/ExpressionStringify.java @@ -274,6 +274,12 @@ public String visit(Expression.NestedList expr, EmptyVisitationContext context) return ""; } + @Override + public String visit(Expression.NestedMap expr, EmptyVisitationContext context) + throws RuntimeException { + return ""; + } + @Override public String visit(FieldReference expr, EmptyVisitationContext context) throws RuntimeException { StringBuilder sb = new StringBuilder("FieldRef#");