diff --git a/isthmus/src/main/java/io/substrait/isthmus/expression/FieldSelectionConverter.java b/isthmus/src/main/java/io/substrait/isthmus/expression/FieldSelectionConverter.java index 4893d1ef5..b18a53dba 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/expression/FieldSelectionConverter.java +++ b/isthmus/src/main/java/io/substrait/isthmus/expression/FieldSelectionConverter.java @@ -12,6 +12,9 @@ import org.apache.calcite.rex.RexLiteral; import org.apache.calcite.rex.RexNode; import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.SqlOperator; +import org.apache.calcite.sql.fun.SqlItemOperator; +import org.apache.calcite.sql.type.SqlTypeName; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -43,7 +46,7 @@ public FieldSelectionConverter(TypeConverter typeConverter) { * * * @@ -65,19 +68,26 @@ public Optional convert( LOGGER .atWarn() .log( - "Found item operator without literal kind/type. This isn't handled well. Reference was {} with toString {}.", + "Found item operator without literal kind/type. This isn't handled well. Reference" + + " was {} with toString {}.", reference.getKind().name(), reference); return Optional.empty(); } - Literal literal = (new LiteralConverter(typeConverter)).convert((RexLiteral) reference); + SqlTypeName containerType = toDereference.getType().getSqlTypeName(); + // A Substrait list reference returns null for an out-of-range index, so it cannot preserve + // the error behavior of OFFSET or ORDINAL. Decline before the array operand is converted. + if (containerType == SqlTypeName.ARRAY && !isSafeArrayOperator(call.getOperator())) { + return Optional.empty(); + } Expression input = topLevelConverter.apply(toDereference); - switch (toDereference.getType().getSqlTypeName()) { + switch (containerType) { case ROW: { + Literal literal = (new LiteralConverter(typeConverter)).convert((RexLiteral) reference); Optional index = toInt(literal); if (index.isEmpty()) { return Optional.empty(); @@ -90,20 +100,33 @@ public Optional convert( } case ARRAY: { - Optional index = toInt(literal); - if (index.isEmpty()) { - return Optional.empty(); + SqlItemOperator operator = (SqlItemOperator) call.getOperator(); + long offset = Integer.MAX_VALUE; + if (!((RexLiteral) reference).isNull()) { + Literal literal = (new LiteralConverter(typeConverter)).convert((RexLiteral) reference); + Optional index = toLong(literal); + if (index.isEmpty()) { + return Optional.empty(); + } + // Negative Substrait offsets count from the end, whereas Calcite returns null below + // the operator's base. Substrait lists have at most Integer.MAX_VALUE elements, so that + // zero-based offset cannot select an element. Keep the reference to evaluate the + // array operand and preserve its element type, including for a null index. + if (index.get() >= operator.offset) { + offset = Math.min(index.get() - operator.offset, Integer.MAX_VALUE); + } } if (input instanceof FieldReference) { - return Optional.of(((FieldReference) input).dereferenceList(index.get())); + return Optional.of(((FieldReference) input).dereferenceList((int) offset)); } else { - return Optional.of(FieldReference.newListReference(index.get(), input)); + return Optional.of(FieldReference.newListReference((int) offset, input)); } } case MAP: { + Literal literal = (new LiteralConverter(typeConverter)).convert((RexLiteral) reference); Optional mapKey = toString(literal); if (mapKey.isEmpty()) { return Optional.empty(); @@ -121,6 +144,14 @@ public Optional convert( return Optional.empty(); } + private static boolean isSafeArrayOperator(SqlOperator operator) { + if (!(operator instanceof SqlItemOperator)) { + return false; + } + SqlItemOperator itemOperator = (SqlItemOperator) operator; + return itemOperator.safe && (itemOperator.offset == 0 || itemOperator.offset == 1); + } + /** * Converts a numeric literal to an integer index. * @@ -141,6 +172,13 @@ private Optional toInt(Expression.Literal l) { return Optional.empty(); } + private Optional toLong(Expression.Literal literal) { + if (literal instanceof Expression.I64Literal) { + return Optional.of(((Expression.I64Literal) literal).value()); + } + return toInt(literal).map(Integer::longValue); + } + /** * Converts a fixed-char literal to a string key. * diff --git a/isthmus/src/test/java/io/substrait/isthmus/NestedStructQueryTest.java b/isthmus/src/test/java/io/substrait/isthmus/NestedStructQueryTest.java index f1400526a..dbcc66592 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/NestedStructQueryTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/NestedStructQueryTest.java @@ -204,7 +204,7 @@ public RelDataType getRowType(RelDataTypeFactory factory) { + " field: 1 # a\n" + " child {\n" + " list_element {\n" - + " offset: 1\n" + + " offset: 0\n" + " }\n" + " }\n" + " }\n" @@ -244,13 +244,13 @@ public RelDataType getRowType(RelDataTypeFactory factory) { + " field: 1 # a\n" + " child {\n" + " list_element {\n" - + " offset: 1\n" + + " offset: 0\n" + " child {\n" + " list_element {\n" - + " offset: 2\n" + + " offset: 1\n" + " child {\n" + " list_element {\n" - + " offset: 3\n" + + " offset: 2\n" + " }\n" + " }\n" + " }\n" @@ -300,7 +300,7 @@ public RelDataType getRowType(RelDataTypeFactory factory) { + " field: 0 # .b\n" + " child {\n" + " list_element {\n" - + " offset: 2\n" + + " offset: 1\n" + " child {\n" + " struct_field {\n" + " field: 0 # .c\n" diff --git a/isthmus/src/test/java/io/substrait/isthmus/expression/FieldSelectionConverterTest.java b/isthmus/src/test/java/io/substrait/isthmus/expression/FieldSelectionConverterTest.java new file mode 100644 index 000000000..7f81f0359 --- /dev/null +++ b/isthmus/src/test/java/io/substrait/isthmus/expression/FieldSelectionConverterTest.java @@ -0,0 +1,364 @@ +package io.substrait.isthmus.expression; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import io.substrait.expression.Expression; +import io.substrait.expression.FieldReference; +import io.substrait.expression.proto.ExpressionProtoConverter; +import io.substrait.expression.proto.ProtoExpressionConverter; +import io.substrait.extension.DefaultExtensionCatalog; +import io.substrait.extension.ExtensionCollector; +import io.substrait.isthmus.TypeConverter; +import io.substrait.isthmus.UserTypeMapper; +import io.substrait.isthmus.utils.UserTypeFactory; +import io.substrait.type.Type; +import io.substrait.type.TypeCreator; +import io.substrait.util.EmptyVisitationContext; +import java.math.BigDecimal; +import java.util.List; +import java.util.Optional; +import java.util.stream.Stream; +import org.apache.calcite.DataContexts; +import org.apache.calcite.jdbc.JavaTypeFactoryImpl; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexCall; +import org.apache.calcite.rex.RexExecutable; +import org.apache.calcite.rex.RexExecutorImpl; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.SqlOperator; +import org.apache.calcite.sql.fun.SqlLibraryOperators; +import org.apache.calcite.sql.fun.SqlStdOperatorTable; +import org.apache.calcite.sql.type.SqlTypeName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.CsvSource; +import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.ValueSource; + +class FieldSelectionConverterTest { + private final JavaTypeFactoryImpl typeFactory = new JavaTypeFactoryImpl(); + private final RexBuilder rexBuilder = new RexBuilder(typeFactory); + private final RelDataType intType = typeFactory.createSqlType(SqlTypeName.INTEGER); + private final RexExpressionConverter converter = new RexExpressionConverter(); + + @ParameterizedTest + @CsvSource({"1, 0, 10", "2, 1, 20", "3, 2, 30", "4, 3,", "0,,", "-1,,", "-2,,"}) + void itemPreservesCalciteIndexing(int index, Integer expectedOffset, Integer expectedValue) { + RexNode call = rexBuilder.makeCall(SqlStdOperatorTable.ITEM, array(), integer(index)); + RexExecutable executable = + RexExecutorImpl.getExecutable(rexBuilder, List.of(call), typeFactory.builder().build()); + executable.setDataContext(DataContexts.EMPTY); + assertEquals(expectedValue, executable.execute()[0]); + + Expression converted = call.accept(converter); + assertEquals( + expectedOffset == null ? Integer.MAX_VALUE : expectedOffset, listOffset(converted)); + } + + @ParameterizedTest + @MethodSource("safeIndexing") + void safeOperatorsUseTheirOwnBase(SqlOperator operator, long index, int expectedOffset) { + Expression expression = + rexBuilder.makeCall(operator, array(), integer(index)).accept(converter); + assertEquals(expectedOffset, listOffset(expression)); + + // The same offset is used when dereferencing an array column rather than an expression. + RexNode column = rexBuilder.makeInputRef(array().getType(), 0); + Expression columnExpression = + rexBuilder.makeCall(operator, column, integer(index)).accept(converter); + io.substrait.proto.Expression proto = toProto(columnExpression); + assertEquals( + expectedOffset, + listOffset(proto.getSelection().getDirectReference().getStructField().getChild())); + } + + private static Stream safeIndexing() { + return Stream.of( + Arguments.of(SqlStdOperatorTable.ITEM, 1L, 0), + Arguments.of(SqlStdOperatorTable.ITEM, 2147483648L, Integer.MAX_VALUE), + Arguments.of(SqlLibraryOperators.SAFE_OFFSET, 0L, 0), + Arguments.of(SqlLibraryOperators.SAFE_OFFSET, 1L, 1), + Arguments.of(SqlLibraryOperators.SAFE_OFFSET, 3L, 3), + Arguments.of(SqlLibraryOperators.SAFE_OFFSET, (long) Integer.MAX_VALUE, Integer.MAX_VALUE), + Arguments.of(SqlLibraryOperators.SAFE_ORDINAL, 1L, 0), + Arguments.of(SqlLibraryOperators.SAFE_ORDINAL, 3L, 2), + Arguments.of(SqlLibraryOperators.SAFE_ORDINAL, 2147483648L, Integer.MAX_VALUE)); + } + + @ParameterizedTest + @MethodSource("invalidLowIndexes") + void invalidLowIndexesAreNull(SqlOperator operator, long index) { + Expression converted = rexBuilder.makeCall(operator, array(), integer(index)).accept(converter); + assertEquals(Integer.MAX_VALUE, listOffset(converted)); + } + + private static Stream invalidLowIndexes() { + return Stream.of( + Arguments.of(SqlStdOperatorTable.ITEM, 0L), + Arguments.of(SqlStdOperatorTable.ITEM, Long.MIN_VALUE), + Arguments.of(SqlLibraryOperators.SAFE_OFFSET, -1L), + Arguments.of(SqlLibraryOperators.SAFE_OFFSET, Long.MIN_VALUE), + Arguments.of(SqlLibraryOperators.SAFE_ORDINAL, 0L), + Arguments.of(SqlLibraryOperators.SAFE_ORDINAL, -1L)); + } + + @Test + void nullIndexIsNull() { + Expression converted = + rexBuilder + .makeCall(SqlStdOperatorTable.ITEM, array(), rexBuilder.makeNullLiteral(intType)) + .accept(converter); + assertEquals(Integer.MAX_VALUE, listOffset(converted)); + } + + @Test + void untypedNullIndexIsNull() { + Expression converted = + rexBuilder + .makeCall( + SqlStdOperatorTable.ITEM, + array(), + rexBuilder.makeNullLiteral(typeFactory.createSqlType(SqlTypeName.NULL))) + .accept(converter); + assertEquals(Integer.MAX_VALUE, listOffset(converted)); + verifyRoundTrip(converted, TypeConverter.DEFAULT.toSubstrait(array().getType())); + } + + @ParameterizedTest + @MethodSource("nullOrInvalidIndexes") + void retainsThrowingArray(SqlOperator operator, Integer index) { + RexNode cast = + rexBuilder.makeAbstractCast(intType, rexBuilder.makeLiteral("not an integer"), false); + RexNode throwingArray = rexBuilder.makeCall(SqlStdOperatorTable.ARRAY_VALUE_CONSTRUCTOR, cast); + RexNode nestedArray = + rexBuilder.makeCall(SqlStdOperatorTable.ARRAY_VALUE_CONSTRUCTOR, throwingArray); + RexNode nestedSelection = + rexBuilder.makeCall(SqlStdOperatorTable.ITEM, nestedArray, integer(1)); + + for (RexNode input : List.of(throwingArray, nestedSelection)) { + RexNode call = rexBuilder.makeCall(operator, input, nullableIndex(index)); + RexExecutable executable = + RexExecutorImpl.getExecutable(rexBuilder, List.of(call), typeFactory.builder().build()); + executable.setDataContext(DataContexts.EMPTY); + assertThrows(NumberFormatException.class, executable::execute); + FieldReference converted = assertInstanceOf(FieldReference.class, call.accept(converter)); + assertEquals( + Integer.MAX_VALUE, ((FieldReference.ListElement) converted.segments().get(0)).offset()); + Expression source = input.accept(converter); + assertEquals( + source instanceof FieldReference + ? ((FieldReference) source).inputExpression().orElseThrow() + : source, + converted.inputExpression().orElseThrow()); + } + } + + @ParameterizedTest + @MethodSource("nullOrInvalidIndexes") + void retainsNonliteralArray(SqlOperator operator, Integer index) { + RexNode text = rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.VARCHAR, 20), 0); + RexNode cast = rexBuilder.makeAbstractCast(intType, text, false); + RexNode array = rexBuilder.makeCall(SqlStdOperatorTable.ARRAY_VALUE_CONSTRUCTOR, cast); + RexNode call = rexBuilder.makeCall(operator, array, nullableIndex(index)); + + FieldReference converted = assertInstanceOf(FieldReference.class, call.accept(converter)); + assertEquals(Integer.MAX_VALUE, listOffset(converted)); + assertEquals(array.accept(converter), converted.inputExpression().orElseThrow()); + } + + @ParameterizedTest + @MethodSource("nullOrInvalidIndexes") + void literalAndColumnArraysRetainReferences(SqlOperator operator, Integer index) { + RexNode column = rexBuilder.makeInputRef(array().getType(), 0); + for (RexNode input : List.of(array(), column)) { + Expression converted = + rexBuilder.makeCall(operator, input, nullableIndex(index)).accept(converter); + FieldReference reference = assertInstanceOf(FieldReference.class, converted); + assertEquals( + Integer.MAX_VALUE, ((FieldReference.ListElement) reference.segments().get(0)).offset()); + verifyRoundTrip(converted, TypeConverter.DEFAULT.toSubstrait(input.getType())); + } + } + + private static Stream nullOrInvalidIndexes() { + return Stream.of( + Arguments.of(SqlStdOperatorTable.ITEM, 0), + Arguments.of(SqlStdOperatorTable.ITEM, -1), + Arguments.of(SqlStdOperatorTable.ITEM, null), + Arguments.of(SqlLibraryOperators.SAFE_OFFSET, -1), + Arguments.of(SqlLibraryOperators.SAFE_OFFSET, null), + Arguments.of(SqlLibraryOperators.SAFE_ORDINAL, 0), + Arguments.of(SqlLibraryOperators.SAFE_ORDINAL, null)); + } + + private RexNode nullableIndex(Integer index) { + return index == null ? rexBuilder.makeNullLiteral(intType) : integer(index); + } + + @Test + void invalidIndexesRemainReferencesInsideCollections() { + RexNode column = rexBuilder.makeInputRef(array().getType(), 0); + RexNode missing = rexBuilder.makeCall(SqlStdOperatorTable.ITEM, column, integer(0)); + Expression list = + rexBuilder + .makeCall(SqlStdOperatorTable.ARRAY_VALUE_CONSTRUCTOR, integer(1), missing) + .accept(converter); + assertInstanceOf(Expression.NestedList.class, list); + verifyRoundTrip(list, TypeConverter.DEFAULT.toSubstrait(column.getType())); + Expression map = + rexBuilder + .makeCall( + SqlStdOperatorTable.MAP_VALUE_CONSTRUCTOR, missing, rexBuilder.makeLiteral("x")) + .accept(converter); + assertInstanceOf(Expression.NestedMap.class, map); + verifyRoundTrip(map, TypeConverter.DEFAULT.toSubstrait(column.getType())); + + RelDataType nestedType = + typeFactory.createArrayType(typeFactory.createArrayType(array().getType(), -1), -1); + RexNode nested = rexBuilder.makeInputRef(nestedType, 0); + for (int index : new int[] {0, 1, 0}) { + nested = rexBuilder.makeCall(SqlStdOperatorTable.ITEM, nested, integer(index)); + } + FieldReference reference = assertInstanceOf(FieldReference.class, nested.accept(converter)); + assertEquals(4, reference.segments().size()); + assertEquals( + Integer.MAX_VALUE, + assertInstanceOf(FieldReference.ListElement.class, reference.segments().get(0)).offset()); + assertEquals( + 0, + assertInstanceOf(FieldReference.ListElement.class, reference.segments().get(1)).offset()); + assertEquals( + Integer.MAX_VALUE, + assertInstanceOf(FieldReference.ListElement.class, reference.segments().get(2)).offset()); + verifyRoundTrip(reference, TypeConverter.DEFAULT.toSubstrait(nestedType)); + } + + @Test + void rowElementTypesRemainIdenticalForValidAndInvalidIndexes() { + RelDataType row = typeFactory.builder().add("x", intType).build(); + RexNode column = rexBuilder.makeInputRef(typeFactory.createArrayType(row, -1), 0); + RexNode first = rexBuilder.makeCall(SqlStdOperatorTable.ITEM, column, integer(1)); + RexNode missing = rexBuilder.makeCall(SqlStdOperatorTable.ITEM, column, integer(0)); + Expression list = + rexBuilder + .makeCall(SqlStdOperatorTable.ARRAY_VALUE_CONSTRUCTOR, missing, first) + .accept(converter); + assertEquals(first.accept(converter).getType(), missing.accept(converter).getType()); + assertInstanceOf(Expression.NestedList.class, list); + verifyRoundTrip(list, TypeConverter.DEFAULT.toSubstrait(column.getType())); + } + + @Test + void requiredUserDefinedElementsCanBeAccessedWithInvalidIndexes() { + UserTypeFactory factory = new UserTypeFactory("extension:test:array", "element"); + TypeConverter types = + new TypeConverter( + new UserTypeMapper() { + @Override + public Type toSubstrait(RelDataType type) { + return factory.isTypeFromFactory(type) + ? factory.createSubstrait(type.isNullable(), factory.getTypeParameters(type)) + : null; + } + + @Override + public RelDataType toCalcite(Type.UserDefined type) { + return factory.createCalcite(type.nullable(), type.typeParameters()); + } + }); + RexExpressionConverter custom = + new RexExpressionConverter(null, CallConverters.defaults(types), null, types); + RexNode column = + rexBuilder.makeInputRef(typeFactory.createArrayType(factory.createCalcite(false), -1), 0); + for (Integer index : new Integer[] {0, null}) { + Expression converted = + rexBuilder + .makeCall(SqlStdOperatorTable.ITEM, column, nullableIndex(index)) + .accept(custom); + assertInstanceOf(FieldReference.class, converted); + verifyRoundTrip(converted, types.toSubstrait(column.getType())); + } + } + + @ParameterizedTest + @ValueSource(longs = {2147483649L, 4294967297L, Long.MAX_VALUE}) + void offsetsPastTheListLengthLimitAreNull(long index) { + RexNode call = rexBuilder.makeCall(SqlStdOperatorTable.ITEM, array(), integer(index)); + Expression converted = call.accept(converter); + assertEquals(Integer.MAX_VALUE, listOffset(converted)); + } + + @ParameterizedTest + @MethodSource("unsafeOperators") + void throwingOperatorsAreRejected(SqlOperator operator) { + RexNode call = rexBuilder.makeCall(operator, array(), integer(4)); + assertThrows(IllegalArgumentException.class, () -> call.accept(converter)); + } + + @ParameterizedTest + @MethodSource("unsafeOperators") + void declinedOperatorsDoNotConvertTheArrayOperand(SqlOperator operator) { + RexNode call = rexBuilder.makeCall(operator, array(), integer(1)); + assertEquals( + Optional.empty(), + new FieldSelectionConverter(TypeConverter.DEFAULT) + .convert( + (RexCall) call, + operand -> { + throw new AssertionError("converted " + operand); + })); + } + + private static Stream unsafeOperators() { + return Stream.of(SqlLibraryOperators.OFFSET, SqlLibraryOperators.ORDINAL); + } + + private RexNode array() { + return rexBuilder.makeCall( + SqlStdOperatorTable.ARRAY_VALUE_CONSTRUCTOR, integer(10), integer(20), integer(30)); + } + + private RexNode integer(long value) { + RelDataType type = + value >= Integer.MIN_VALUE && value <= Integer.MAX_VALUE + ? intType + : typeFactory.createSqlType(SqlTypeName.BIGINT); + return rexBuilder.makeExactLiteral(BigDecimal.valueOf(value), type); + } + + private int listOffset(Expression expression) { + assertInstanceOf(FieldReference.class, expression); + return listOffset(toProto(expression).getSelection().getDirectReference()); + } + + private int listOffset(io.substrait.proto.Expression.ReferenceSegment segment) { + // Protobuf returns offset 0 for a segment that is not a list element. + assertTrue(segment.hasListElement()); + return segment.getListElement().getOffset(); + } + + private io.substrait.proto.Expression toProto(Expression expression) { + return expression.accept( + new ExpressionProtoConverter(new ExtensionCollector(), null), + EmptyVisitationContext.INSTANCE); + } + + private void verifyRoundTrip(Expression expression, Type inputType) { + ExtensionCollector extensions = new ExtensionCollector(); + io.substrait.proto.Expression proto = + new ExpressionProtoConverter(extensions, null).toProto(expression); + Expression restored = + new ProtoExpressionConverter( + extensions, + DefaultExtensionCatalog.DEFAULT_COLLECTION, + TypeCreator.REQUIRED.struct(inputType), + null) + .from(proto); + assertEquals(expression, restored); + } +}