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..8d74a876d 100644 --- a/core/src/main/java/io/substrait/expression/Expression.java +++ b/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. + * + *

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 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 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(); + } + } + /** 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..2b79abdf1 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 values the key-value pairs in the nested map + * @return a NestedMap expression + */ + public static Expression.NestedMap nestedMap( + boolean nullable, Map values) { + return Expression.NestedMap.builder().nullable(nullable).putAllValues(values).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(); + 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) { 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..b35a3cc40 100644 --- a/core/src/main/java/io/substrait/expression/proto/ProtoExpressionConverter.java +++ b/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 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: + // A LinkedHashMap keeps the key-value pairs in the order the producer emitted them. + Map 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()); } } diff --git a/core/src/main/java/io/substrait/relation/ExpressionCopyOnWriteVisitor.java b/core/src/main/java/io/substrait/relation/ExpressionCopyOnWriteVisitor.java index 1a8923083..e8b6345cf 100644 --- a/core/src/main/java/io/substrait/relation/ExpressionCopyOnWriteVisitor.java +++ b/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 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; + // A LinkedHashMap keeps the key-value pairs in their original order. + Map values = new LinkedHashMap<>(); + for (Map.Entry keyValue : expr.values().entrySet()) { + Optional key = keyValue.getKey().accept(this, context); + Optional 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. * 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..ffb113128 --- /dev/null +++ b/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(); + } +} 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..456b396c1 --- /dev/null +++ b/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(); + } +} 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#");