feat: convert sql expression into proto extended expressions

.editorconfig

@@ -10,7 +10,7 @@ trim_trailing_whitespace = true
 [*.{yaml,yml}]
 indent_size = 2
 
-[{**/*.sql,**/OuterReferenceResolver.md,gradlew.bat}]
+[{**/*.sql,**/OuterReferenceResolver.md,gradlew.bat,**/*.parquet}]
 charset = unset
 end_of_line = unset
 insert_final_newline = unset

.pre-commit-config.yaml

@@ -1,11 +1,11 @@
 repos:
 - repo: https://github.com/adrienverge/yamllint.git
-  rev: v1.26.0
+  rev: v1.33.0
   hooks:
   - id: yamllint
     args: [-c=.yamllint.yaml]
 - repo: https://github.com/alessandrojcm/commitlint-pre-commit-hook
-  rev: v8.0.0
+  rev: v9.9.0
   hooks:
   - id: commitlint
     stages: [commit-msg]

isthmus/build.gradle.kts

@@ -72,6 +72,7 @@ java {
 }
 
 var CALCITE_VERSION = "1.34.0"
+var ARROW_VERSION = "14.0.0"
 
 dependencies {
   implementation(project(":core"))
@@ -94,6 +95,8 @@ dependencies {
   implementation("org.immutables:value-annotations:2.8.8")
   annotationProcessor("org.immutables:value:2.8.8")
   testImplementation("org.apache.calcite:calcite-plus:${CALCITE_VERSION}")
+  testImplementation("org.apache.arrow:arrow-dataset:${ARROW_VERSION}")
+  testImplementation("org.apache.arrow:arrow-memory-netty:${ARROW_VERSION}")
   annotationProcessor("com.github.bsideup.jabel:jabel-javac-plugin:0.4.2")
   compileOnly("com.github.bsideup.jabel:jabel-javac-plugin:0.4.2")
 }

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

@@ -0,0 +1,165 @@
+package io.substrait.isthmus;
+
+import com.github.bsideup.jabel.Desugar;
+import io.substrait.extendedexpression.ExtendedExpressionProtoConverter;
+import io.substrait.extendedexpression.ImmutableExpressionReference;
+import io.substrait.extendedexpression.ImmutableExtendedExpression;
+import io.substrait.extension.SimpleExtension;
+import io.substrait.isthmus.expression.RexExpressionConverter;
+import io.substrait.isthmus.expression.ScalarFunctionConverter;
+import io.substrait.proto.ExtendedExpression;
+import io.substrait.type.NamedStruct;
+import io.substrait.type.Type;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.calcite.jdbc.CalciteSchema;
+import org.apache.calcite.prepare.CalciteCatalogReader;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rel.type.RelDataTypeField;
+import org.apache.calcite.rex.RexInputRef;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.SqlNode;
+import org.apache.calcite.sql.parser.SqlParseException;
+import org.apache.calcite.sql.parser.SqlParser;
+import org.apache.calcite.sql.validate.SqlValidator;
+import org.apache.calcite.sql2rel.SqlToRelConverter;
+import org.apache.calcite.sql2rel.StandardConvertletTable;
+
+public class SqlExpressionToSubstrait extends SqlConverterBase {
+
+  protected final RexExpressionConverter rexConverter;
+
+  public SqlExpressionToSubstrait() {
+    this(FEATURES_DEFAULT, EXTENSION_COLLECTION);
+  }
+
+  public SqlExpressionToSubstrait(
+      FeatureBoard features, SimpleExtension.ExtensionCollection extensions) {
+    super(features);
+    ScalarFunctionConverter scalarFunctionConverter =
+        new ScalarFunctionConverter(extensions.scalarFunctions(), factory);
+    this.rexConverter = new RexExpressionConverter(scalarFunctionConverter);
+  }
+
+  @Desugar
+  private record Result(
+      SqlValidator validator,
+      CalciteCatalogReader catalogReader,
+      Map<String, RelDataType> nameToTypeMap,
+      Map<String, RexNode> nameToNodeMap) {}
+
+  /**
+   * Converts the given SQL expression string to an {@link io.substrait.proto.ExtendedExpression }
+   *
+   * @param sqlExpression a SQL expression
+   * @param createStatements table creation statements defining fields referenced by the expression
+   * @return a {@link io.substrait.proto.ExtendedExpression }
+   * @throws SqlParseException
+   */
+  public ExtendedExpression convert(String sqlExpression, List<String> createStatements)
+      throws SqlParseException {
+    var result = registerCreateTablesForExtendedExpression(createStatements);
+    return executeInnerSQLExpression(
+        sqlExpression,
+        result.validator(),
+        result.catalogReader(),
+        result.nameToTypeMap(),
+        result.nameToNodeMap());
+  }
+
+  private ExtendedExpression executeInnerSQLExpression(
+      String sqlExpression,
+      SqlValidator validator,
+      CalciteCatalogReader catalogReader,
+      Map<String, RelDataType> nameToTypeMap,
+      Map<String, RexNode> nameToNodeMap)
+      throws SqlParseException {
+    RexNode rexNode =
+        sqlToRexNode(sqlExpression, validator, catalogReader, nameToTypeMap, nameToNodeMap);
+    NamedStruct namedStruct = toNamedStruct(nameToTypeMap);
+
+    ImmutableExpressionReference expressionReference =
+        ImmutableExpressionReference.builder()
+            .expression(rexNode.accept(this.rexConverter))
+            .addOutputNames("new-column")
+            .build();
+
+    List<io.substrait.extendedexpression.ExtendedExpression.ExpressionReference>
+        expressionReferences = new ArrayList<>();
+    expressionReferences.add(expressionReference);
+
+    ImmutableExtendedExpression.Builder extendedExpression =
+        ImmutableExtendedExpression.builder()
+            .referredExpressions(expressionReferences)
+            .baseSchema(namedStruct);
+
+    return new ExtendedExpressionProtoConverter().toProto(extendedExpression.build());
+  }
+
+  private RexNode sqlToRexNode(
+      String sql,
+      SqlValidator validator,
+      CalciteCatalogReader catalogReader,
+      Map<String, RelDataType> nameToTypeMap,
+      Map<String, RexNode> nameToNodeMap)
+      throws SqlParseException {
+    SqlParser parser = SqlParser.create(sql, parserConfig);
+    SqlNode sqlNode = parser.parseExpression();
+    SqlNode validSQLNode = validator.validateParameterizedExpression(sqlNode, nameToTypeMap);
+    SqlToRelConverter converter = createSqlToRelConverter(validator, catalogReader);
+    return converter.convertExpression(validSQLNode, nameToNodeMap);
+  }
+
+  SqlToRelConverter createSqlToRelConverter(
+      SqlValidator validator, CalciteCatalogReader catalogReader) {
+    return new SqlToRelConverter(
+        null,
+        validator,
+        catalogReader,
+        relOptCluster,
+        StandardConvertletTable.INSTANCE,
+        converterConfig);
+  }
+
+  private Result registerCreateTablesForExtendedExpression(List<String> tables)
+      throws SqlParseException {
+    Map<String, RelDataType> nameToTypeMap = new LinkedHashMap<>();
+    Map<String, RexNode> nameToNodeMap = new HashMap<>();
+    CalciteSchema rootSchema = CalciteSchema.createRootSchema(false);
+    CalciteCatalogReader catalogReader =
+        new CalciteCatalogReader(rootSchema, List.of(), factory, config);
+    SqlValidator validator = Validator.create(factory, catalogReader, SqlValidator.Config.DEFAULT);
+    if (tables != null) {
+      for (String tableDef : tables) {
+        List<DefinedTable> tList = parseCreateTable(factory, validator, tableDef);
+        for (DefinedTable t : tList) {
+          rootSchema.add(t.getName(), t);
+          for (RelDataTypeField field : t.getRowType(factory).getFieldList()) {
+            nameToTypeMap.put(
+                field.getName(), field.getType()); // to validate the sql expression tree
+            nameToNodeMap.put(
+                field.getName(),
+                new RexInputRef(
+                    field.getIndex(), field.getType())); // to convert sql expression into RexNode
+          }
+        }
+      }
+    }
+    return new Result(validator, catalogReader, nameToTypeMap, nameToNodeMap);
+  }
+
+  private NamedStruct toNamedStruct(Map<String, RelDataType> nameToTypeMap) {
+    var names = new ArrayList<String>();
+    var types = new ArrayList<Type>();
+    for (Map.Entry<String, RelDataType> entry : nameToTypeMap.entrySet()) {
+      String k = entry.getKey();
+      RelDataType v = entry.getValue();
+      names.add(k);
+      types.add(TypeConverter.DEFAULT.toSubstrait(v));
+    }
+    return NamedStruct.of(names, Type.Struct.builder().fields(types).nullable(false).build());
+  }
+}

isthmus/src/test/java/io/substrait/isthmus/ExtendedExpressionTestBase.java

@@ -0,0 +1,55 @@
+package io.substrait.isthmus;
+
+import com.google.common.base.Charsets;
+import com.google.common.io.Resources;
+import io.substrait.extendedexpression.ExtendedExpressionProtoConverter;
+import io.substrait.extendedexpression.ProtoExtendedExpressionConverter;
+import io.substrait.proto.ExtendedExpression;
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.List;
+import org.apache.calcite.sql.parser.SqlParseException;
+import org.junit.jupiter.api.Assertions;
+
+public class ExtendedExpressionTestBase {
+  public static String asString(String resource) throws IOException {
+    return Resources.toString(Resources.getResource(resource), Charsets.UTF_8);
+  }
+
+  public static List<String> tpchSchemaCreateStatements() throws IOException {
+    String[] values = asString("tpch/schema.sql").split(";");
+    return Arrays.stream(values)
+        .filter(t -> !t.trim().isBlank())
+        .collect(java.util.stream.Collectors.toList());
+  }
+
+  protected ExtendedExpression assertProtoExtendedExpressionRoundrip(String query)
+      throws IOException, SqlParseException {
+    return assertProtoExtendedExpressionRoundrip(query, new SqlExpressionToSubstrait());
+  }
+
+  protected ExtendedExpression assertProtoExtendedExpressionRoundrip(
+      String query, SqlExpressionToSubstrait s) throws IOException, SqlParseException {
+    return assertProtoExtendedExpressionRoundrip(query, s, tpchSchemaCreateStatements());
+  }
+
+  protected ExtendedExpression assertProtoExtendedExpressionRoundrip(
+      String query, SqlExpressionToSubstrait s, List<String> creates)
+      throws SqlParseException, IOException {
+    // proto initial extended expression
+    ExtendedExpression extendedExpressionProtoInitial = s.convert(query, creates);
+
+    // pojo final extended expression
+    io.substrait.extendedexpression.ExtendedExpression extendedExpressionPojoFinal =
+        new ProtoExtendedExpressionConverter().from(extendedExpressionProtoInitial);
+
+    // proto final extended expression
+    ExtendedExpression extendedExpressionProtoFinal =
+        new ExtendedExpressionProtoConverter().toProto(extendedExpressionPojoFinal);
+
+    // round-trip to validate extended expression proto initial equals to final
+    Assertions.assertEquals(extendedExpressionProtoFinal, extendedExpressionProtoInitial);
+
+    return extendedExpressionProtoInitial;
+  }
+}

isthmus/src/test/java/io/substrait/isthmus/SimpleExtendedExpressionsTest.java

@@ -0,0 +1,32 @@
+package io.substrait.isthmus;
+
+import io.substrait.proto.ExtendedExpression;
+import java.io.IOException;
+import java.util.stream.Stream;
+import org.apache.calcite.sql.parser.SqlParseException;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+
+public class SimpleExtendedExpressionsTest extends ExtendedExpressionTestBase {
+
+  private static Stream<Arguments> expressionTypeProvider() {
+    return Stream.of(
+        Arguments.of("2"), // I32LiteralExpression
+        Arguments.of("L_ORDERKEY"), // FieldReferenceExpression
+        Arguments.of("L_ORDERKEY > 10"), // ScalarFunctionExpressionFilter
+        Arguments.of("L_ORDERKEY + 10"), // ScalarFunctionExpressionProjection
+        Arguments.of("L_ORDERKEY IN (10, 20)"), // ScalarFunctionExpressionIn
+        Arguments.of("L_ORDERKEY is not null"), // ScalarFunctionExpressionIsNotNull
+        Arguments.of("L_ORDERKEY is null") // ScalarFunctionExpressionIsNull
+        );
+  }
+
+  @ParameterizedTest
+  @MethodSource("expressionTypeProvider")
+  public void testExtendedExpressionsRoundTrip(String sqlExpression)
+      throws SqlParseException, IOException {
+
+    ExtendedExpression extendedExpression = assertProtoExtendedExpressionRoundrip(sqlExpression);
+  }
+}