feat: convert sql expression into proto extended expressions

core/src/main/java/io/substrait/extension/ExtensionCollector.java

@@ -1,5 +1,6 @@
 package io.substrait.extension;
 
+import io.substrait.proto.ExtendedExpression;
 import io.substrait.proto.Plan;
 import io.substrait.proto.SimpleExtensionDeclaration;
 import io.substrait.proto.SimpleExtensionURI;
@@ -98,6 +99,54 @@ public void addExtensionsToPlan(Plan.Builder builder) {
     builder.addAllExtensions(extensionList);
   }
 
+  public void addExtensionsToExtendedExpression(ExtendedExpression.Builder builder) {
+    var uriPos = new AtomicInteger(1);
+    var uris = new HashMap<String, SimpleExtensionURI>();
+
+    var extensionList = new ArrayList<SimpleExtensionDeclaration>();
+    for (var e : funcMap.forwardMap.entrySet()) {
+      SimpleExtensionURI uri =
+          uris.computeIfAbsent(
+              e.getValue().namespace(),
+              k ->
+                  SimpleExtensionURI.newBuilder()
+                      .setExtensionUriAnchor(uriPos.getAndIncrement())
+                      .setUri(k)
+                      .build());
+      var decl =
+          SimpleExtensionDeclaration.newBuilder()
+              .setExtensionFunction(
+                  SimpleExtensionDeclaration.ExtensionFunction.newBuilder()
+                      .setFunctionAnchor(e.getKey())
+                      .setName(e.getValue().key())
+                      .setExtensionUriReference(uri.getExtensionUriAnchor()))
+              .build();
+      extensionList.add(decl);
+    }
+    for (var e : typeMap.forwardMap.entrySet()) {
+      SimpleExtensionURI uri =
+          uris.computeIfAbsent(
+              e.getValue().namespace(),
+              k ->
+                  SimpleExtensionURI.newBuilder()
+                      .setExtensionUriAnchor(uriPos.getAndIncrement())
+                      .setUri(k)
+                      .build());
+      var decl =
+          SimpleExtensionDeclaration.newBuilder()
+              .setExtensionType(
+                  SimpleExtensionDeclaration.ExtensionType.newBuilder()
+                      .setTypeAnchor(e.getKey())
+                      .setName(e.getValue().key())
+                      .setExtensionUriReference(uri.getExtensionUriAnchor()))
+              .build();
+      extensionList.add(decl);
+    }
+
+    builder.addAllExtensionUris(uris.values());
+    builder.addAllExtensions(extensionList);
+  }
+
   /** We don't depend on guava... */
   private static class BidiMap<T1, T2> {
     private final Map<T1, T2> forwardMap;

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

@@ -1,22 +1,44 @@
 package io.substrait.isthmus;
 
+import com.github.bsideup.jabel.Desugar;
 import com.google.common.annotations.VisibleForTesting;
 import io.substrait.extension.ExtensionCollector;
+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 io.substrait.type.Type;
+import io.substrait.type.TypeCreator;
+import io.substrait.type.proto.TypeProtoConverter;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 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;
 import org.apache.calcite.sql.parser.SqlParser;
+import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.calcite.sql.validate.SqlValidator;
 import org.apache.calcite.sql2rel.SqlToRelConverter;
 import org.apache.calcite.sql2rel.StandardConvertletTable;
@@ -48,6 +70,12 @@ public Plan execute(String sql, String name, Schema schema) throws SqlParseExcep
     return executeInner(sql, factory, pair.left, pair.right);
   }
 
