diff --git a/isthmus/src/main/java/io/substrait/isthmus/SqlExpressionToSubstrait.java b/isthmus/src/main/java/io/substrait/isthmus/SqlExpressionToSubstrait.java index 620d12927..35b48c30e 100644 --- a/isthmus/src/main/java/io/substrait/isthmus/SqlExpressionToSubstrait.java +++ b/isthmus/src/main/java/io/substrait/isthmus/SqlExpressionToSubstrait.java @@ -208,6 +208,8 @@ private Result registerCreateTablesForExtendedExpression(List tables) for (SubstraitTable t : tList) { rootSchema.add(t.getName(), t); for (RelDataTypeField field : t.getRowType(factory).getFieldList()) { + // Field references index the combined base schema in insertion order. + int fieldIndex = nameToTypeMap.size(); nameToTypeMap.merge( // to validate the sql expression tree field.getName(), field.getType(), @@ -217,7 +219,7 @@ private Result registerCreateTablesForExtendedExpression(List tables) }); nameToNodeMap.merge( // to convert sql expression into RexNode field.getName(), - new RexInputRef(field.getIndex(), field.getType()), + new RexInputRef(fieldIndex, field.getType()), (v1, v2) -> { throw new IllegalArgumentException( "There is no support for duplicate column names: " + field.getName()); diff --git a/isthmus/src/test/java/io/substrait/isthmus/SimpleExtendedExpressionsTest.java b/isthmus/src/test/java/io/substrait/isthmus/SimpleExtendedExpressionsTest.java index cc31ce2e9..f753a5243 100644 --- a/isthmus/src/test/java/io/substrait/isthmus/SimpleExtendedExpressionsTest.java +++ b/isthmus/src/test/java/io/substrait/isthmus/SimpleExtendedExpressionsTest.java @@ -6,9 +6,11 @@ import static org.junit.jupiter.api.Assertions.assertTrue; import io.substrait.isthmus.expression.RexExpressionConverter; +import io.substrait.proto.Expression; import io.substrait.proto.Expression.RexTypeCase; import io.substrait.proto.ExtendedExpression; import java.io.IOException; +import java.util.List; import java.util.stream.Stream; import org.apache.calcite.sql.parser.SqlParseException; import org.junit.jupiter.api.Test; @@ -20,6 +22,78 @@ class SimpleExtendedExpressionsTest extends ExtendedExpressionTestBase { private static final String MARKER = "provider hook reached"; + private static final String TABLE_A = "CREATE TABLE A (A1 BIGINT, A2 BIGINT, A3 BIGINT)"; + private static final String TABLE_B = "CREATE TABLE B (B1 BIGINT, B2 BIGINT)"; + private static final String TABLE_C = "CREATE TABLE C (C1 BIGINT)"; + + private static Stream columnSchemaProvider() { + return Stream.of( + Arguments.of(List.of(TABLE_A), List.of("A1", "A2", "A3")), + Arguments.of( + List.of(TABLE_A, TABLE_B, TABLE_C), List.of("A1", "A2", "A3", "B1", "B2", "C1")), + Arguments.of( + List.of(TABLE_A + ";" + TABLE_B + ";" + TABLE_C), + List.of("A1", "A2", "A3", "B1", "B2", "C1")), + Arguments.of( + List.of(TABLE_B, TABLE_C, TABLE_A), List.of("B1", "B2", "C1", "A1", "A2", "A3"))); + } + + @ParameterizedTest + @MethodSource("columnSchemaProvider") + void fieldReferencesIndexTheCombinedSchema(List tables, List columnNames) + throws SqlParseException { + // Reverse the expression order so a reference's index cannot accidentally be its position + // in the output expression list. All columns have the same type, so types cannot detect this. + String[] expressions = new String[columnNames.size()]; + for (int index = 0; index < expressions.length; index++) { + expressions[index] = columnNames.get(columnNames.size() - index - 1); + } + ExtendedExpression converted = new SqlExpressionToSubstrait().convert(expressions, tables); + + assertEquals(columnNames, converted.getBaseSchema().getNamesList()); + assertEquals(columnNames.size(), converted.getBaseSchema().getStruct().getTypesCount()); + assertEquals(expressions.length, converted.getReferredExprCount()); + for (int index = 0; index < expressions.length; index++) { + assertEquals( + columnNames.size() - index - 1, + selectedField(converted.getReferredExpr(index).getExpression()), + expressions[index]); + } + } + + @Test + void functionArgumentsIndexTheCombinedSchema() throws SqlParseException { + ExtendedExpression converted = + new SqlExpressionToSubstrait() + .convert(new String[] {"A1 = B1", "B2 + A3"}, List.of(TABLE_A, TABLE_B)); + + Expression.ScalarFunction filter = + converted.getReferredExpr(0).getExpression().getScalarFunction(); + assertEquals(0, selectedField(filter.getArguments(0).getValue())); + assertEquals(3, selectedField(filter.getArguments(1).getValue())); + Expression.ScalarFunction projection = + converted.getReferredExpr(1).getExpression().getScalarFunction(); + assertEquals(4, selectedField(projection.getArguments(0).getValue())); + assertEquals(2, selectedField(projection.getArguments(1).getValue())); + } + + @Test + void eachConversionBuildsItsOwnColumnIndices() throws SqlParseException { + SqlExpressionToSubstrait converter = new SqlExpressionToSubstrait(); + ExtendedExpression multipleTables = converter.convert("B2", List.of(TABLE_A, TABLE_B)); + ExtendedExpression singleTable = converter.convert("B2", List.of(TABLE_B)); + + assertEquals(4, selectedField(multipleTables.getReferredExpr(0).getExpression())); + assertEquals(1, selectedField(singleTable.getReferredExpr(0).getExpression())); + } + + private static int selectedField(Expression expression) { + assertEquals(RexTypeCase.SELECTION, expression.getRexTypeCase()); + assertTrue(expression.getSelection().hasRootReference()); + assertTrue(expression.getSelection().getDirectReference().hasStructField()); + return expression.getSelection().getDirectReference().getStructField().getField(); + } + private static Stream expressionTypeProvider() { return Stream.of( Arguments.of("2"), // I32LiteralExpression