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/SqlConverterBase.java

@@ -1,11 +1,11 @@
 package io.substrait.isthmus;
 
+import com.github.bsideup.jabel.Desugar;
 import io.substrait.extension.SimpleExtension;
 import io.substrait.isthmus.calcite.SubstraitOperatorTable;
 import io.substrait.type.NamedStruct;
 import java.io.IOException;
-import java.util.ArrayList;
-import java.util.List;
+import java.util.*;
 import java.util.function.Function;
 import org.apache.calcite.config.CalciteConnectionConfig;
 import org.apache.calcite.config.CalciteConnectionProperty;
@@ -22,7 +22,10 @@
 import org.apache.calcite.rel.metadata.RelMetadataQuery;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeFactory;
+import org.apache.calcite.rel.type.RelDataTypeField;
 import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexInputRef;
+import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.schema.Schema;
 import org.apache.calcite.schema.Table;
 import org.apache.calcite.schema.impl.AbstractTable;
@@ -86,8 +89,9 @@ protected SqlConverterBase(FeatureBoard features) {
     EXTENSION_COLLECTION = defaults;
   }
 
-  Pair<SqlValidator, CalciteCatalogReader> registerCreateTables(List<String> tables)
-      throws SqlParseException {
+  Result registerCreateTables(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);
@@ -97,12 +101,27 @@ Pair<SqlValidator, CalciteCatalogReader> registerCreateTables(List<String> table
         List<DefinedTable> tList = parseCreateTable(factory, validator, tableDef);
         for (DefinedTable t : tList) {
           rootSchema.add(t.getName(), t);
+          for (RelDataTypeField field : t.type.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 Pair.of(validator, catalogReader);
+    return new Result(validator, catalogReader, nameToTypeMap, nameToNodeMap);
   }
 
+  @Desugar
+  public record Result(
+      SqlValidator validator,
+      CalciteCatalogReader catalogReader,
+      Map<String, RelDataType> nameToTypeMap,
+      Map<String, RexNode> nameToNodeMap) {}
+
   Pair<SqlValidator, CalciteCatalogReader> registerCreateTables(
       Function<List<String>, NamedStruct> tableLookup) throws SqlParseException {
     Function<List<String>, Table> lookup =

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

@@ -1,18 +1,37 @@
 package io.substrait.isthmus;
 
+import com.github.bsideup.jabel.Desugar;
 import com.google.common.annotations.VisibleForTesting;
 import io.substrait.extension.ExtensionCollector;
+import io.substrait.extension.SimpleExtension;
+import io.substrait.isthmus.expression.RexExpressionConverter;
+import io.substrait.isthmus.expression.ScalarFunctionConverter;
+import io.substrait.proto.Expression;
+import io.substrait.proto.Expression.ScalarFunction;
+import io.substrait.proto.ExpressionReference;
+import io.substrait.proto.ExtendedExpression;
+import io.substrait.proto.FunctionArgument;
 import io.substrait.proto.Plan;
 import io.substrait.proto.PlanRel;
+import io.substrait.proto.SimpleExtensionDeclaration;
+import io.substrait.proto.SimpleExtensionURI;
 import io.substrait.relation.RelProtoConverter;
 import io.substrait.type.NamedStruct;
-import java.util.List;
+import io.substrait.type.TypeCreator;
+import io.substrait.type.proto.TypeProtoConverter;
+import java.io.IOException;
+import java.util.*;
 import java.util.function.Function;
 import org.apache.calcite.plan.hep.HepPlanner;
 import org.apache.calcite.plan.hep.HepProgram;
 import org.apache.calcite.prepare.CalciteCatalogReader;
 import org.apache.calcite.rel.RelRoot;
+import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeFactory;
+import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.rex.RexInputRef;
+import org.apache.calcite.rex.RexLiteral;
+import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.schema.Schema;
 import org.apache.calcite.sql.SqlNode;
 import org.apache.calcite.sql.parser.SqlParseException;
@@ -24,6 +43,12 @@
 /** Take a SQL statement and a set of table definitions and return a substrait plan. */
 public class SqlToSubstrait extends SqlConverterBase {
 
+  private final ScalarFunctionConverter functionConverter =
+      new ScalarFunctionConverter(EXTENSION_COLLECTION.scalarFunctions(), factory);
+
+  private final RexExpressionConverter rexExpressionConverter =
+      new RexExpressionConverter(functionConverter);
+
   public SqlToSubstrait() {
     this(null);
   }
@@ -39,19 +64,30 @@ public Plan execute(String sql, Function<List<String>, NamedStruct> tableLookup)
   }
 
   public Plan execute(String sql, List<String> tables) throws SqlParseException {
-    var pair = registerCreateTables(tables);
-    return executeInner(sql, factory, pair.left, pair.right);
+    var result = registerCreateTables(tables);
+    return executeInner(sql, factory, result.validator(), result.catalogReader());
   }
 
   public Plan execute(String sql, String name, Schema schema) throws SqlParseException {
     var pair = registerSchema(name, schema);
     return executeInner(sql, factory, pair.left, pair.right);
   }
 
+  public ExtendedExpression executeExpression(String expr, List<String> tables)
+      throws SqlParseException {
+    var result = registerCreateTables(tables);
+    return executeInnerExpression(
+        expr,
+        result.validator(),
+        result.catalogReader(),
+        result.nameToTypeMap(),
+        result.nameToNodeMap());
+  }
+
   // Package protected for testing
   List<RelRoot> sqlToRelNode(String sql, List<String> tables) throws SqlParseException {
-    var pair = registerCreateTables(tables);
-    return sqlToRelNode(sql, pair.left, pair.right);
+    var result = registerCreateTables(tables);
+    return sqlToRelNode(sql, result.validator(), result.catalogReader());
   }
 
   // Package protected for testing
@@ -91,6 +127,150 @@ private Plan executeInner(
     return plan.build();
   }
 
+  private ExtendedExpression executeInnerExpression(
+      String sql,
+      SqlValidator validator,
+      CalciteCatalogReader catalogReader,
+      Map<String, RelDataType> nameToTypeMap,
+      Map<String, RexNode> nameToNodeMap)
+      throws SqlParseException {
+    ExtendedExpression.Builder extendedExpressionBuilder = ExtendedExpression.newBuilder();
+    ExtensionCollector functionCollector = new ExtensionCollector();
+    RexNode rexNode = sqlToRexNode(sql, validator, catalogReader, nameToTypeMap, nameToNodeMap);
+    ResulTraverseRowExpression result = new TraverseRexNode().getRowExpression(rexNode);
+    io.substrait.proto.Type output =
+        TypeCreator.NULLABLE.BOOLEAN.accept(new TypeProtoConverter(functionCollector));
+    List<FunctionArgument> functionArgumentList = new ArrayList<>();
+    result
+        .expressionBuilderMap()
+        .forEach(
+            (k, v) -> {
+              functionArgumentList.add(FunctionArgument.newBuilder().setValue(v).build());
+            });
+
+    ScalarFunction.Builder scalarFunctionBuilder =
+        ScalarFunction.newBuilder()
+            .setFunctionReference(1) // rel_01
+            .setOutputType(output)
+            .addAllArguments(functionArgumentList);
+
+    Expression.Builder expressionBuilder =
+        Expression.newBuilder().setScalarFunction(scalarFunctionBuilder);
+
+    ExpressionReference.Builder expressionReferenceBuilder =
+        ExpressionReference.newBuilder()
+            .setExpression(expressionBuilder)
+            .addOutputNames(result.ref().getName());
+
+    extendedExpressionBuilder.addReferredExpr(0, expressionReferenceBuilder);
+
+    io.substrait.expression.Expression.ScalarFunctionInvocation func =
+        (io.substrait.expression.Expression.ScalarFunctionInvocation)
+            rexNode.accept(rexExpressionConverter);
+    String declaration = func.declaration().key(); // values example: gt:any_any, add:i64_i64
+
+    // this is not mandatory to be defined; it is working without this definition. It is
+    // only created here to create a proto message that has the correct semantics
+    HashMap<String, SimpleExtensionURI> extensionUris = new HashMap<>();
+    SimpleExtensionURI simpleExtensionURI;
+    try {
+      simpleExtensionURI =
+          SimpleExtensionURI.newBuilder()
+              .setExtensionUriAnchor(1) // rel_02
+              .setUri(
+                  SimpleExtension.loadDefaults().scalarFunctions().stream()
+                      .filter(s -> s.toString().equalsIgnoreCase(declaration))
+                      .findFirst()
+                      .orElseThrow(
+                          () ->
+                              new IllegalArgumentException(
+                                  String.format("Failed to get URI resource for %s.", declaration)))
+                      .uri())
+              .build();
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+    extensionUris.put("uri", simpleExtensionURI);
+
+    ArrayList<SimpleExtensionDeclaration> extensions = new ArrayList<>();
+    SimpleExtensionDeclaration extensionFunctionLowerThan =
+        SimpleExtensionDeclaration.newBuilder()
+            .setExtensionFunction(
+                SimpleExtensionDeclaration.ExtensionFunction.newBuilder()
+                    .setFunctionAnchor(scalarFunctionBuilder.getFunctionReference()) // rel_01
+                    .setName(declaration)
+                    .setExtensionUriReference(simpleExtensionURI.getExtensionUriAnchor())) // rel_02
+            .build();
+    extensions.add(extensionFunctionLowerThan);
+
+    System.out.println(
+        "extendedExpressionBuilder.getExtensionUrisList(): "
+            + extendedExpressionBuilder.getExtensionUrisList());
+    // adding it for semantic purposes, it is not mandatory or needed
+    extendedExpressionBuilder.addAllExtensionUris(extensionUris.values());
+    extendedExpressionBuilder.addAllExtensions(extensions);
+
+    NamedStruct namedStruct = TypeConverter.DEFAULT.toNamedStruct(nameToTypeMap);
+    extendedExpressionBuilder.setBaseSchema(
+        namedStruct.toProto(new TypeProtoConverter(functionCollector)));
+
+    return extendedExpressionBuilder.build();
+  }
+
+  class TraverseRexNode {
+    RexInputRef ref = null;
+    int control = 0;
+    Expression.Builder referenceBuilder = null;
+    Expression.Builder literalBuilder = null;
+    Map<Integer, Expression.Builder> expressionBuilderMap = new LinkedHashMap<>();
+
+    ResulTraverseRowExpression getRowExpression(RexNode rexNode) {
+      switch (rexNode.getClass().getSimpleName().toUpperCase()) {
+        case "REXCALL":
+          for (RexNode rexInternal : ((RexCall) rexNode).operands) {
+            getRowExpression(rexInternal);
+          }
+          ;
+          break;
+        case "REXINPUTREF":
+          ref = (RexInputRef) rexNode;
+          referenceBuilder =
+              Expression.newBuilder()
+                  .setSelection(
+                      Expression.FieldReference.newBuilder()
+                          .setDirectReference(
+                              Expression.ReferenceSegment.newBuilder()
+                                  .setStructField(
+                                      Expression.ReferenceSegment.StructField.newBuilder()
+                                          .setField(ref.getIndex()))));
+          expressionBuilderMap.put(control, referenceBuilder);
+          control++;
+          break;
+        case "REXLITERAL":
+          RexLiteral literal = (RexLiteral) rexNode;
+          literalBuilder =
+              Expression.newBuilder()
+                  .setLiteral(
+                      Expression.Literal.newBuilder().setI32(literal.getValueAs(Integer.class)));
+          expressionBuilderMap.put(control, literalBuilder);
+          control++;
+          break;
+        default:
+          throw new AssertionError(
+              "Unsupported type for: " + rexNode.getClass().getSimpleName().toUpperCase());
+      }
+      return new ResulTraverseRowExpression(
+          ref, referenceBuilder, literalBuilder, expressionBuilderMap);
+    }
+  }
+
+  @Desugar
+  private record ResulTraverseRowExpression(
+      RexInputRef ref,
+      Expression.Builder referenceBuilder,
+      Expression.Builder literalBuilder,
+      Map<Integer, Expression.Builder> expressionBuilderMap) {}
+
   private List<RelRoot> sqlToRelNode(
       String sql, SqlValidator validator, CalciteCatalogReader catalogReader)
       throws SqlParseException {
@@ -107,6 +287,20 @@ private List<RelRoot> sqlToRelNode(
     return roots;
   }
 
+  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);
+  }
+
   @VisibleForTesting
   SqlToRelConverter createSqlToRelConverter(
       SqlValidator validator, CalciteCatalogReader catalogReader) {

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

@@ -22,8 +22,9 @@ public SubstraitToSql() {
 
   public RelNode substraitRelToCalciteRel(Rel relRoot, List<String> tables)
       throws SqlParseException {
-    var pair = registerCreateTables(tables);
-    return SubstraitRelNodeConverter.convert(relRoot, relOptCluster, pair.right, parserConfig);
+    var result = registerCreateTables(tables);
+    return SubstraitRelNodeConverter.convert(
+        relRoot, relOptCluster, result.catalogReader(), parserConfig);
   }
 
   public RelNode substraitRelToCalciteRel(