+  public ExtendedExpression executeExpression(String expr, List<String> tables)
+      throws SqlParseException {
+    var pair = registerCreateTables(tables);
+    return executeInnerExpression(expr, pair.left, pair.right);
+  }
+
   // Package protected for testing
   List<RelRoot> sqlToRelNode(String sql, List<String> tables) throws SqlParseException {
     var pair = registerCreateTables(tables);
@@ -91,6 +119,138 @@ private Plan executeInner(
     return plan.build();
   }
 
+  private ExtendedExpression executeInnerExpression(
+      String sql, SqlValidator validator, CalciteCatalogReader catalogReader)
+      throws SqlParseException {
+    ExtendedExpression.Builder extendedExpressionBuilder = ExtendedExpression.newBuilder();
+    ExtensionCollector functionCollector = new ExtensionCollector();
+    sqlToRexNode(sql, validator, catalogReader)
+        .forEach(
+            rexNode -> {
+              // FIXME! Implement it dynamically for more expression types
+              ResulTraverseRowExpression result = TraverseRexNode.getRowExpression(rexNode);
+
+              // FIXME! Get output type dynamically:
+              // final static Map<String, Type> getTypeCreator = new HashMap<>(){{put("BOOLEAN",
+              // TypeCreator.of(true).BOOLEAN);}};
+              // getTypeCreator.get(rexNode.getType()).accept(...)
+              io.substrait.proto.Type output =
+                  TypeCreator.NULLABLE.BOOLEAN.accept(new TypeProtoConverter(functionCollector));
+
+              // FIXME! setFunctionReference, addArguments(index: 0, 1)
+              Expression.Builder expressionBuilder =
+                  Expression.newBuilder()
+                      .setScalarFunction(
+                          ScalarFunction.newBuilder()
+                              .setFunctionReference(1)
+                              .setOutputType(output)
+                              .addArguments(
+                                  0,
+                                  FunctionArgument.newBuilder().setValue(result.referenceBuilder()))
+                              .addArguments(
+                                  1,
+                                  FunctionArgument.newBuilder()
+                                      .setValue(result.expressionBuilderLiteral())));
+              ExpressionReference.Builder expressionReferenceBuilder =
+                  ExpressionReference.newBuilder()
+                      .setExpression(expressionBuilder)
+                      .addOutputNames(result.ref().getName());
+
+              // FIXME! Get schema dynamically
+              // (as the same for Plan with:
+              // TypeConverter.DEFAULT.toNamedStruct(rexNode.getType());)
+              List<String> columnNames =
+                  Arrays.asList("N_NATIONKEY", "N_NAME", "N_REGIONKEY", "N_COMMENT");
+              List<Type> dataTypes =
+                  Arrays.asList(
+                      TypeCreator.NULLABLE.I32,
+                      TypeCreator.NULLABLE.STRING,
+                      TypeCreator.NULLABLE.I32,
+                      TypeCreator.NULLABLE.STRING);
+              NamedStruct namedStruct =
+                  NamedStruct.of(
+                      columnNames, Type.Struct.builder().fields(dataTypes).nullable(false).build());
+
+              extendedExpressionBuilder
+                  .addReferredExpr(0, expressionReferenceBuilder)
+                  .setBaseSchema(namedStruct.toProto(new TypeProtoConverter(functionCollector)));
+
+              // Extensions URI FIXME! Populate/create this dynamically
+              HashMap<String, SimpleExtensionURI> extensionUris = new HashMap<>();
+              extensionUris.put(
+                  "key-001",
+                  SimpleExtensionURI.newBuilder()
+                      .setExtensionUriAnchor(1)
+                      .setUri("/functions_comparison.yaml")
+                      .build());
+
+              // Extensions FIXME! Populate/create this dynamically, maybe use rexNode.getKind()
+              ArrayList<SimpleExtensionDeclaration> extensions = new ArrayList<>();
+              SimpleExtensionDeclaration extensionFunctionLowerThan =
+                  SimpleExtensionDeclaration.newBuilder()
+                      .setExtensionFunction(
+                          SimpleExtensionDeclaration.ExtensionFunction.newBuilder()
+                              .setFunctionAnchor(1)
+                              .setName("gt:any_any")
+                              .setExtensionUriReference(1))
+                      .build();
+              extensions.add(extensionFunctionLowerThan);
+
+              extendedExpressionBuilder.addAllExtensionUris(extensionUris.values());
+              extendedExpressionBuilder.addAllExtensions(extensions);
+            });
+    return extendedExpressionBuilder.build();
+  }
+
+  static class TraverseRexNode {
+    static RexInputRef ref = null;
+    static Expression.Builder referenceBuilder = null;
+    static Expression.Builder expressionBuilderLiteral = null;
+
+    static 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()))));
+          break;
+        case "REXLITERAL":
+          RexLiteral literal = (RexLiteral) rexNode;
+          expressionBuilderLiteral =
+              Expression.newBuilder()
+                  .setLiteral(
+                      Expression.Literal.newBuilder().setI32(literal.getValueAs(Integer.class)));
+          break;
+        default:
+          throw new AssertionError(
+              "Unsupported type for: " + rexNode.getClass().getSimpleName().toUpperCase());
+      }
+      ResulTraverseRowExpression result =
+          new ResulTraverseRowExpression(ref, referenceBuilder, expressionBuilderLiteral);
+      return result;
+    }
+  }
+
+  @Desugar
+  private record ResulTraverseRowExpression(
+      RexInputRef ref,
+      Expression.Builder referenceBuilder,
+      Expression.Builder expressionBuilderLiteral) {}
+
   private List<RelRoot> sqlToRelNode(
       String sql, SqlValidator validator, CalciteCatalogReader catalogReader)
       throws SqlParseException {
@@ -107,6 +267,42 @@ private List<RelRoot> sqlToRelNode(
     return roots;
   }
 
+  private List<RexNode> sqlToRexNode(
+      String sql, SqlValidator validator, CalciteCatalogReader catalogReader)
+      throws SqlParseException {
+    SqlParser parser = SqlParser.create(sql, parserConfig);
+    SqlNode sqlNode = parser.parseExpression();
+    Result result = getResult(validator);
+    SqlNode validSQLNode =
+        validator.validateParameterizedExpression(
+            sqlNode,
+            result.nameToTypeMap()); // FIXME! It may be optional to include this validation
+    SqlToRelConverter converter = createSqlToRelConverter(validator, catalogReader);
+    RexNode rexNode = converter.convertExpression(validSQLNode, result.nameToNodeMap());
+
+    return Collections.singletonList(rexNode);
+  }
+
+  private static Result getResult(SqlValidator validator) {
+    // FIXME! Needs to be created dinamycally, this is for PoC purpose
+    HashMap<String, RexNode> nameToNodeMap = new HashMap<>();
+    nameToNodeMap.put(
+        "N_NATIONKEY",
+        new RexInputRef(0, validator.getTypeFactory().createSqlType(SqlTypeName.BIGINT)));
+    nameToNodeMap.put(
+        "N_REGIONKEY",
+        new RexInputRef(1, validator.getTypeFactory().createSqlType(SqlTypeName.BIGINT)));
+    final Map<String, RelDataType> nameToTypeMap = new HashMap<>();
+    for (Map.Entry<String, RexNode> entry : nameToNodeMap.entrySet()) {
+      nameToTypeMap.put(entry.getKey(), entry.getValue().getType());
+    }
+    Result result = new Result(nameToNodeMap, nameToTypeMap);
+    return result;
+  }
+
+  private @Desugar record Result(
+      HashMap<String, RexNode> nameToNodeMap, Map<String, RelDataType> nameToTypeMap) {}
+
   @VisibleForTesting
   SqlToRelConverter createSqlToRelConverter(
       SqlValidator validator, CalciteCatalogReader catalogReader) {

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

@@ -0,0 +1,51 @@
+package io.substrait.isthmus;
+
+import com.google.common.base.Charsets;
+import com.google.common.io.Resources;
+import com.google.protobuf.InvalidProtocolBufferException;
+import com.google.protobuf.util.JsonFormat;
+import io.substrait.proto.ExtendedExpression;
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.List;
+import org.apache.calcite.sql.parser.SqlParseException;
+
+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 SqlToSubstrait());
+  }
+
+  protected ExtendedExpression assertProtoExtendedExpressionRoundrip(String query, SqlToSubstrait s)
+      throws IOException, SqlParseException {
+    return assertProtoExtendedExpressionRoundrip(query, s, tpchSchemaCreateStatements());
+  }
+
+  protected ExtendedExpression assertProtoExtendedExpressionRoundrip(
+      String query, SqlToSubstrait s, List<String> creates) throws SqlParseException {
+    io.substrait.proto.ExtendedExpression protoExtendedExpression =
+        s.executeExpression(query, creates);
+
+    try {
+      String ee = JsonFormat.printer().print(protoExtendedExpression);
+      System.out.println("Proto Extended Expression: \n" + ee);
+
+      // FIXME! Implement test validation as the same as proto Plan implementation
+    } catch (InvalidProtocolBufferException e) {
+      throw new RuntimeException(e);
+    }
+
+    return protoExtendedExpression;
+  }
+}

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

@@ -0,0 +1,13 @@
+package io.substrait.isthmus;
+
+import java.io.IOException;
+import org.apache.calcite.sql.parser.SqlParseException;
+import org.junit.jupiter.api.Test;
+
+public class SimpleExtendedExpressionsTest extends ExtendedExpressionTestBase {
+
+  @Test
+  public void filter() throws IOException, SqlParseException {
+    assertProtoExtendedExpressionRoundrip("N_NATIONKEY > 18");
+  }
+}