fix(isthmus): index extended expression fields against the combined schema

isthmus/src/main/java/io/substrait/isthmus/SqlExpressionToSubstrait.java

@@ -208,6 +208,8 @@ private Result registerCreateTablesForExtendedExpression(List<String> 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<String> 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());

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<Arguments> 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<String> tables, List<String> 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<Arguments> expressionTypeProvider() {
     return Stream.of(
         Arguments.of("2"), // I32LiteralExpression