diff --git a/isthmus/src/main/java/io/substrait/isthmus/expression/CallConverters.java b/isthmus/src/main/java/io/substrait/isthmus/expression/CallConverters.java index 4b00c15da..7d78bc3e6 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/expression/CallConverters.java +++ b/isthmus/src/main/java/io/substrait/isthmus/expression/CallConverters.java @@ -127,14 +127,18 @@ else if (operand instanceof Expression.StructLiteral }; /** - * Converts Calcite ROW constructors into Substrait {@link Expression.StructLiteral}s. + * Converts Calcite ROW constructors into Substrait {@link Expression.StructLiteral}s, or into + * {@link Expression.NestedStruct}s when the fields are not all literals. * - *

ROW values are always concrete (never null themselves) - if a value is actually null, use - * NullLiteral instead of StructLiteral. Therefore, the resulting StructLiteral always has - * nullable=false. The ROW's type may be nullable (for regular structs) or non-nullable (for UDT - * struct encoding), but the value itself is always concrete. + *

Either way the struct takes its nullability from the ROW's type, which is where Calcite + * keeps it. On a Substrait literal, {@code nullable} marks the literal's type as nullable rather + * than the value as null - a value that is actually null is a NullLiteral - so a nullable ROW of + * literals does not have to give up being a StructLiteral. The UDT struct encoding depends on + * that: it builds a deliberately non-nullable ROW, keeping the user-defined type's own + * nullability in the REINTERPRET target type, and so still arrives here as a StructLiteral. * - *

Each literal's nullability is set to match its field type's nullability. + *

Each literal's nullability is set to match its field type's nullability. Note that Calcite + * makes every field of a nullable record type nullable, so a nullable ROW widens its fields. */ public static final SimpleCallConverter ROW = (call, visitor) -> { @@ -145,10 +149,12 @@ else if (operand instanceof Expression.StructLiteral List operands = call.getOperands().stream().map(visitor).collect(Collectors.toList()); if (!operands.stream().allMatch(expr -> expr instanceof Expression.Literal)) { - throw new IllegalArgumentException("ROW operands must be literals."); + return Expression.NestedStruct.builder() + .nullable(call.getType().isNullable()) + .fields(operands) + .build(); } - // ROW types are never nullable (struct literals are always concrete values). // Field nullability comes from individual field types, so match literal nullability // to field type nullability. List fieldTypes = call.getType().getFieldList(); @@ -162,9 +168,7 @@ else if (operand instanceof Expression.StructLiteral }) .collect(Collectors.toList()); - // Struct literals are always concrete values (never null). - // For UDT struct literals, struct-level nullability is in the REINTERPRET target type. - return ExpressionCreator.struct(false, literals); + return ExpressionCreator.struct(call.getType().isNullable(), literals); }; /** diff --git a/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java b/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java index e3c42d58e..0174b5eed 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java +++ b/isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java @@ -437,7 +437,19 @@ public RexNode visit(Expression.StructLiteral expr, Context context) throws Runt public RexNode visit(Expression.ListLiteral expr, Context context) throws RuntimeException { List args = expr.values().stream().map(l -> l.accept(this, context)).collect(Collectors.toList()); - return rexBuilder.makeCall(SqlStdOperatorTable.ARRAY_VALUE_CONSTRUCTOR, args); + // expr.getType() carries the nullability of the list itself, which Calcite would otherwise + // infer as non-nullable from the elements + RelDataType listType = typeConverter.toCalcite(typeFactory, expr.getType()); + return rexBuilder.makeCall(listType, SqlStdOperatorTable.ARRAY_VALUE_CONSTRUCTOR, args); + } + + @Override + public RexNode visit(Expression.NestedStruct expr, Context context) { + List fieldNodes = + expr.fields().stream().map(f -> f.accept(this, context)).collect(Collectors.toList()); + // expr.getType() carries the nullability of the NestedStruct itself + RelDataType structType = typeConverter.toCalcite(typeFactory, expr.getType()); + return rexBuilder.makeCall(structType, SqlStdOperatorTable.ROW, fieldNodes); } @Override @@ -475,7 +487,25 @@ public RexNode visit(Expression.MapLiteral expr, Context context) throws Runtime entry.getKey().accept(this, context), entry.getValue().accept(this, context))) .collect(Collectors.toList()); - return rexBuilder.makeCall(SqlStdOperatorTable.MAP_VALUE_CONSTRUCTOR, args); + // expr.getType() carries the nullability of the map itself, which Calcite would otherwise + // infer as non-nullable from the keys and values + RelDataType mapType = typeConverter.toCalcite(typeFactory, expr.getType()); + return rexBuilder.makeCall(mapType, SqlStdOperatorTable.MAP_VALUE_CONSTRUCTOR, args); + } + + @Override + public RexNode visit(Expression.NestedMap expr, Context context) { + List args = + expr.values().entrySet().stream() + .flatMap( + entry -> + Stream.of( + entry.getKey().accept(this, context), + entry.getValue().accept(this, context))) + .collect(Collectors.toList()); + // expr.getType() carries the nullability of the NestedMap itself + RelDataType mapType = typeConverter.toCalcite(typeFactory, expr.getType()); + return rexBuilder.makeCall(mapType, SqlStdOperatorTable.MAP_VALUE_CONSTRUCTOR, args); } @Override diff --git a/isthmus/src/main/java/io/substrait/isthmus/expression/SqlMapValueConstructorCallConverter.java b/isthmus/src/main/java/io/substrait/isthmus/expression/SqlMapValueConstructorCallConverter.java index 5d2960e90..1cc6c9109 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/expression/SqlMapValueConstructorCallConverter.java +++ b/isthmus/src/main/java/io/substrait/isthmus/expression/SqlMapValueConstructorCallConverter.java @@ -3,21 +3,25 @@ import io.substrait.expression.Expression; import io.substrait.expression.ExpressionCreator; import io.substrait.isthmus.CallConverter; -import java.util.HashMap; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.function.Function; +import java.util.stream.Collectors; import org.apache.calcite.rex.RexCall; import org.apache.calcite.rex.RexNode; import org.apache.calcite.sql.SqlOperator; import org.apache.calcite.sql.fun.SqlMapValueConstructor; /** - * Converts Calcite {@link SqlMapValueConstructor} calls into Substrait map literals. + * Converts Calcite {@link SqlMapValueConstructor} calls into Substrait map expressions. * - *

Expects an even-numbered operand list (key/value pairs) and produces an {@link Expression} map - * literal via {@link ExpressionCreator}. + *

Expects an even-numbered operand list (key/value pairs) and produces an {@link + * Expression.MapLiteral} when every key and value is a literal, and an {@link Expression.NestedMap} + * otherwise. Either way the map takes its nullability from the call's type, which is where Calcite + * keeps it: on a Substrait literal, {@code nullable} marks the literal's type as nullable rather + * than the value as null, so a nullable map of literals is still a MapLiteral. */ public class SqlMapValueConstructorCallConverter implements CallConverter { @@ -26,38 +30,53 @@ public SqlMapValueConstructorCallConverter() {} /** * Attempts to convert a Calcite {@link RexCall} representing a {@link SqlMapValueConstructor} - * into a Substrait map literal. + * into a Substrait map expression. * * @param call The Calcite call to convert. * @param topLevelConverter Function for converting {@link RexNode} operands to Substrait {@link * Expression}s. * @return An {@link Optional} containing the converted {@link Expression} if the operator is a * {@link SqlMapValueConstructor}; otherwise {@link Optional#empty()}. - * @throws ClassCastException if operands converted by {@code topLevelConverter} are not {@link - * Expression.Literal} instances. - * @throws AssertionError if the number of operands is not even (expecting key/value pairs). + * @throws IllegalArgumentException if the number of operands is not even (expecting key/value + * pairs). */ @Override public Optional convert( RexCall call, Function topLevelConverter) { SqlOperator operator = call.getOperator(); if (operator instanceof SqlMapValueConstructor) { - return toMapLiteral(call, topLevelConverter); + return toMap(call, topLevelConverter); } return Optional.empty(); } - private Optional toMapLiteral( + private Optional toMap( RexCall call, Function topLevelConverter) { - List literals = - call.operands.stream() - .map(t -> ((Expression.Literal) topLevelConverter.apply(t))) - .collect(java.util.stream.Collectors.toList()); - Map items = new HashMap<>(); - assert literals.size() % 2 == 0; - for (int i = 0; i < literals.size(); i += 2) { - items.put(literals.get(i), literals.get(i + 1)); + if (call.operands.size() % 2 != 0) { + throw new IllegalArgumentException( + String.format( + "A map value constructor takes key/value pairs, so it must have an even number of" + + " operands, but it has %d.", + call.operands.size())); } - return Optional.of(ExpressionCreator.map(false, items)); + + List expressions = + call.operands.stream().map(topLevelConverter).collect(Collectors.toList()); + + // The maps below are LinkedHashMaps so that the pairs keep the order they were written in. + if (expressions.stream().allMatch(e -> e instanceof Expression.Literal)) { + Map literals = new LinkedHashMap<>(); + for (int i = 0; i < expressions.size(); i += 2) { + literals.put( + (Expression.Literal) expressions.get(i), (Expression.Literal) expressions.get(i + 1)); + } + return Optional.of(ExpressionCreator.map(call.getType().isNullable(), literals)); + } + + Map values = new LinkedHashMap<>(); + for (int i = 0; i < expressions.size(); i += 2) { + values.put(expressions.get(i), expressions.get(i + 1)); + } + return Optional.of(ExpressionCreator.nestedMap(call.getType().isNullable(), values)); } } diff --git a/isthmus/src/test/java/io/substrait/isthmus/NestedExpressionsTest.java b/isthmus/src/test/java/io/substrait/isthmus/NestedExpressionsTest.java index 8fde8f573..4d978df71 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/NestedExpressionsTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/NestedExpressionsTest.java @@ -1,9 +1,11 @@ package io.substrait.isthmus; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; import com.google.protobuf.ByteString; import io.substrait.expression.Expression; +import io.substrait.expression.ExpressionCreator; import io.substrait.expression.ImmutableExpression; import io.substrait.relation.Project; import io.substrait.relation.Rel; @@ -11,7 +13,9 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.Map; import org.apache.calcite.rel.RelNode; +import org.apache.calcite.sql.parser.SqlParseException; import org.junit.jupiter.api.Test; class NestedExpressionsTest extends PlanTestBase { @@ -155,4 +159,256 @@ void nullableNestedListTest() { assertFullRoundTrip(project); } + + @Test + void nestedStructWithLiteralsTest() { + Expression.NestedStruct literalNestedStruct = + Expression.NestedStruct.builder() + .addFields(literalExpression) + .addFields(sb.i32(12)) + .build(); + + Project project = + Project.builder().expressions(List.of(literalNestedStruct)).input(emptyTable).build(); + + RelNode relNode = substraitToCalcite.convert(project); // substrait rel to calcite + Rel substraitRel = SubstraitRelVisitor.convert(relNode, extensions); // calcite to substrait + Expression roundTripped = ((Project) substraitRel).getExpressions().get(0); + assertEquals(ImmutableExpression.StructLiteral.class, roundTripped.getClass()); + Expression.StructLiteral structLiteral = (Expression.StructLiteral) roundTripped; + assertEquals(literalNestedStruct.fields(), structLiteral.fields()); + } + + @Test + void nullableNestedStructWithLiteralsTest() { + // An all-literal struct collapses to a StructLiteral, but its nullability has to survive the + // collapse: on a Substrait literal, nullable describes the type, not a null value. + Expression.NestedStruct literalNestedStruct = + Expression.NestedStruct.builder() + .addFields(literalExpression) + .addFields(sb.i32(12)) + .nullable(true) + .build(); + + Project project = + Project.builder().expressions(List.of(literalNestedStruct)).input(emptyTable).build(); + + RelNode relNode = substraitToCalcite.convert(project); // substrait rel to calcite + Rel substraitRel = SubstraitRelVisitor.convert(relNode, extensions); // calcite to substrait + Expression roundTripped = ((Project) substraitRel).getExpressions().get(0); + assertEquals(ImmutableExpression.StructLiteral.class, roundTripped.getClass()); + assertTrue(((Expression.StructLiteral) roundTripped).nullable()); + } + + @Test + void nullableStructLiteralTest() { + // The same nullability, on a value that is a StructLiteral to begin with. Its fields are + // nullable because Calcite makes every field of a nullable record type nullable. + Expression.StructLiteral structLiteral = + ExpressionCreator.struct(true, ExpressionCreator.i32(true, 7)); + + Project project = + Project.builder().expressions(List.of(structLiteral)).input(emptyTable).build(); + + assertFullRoundTrip(project); + } + + @Test + void nestedStructWithNonLiteralsTest() { + Expression.NestedStruct nonLiteralNestedStruct = + Expression.NestedStruct.builder() + .addFields(nonLiteralExpression) + .addFields(nonLiteralExpression2) + .build(); + + Project project = + Project.builder() + .expressions(List.of(nonLiteralNestedStruct)) + .input(commonTable) + // project only the nestedStruct expression and exclude the 5 input columns + .remap(Rel.Remap.of(Collections.singleton(5))) + .build(); + + assertFullRoundTrip(project); + } + + @Test + void heterogeneouslyTypedNestedStructTest() { + Expression.NestedStruct nestedStruct = + Expression.NestedStruct.builder() + .addFields(nonLiteralExpression) + .addFields(fieldRef1) + .addFields(literalExpression) + .build(); + + Project project = + Project.builder() + .expressions(List.of(nestedStruct)) + .input(commonTable) + .remap(Rel.Remap.of(Collections.singleton(5))) + .build(); + + assertFullRoundTrip(project); + } + + @Test + void nullableNestedStructTest() { + Expression.NestedStruct nestedStruct = + Expression.NestedStruct.builder() + .addFields(nonLiteralExpression) + .addFields(nonLiteralExpression2) + .nullable(true) + .build(); + + Project project = + Project.builder().expressions(List.of(nestedStruct)).input(emptyTable).build(); + + assertFullRoundTrip(project); + } + + @Test + void nestedMapWithLiteralsTest() { + // keys deliberately out of natural order, so that the assertion on key order below would catch + // a map that no longer preserves the order the pairs were written in + Expression.NestedMap literalNestedMap = + Expression.NestedMap.builder() + .putValues(sb.str("zzz"), literalExpression) + .putValues(sb.str("aaa"), literalExpression) + .putValues(sb.str("mmm"), literalExpression) + .build(); + + Project project = + Project.builder().expressions(List.of(literalNestedMap)).input(emptyTable).build(); + + RelNode relNode = substraitToCalcite.convert(project); // substrait rel to calcite + Rel substraitRel = SubstraitRelVisitor.convert(relNode, extensions); // calcite to substrait + Expression roundTripped = ((Project) substraitRel).getExpressions().get(0); + assertEquals(ImmutableExpression.MapLiteral.class, roundTripped.getClass()); + Expression.MapLiteral mapLiteral = (Expression.MapLiteral) roundTripped; + assertEquals(literalNestedMap.values(), mapLiteral.values()); + // Map.equals ignores order, so compare the key sequences directly + assertEquals( + new ArrayList<>(literalNestedMap.values().keySet()), + new ArrayList<>(mapLiteral.values().keySet())); + } + + @Test + void nullableNestedMapWithLiteralsTest() { + // An all-literal map collapses to a MapLiteral, but its nullability has to survive the + // collapse: on a Substrait literal, nullable describes the type, not a null value. + Expression.NestedMap literalNestedMap = + Expression.NestedMap.builder() + .putValues(sb.str("a"), literalExpression) + .putValues(sb.str("b"), literalExpression) + .nullable(true) + .build(); + + Project project = + Project.builder().expressions(List.of(literalNestedMap)).input(emptyTable).build(); + + RelNode relNode = substraitToCalcite.convert(project); // substrait rel to calcite + Rel substraitRel = SubstraitRelVisitor.convert(relNode, extensions); // calcite to substrait + Expression roundTripped = ((Project) substraitRel).getExpressions().get(0); + assertEquals(ImmutableExpression.MapLiteral.class, roundTripped.getClass()); + assertTrue(((Expression.MapLiteral) roundTripped).nullable()); + } + + @Test + void nullableMapLiteralTest() { + // The same nullability, on a value that is a MapLiteral to begin with. + Expression.MapLiteral mapLiteral = + ExpressionCreator.map( + true, Map.of(ExpressionCreator.string(false, "a"), ExpressionCreator.i32(false, 1))); + + Project project = Project.builder().expressions(List.of(mapLiteral)).input(emptyTable).build(); + + assertFullRoundTrip(project); + } + + @Test + void nullableListLiteralTest() { + // And on a ListLiteral, the third of the three literal containers. + Expression.ListLiteral listLiteral = + ExpressionCreator.list(true, ExpressionCreator.i32(false, 1)); + + Project project = Project.builder().expressions(List.of(listLiteral)).input(emptyTable).build(); + + assertFullRoundTrip(project); + } + + @Test + void nestedMapWithNonLiteralsTest() { + Expression.NestedMap nonLiteralNestedMap = + Expression.NestedMap.builder() + .putValues(sb.str("a"), nonLiteralExpression) + .putValues(sb.str("b"), nonLiteralExpression2) + .build(); + + Project project = + Project.builder() + .expressions(List.of(nonLiteralNestedMap)) + .input(commonTable) + // project only the nestedMap expression and exclude the 5 input columns + .remap(Rel.Remap.of(Collections.singleton(5))) + .build(); + + assertFullRoundTrip(project); + } + + @Test + void nestedMapWithFieldReferenceTest() { + Expression.NestedMap nestedMapWithField = + Expression.NestedMap.builder().putValues(fieldRef1, fieldRef2).build(); + + Project project = + Project.builder() + .expressions(List.of(nestedMapWithField)) + .input(commonTable) + .remap(Rel.Remap.of(Collections.singleton(5))) + .build(); + + assertFullRoundTrip(project); + } + + @Test + void nullableNestedMapTest() { + Expression.NestedMap nestedMap = + Expression.NestedMap.builder() + .putValues(sb.str("a"), nonLiteralExpression) + .nullable(true) + .build(); + + Project project = Project.builder().expressions(List.of(nestedMap)).input(emptyTable).build(); + + assertFullRoundTrip(project); + } + + @Test + void nestedStructOfNestedTypesTest() { + Expression.NestedList list = + Expression.NestedList.builder() + .addValues(nonLiteralExpression) + .addValues(nonLiteralExpression2) + .build(); + Expression.NestedMap map = + Expression.NestedMap.builder().putValues(sb.str("a"), nonLiteralExpression).build(); + + Expression.NestedStruct nestedStruct = + Expression.NestedStruct.builder().addFields(list).addFields(map).build(); + + Project project = + Project.builder().expressions(List.of(nestedStruct)).input(emptyTable).build(); + + assertFullRoundTrip(project); + } + + @Test + void rowConstructorFromSqlTest() throws SqlParseException { + assertFullRoundTrip("SELECT ROW(a + 1, b) FROM t", "CREATE TABLE t (a INT, b INT)"); + } + + @Test + void mapConstructorFromSqlTest() throws SqlParseException { + assertFullRoundTrip("SELECT MAP['key', a + 1] FROM t", "CREATE TABLE t (a INT)"); + } } diff --git a/isthmus/src/test/java/io/substrait/isthmus/UserDefinedLiteralRoundtripTest.java b/isthmus/src/test/java/io/substrait/isthmus/UserDefinedLiteralRoundtripTest.java index 4678b4b16..6d80e6985 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/UserDefinedLiteralRoundtripTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/UserDefinedLiteralRoundtripTest.java @@ -227,6 +227,23 @@ void nestedMixedEncodingsRoundTrip() { pointStructLiteral(1, 2), pointAnyLiteral("p2-any"), pointStructLiteral(3, 4))); } + @Test + void structEncodedUdtWithNullableStructFieldRoundTrip() { + // A struct-encoded UDT is a Calcite ROW that CallConverters.REINTERPRET recognises by its + // operand being a StructLiteral. A nullable struct field produces a nullable inner ROW, so + // anything that stops a nullable ROW of literals from converting back to a StructLiteral also + // breaks the enclosing user-defined literal. + assertRoundTrip( + ExpressionCreator.userDefinedLiteralStruct( + false, + NESTED_TYPES_URN, + "point", + Collections.emptyList(), + Arrays.asList( + ExpressionCreator.struct(true, ExpressionCreator.i32(true, 1)), + ExpressionCreator.i32(false, 100)))); + } + @Test void parameterizedUdtRoundTrip() { Type.Parameter typeParam =