feat(isthmus): add support for scalar subqueries

Signed-off-by: Niels Pardon <[email protected]>

Author: nielspardon

Diff for: isthmus/src/main/java/io/substrait/isthmus/expression/ExpressionRexConverter.java

@@ -5,6 +5,7 @@
 import io.substrait.expression.EnumArg;
 import io.substrait.expression.Expression;
 import io.substrait.expression.Expression.FailureBehavior;
+import io.substrait.expression.Expression.ScalarSubquery;
 import io.substrait.expression.Expression.SingleOrList;
 import io.substrait.expression.Expression.Switch;
 import io.substrait.expression.FieldReference;
@@ -538,4 +539,10 @@ public RexNode visitEnumArg(SimpleExtension.Function fnDef, int argIdx, EnumArg
                         "EnumArg(value=%s) not handled by visitor type %s.",
                         e.value(), this.getClass().getCanonicalName())));
   }
+
+  @Override
+  public RexNode visit(ScalarSubquery expr) throws RuntimeException {
+    RelNode inputRelnode = expr.input().accept(relNodeConverter);
+    return RexSubQuery.scalar(inputRelnode);
+  }
 }

Diff for: isthmus/src/test/java/io/substrait/isthmus/PlanTestBase.java

@@ -35,7 +35,7 @@
 
 public class PlanTestBase {
   protected final SimpleExtension.ExtensionCollection extensions = SimpleExtension.loadDefaults();
-  protected final RelCreator creator = new RelCreator();
+  protected final RelCreator creator = new RelCreator(tpchSchemaCreateStatements());
   protected final RelBuilder builder = creator.createRelBuilder();
   protected final RexBuilder rex = creator.rex();
   protected final RelDataTypeFactory typeFactory = creator.typeFactory();
@@ -47,11 +47,16 @@ 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());
+  public static List<String> tpchSchemaCreateStatements() {
+    String[] values;
+    try {
+      values = asString("tpch/schema.sql").split(";");
+      return Arrays.stream(values)
+          .filter(t -> !t.trim().isBlank())
+          .collect(java.util.stream.Collectors.toList());
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
   }
 
   protected Plan assertProtoPlanRoundrip(String query) throws IOException, SqlParseException {

Diff for: isthmus/src/test/java/io/substrait/isthmus/RelCreator.java

@@ -1,6 +1,7 @@
 package io.substrait.isthmus;
 
 import java.util.Arrays;
+import java.util.List;
 import org.apache.calcite.config.CalciteConnectionConfig;
 import org.apache.calcite.config.CalciteConnectionProperty;
 import org.apache.calcite.jdbc.CalciteSchema;
@@ -25,19 +26,37 @@
 import org.apache.calcite.sql2rel.SqlToRelConverter;
 import org.apache.calcite.sql2rel.StandardConvertletTable;
 import org.apache.calcite.tools.RelBuilder;
+import org.apache.calcite.util.Pair;
 
-public class RelCreator {
+public class RelCreator extends SqlConverterBase {
   static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(RelCreator.class);
 
   private RelOptCluster cluster;
   private CalciteCatalogReader catalog;
+  private SqlValidator validator;
 
   public RelCreator() {
+    super(SqlConverterBase.FEATURES_DEFAULT);
     CalciteSchema schema = CalciteSchema.createRootSchema(false);
     RelDataTypeFactory factory = new JavaTypeFactoryImpl(SubstraitTypeSystem.TYPE_SYSTEM);
     CalciteConnectionConfig config =
         CalciteConnectionConfig.DEFAULT.set(CalciteConnectionProperty.CASE_SENSITIVE, "false");
-    catalog = new CalciteCatalogReader(schema, Arrays.asList(), factory, config);
+    this.validator = new Validator(catalog, cluster.getTypeFactory(), SqlValidator.Config.DEFAULT);
+    this.catalog = new CalciteCatalogReader(schema, Arrays.asList(), factory, config);
+    VolcanoPlanner planner = new VolcanoPlanner(RelOptCostImpl.FACTORY, Contexts.EMPTY_CONTEXT);
+    cluster = RelOptCluster.create(planner, new RexBuilder(factory));
+  }
+
+  public RelCreator(List<String> creates) {
+    super(SqlConverterBase.FEATURES_DEFAULT);
+    RelDataTypeFactory factory = new JavaTypeFactoryImpl(SubstraitTypeSystem.TYPE_SYSTEM);
+    try {
+      Pair<SqlValidator, CalciteCatalogReader> pair = this.registerCreateTables(creates);
+      this.validator = pair.left;
+      this.catalog = pair.right;
+    } catch (SqlParseException e) {
+      throw new RuntimeException(e);
+    }
     VolcanoPlanner planner = new VolcanoPlanner(RelOptCostImpl.FACTORY, Contexts.EMPTY_CONTEXT);
     cluster = RelOptCluster.create(planner, new RexBuilder(factory));
   }
@@ -51,8 +70,6 @@ public RelRoot parse(String sql) {
           () ->
               new RelMetadataQuery(
                   new ProxyingMetadataHandlerProvider(DefaultRelMetadataProvider.INSTANCE)));
-      SqlValidator validator =
-          new Validator(catalog, cluster.getTypeFactory(), SqlValidator.Config.DEFAULT);
 
       SqlToRelConverter.Config converterConfig =
           SqlToRelConverter.config().withTrimUnusedFields(true).withExpand(false);

Diff for: isthmus/src/test/java/io/substrait/isthmus/SubstraitExpressionConverterTest.java

@@ -1,16 +1,22 @@
 package io.substrait.isthmus;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import io.substrait.dsl.SubstraitBuilder;
+import io.substrait.expression.Expression;
 import io.substrait.isthmus.expression.ExpressionRexConverter;
-import io.substrait.isthmus.expression.ScalarFunctionConverter;
-import io.substrait.isthmus.expression.WindowFunctionConverter;
 import io.substrait.relation.Rel;
+import io.substrait.relation.Rel.Remap;
 import io.substrait.type.Type;
 import io.substrait.type.TypeCreator;
 import java.util.List;
+import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexSubQuery;
+import org.apache.calcite.sql.SqlKind;
 import org.junit.jupiter.api.Test;
 
 public class SubstraitExpressionConverterTest extends PlanTestBase {
@@ -20,17 +26,19 @@ public class SubstraitExpressionConverterTest extends PlanTestBase {
 
   final SubstraitBuilder b = new SubstraitBuilder(extensions);
 
-  final ExpressionRexConverter converter =
-      new ExpressionRexConverter(
-          typeFactory,
-          new ScalarFunctionConverter(extensions.scalarFunctions(), typeFactory),
-          new WindowFunctionConverter(extensions.windowFunctions(), typeFactory),
-          TypeConverter.DEFAULT);
+  final ExpressionRexConverter converter;
 
   final List<Type> commonTableType = List.of(R.I32, R.FP32, N.STRING, N.BOOLEAN);
   final Rel commonTable =
       b.namedScan(List.of("example"), List.of("a", "b", "c", "d"), commonTableType);
 
+  final SubstraitRelNodeConverter relNodeConverter =
+      new SubstraitRelNodeConverter(extensions, typeFactory, builder);
+
+  public SubstraitExpressionConverterTest() {
+    converter = relNodeConverter.expressionRexConverter;
+  }
+
   @Test
   public void switchExpression() {
     var expr =
@@ -43,6 +51,54 @@ public void switchExpression() {
     assertTypeMatch(calciteExpr.getType(), N.BOOLEAN);
   }
 
+  @Test
+  public void scalarSubQuery() {
+    /*
+     * equivalent to:
+     *
+     * select
+     *   r_regionkey
+     * from
+     *   region
+     * where
+     *   r_name = 'EUROPE'
+     */
+    Rel subQueryRel =
+        b.project(
+            input -> {
+              return List.of(b.fieldReference(input, 0));
+            },
+            // currently, all columns of the input are emitted first and then the expressions of
+            // this project unless remap is configured
+            Remap.of(List.of(3)),
+            b.filter(
+                input -> {
+                  return b.equal(
+                      b.fieldReference(input, 1),
+                      Expression.StrLiteral.builder().nullable(false).value("EUROPE").build());
+                },
+                b.namedScan(
+                    List.of("REGION"),
+                    List.of("r_regionkey", "r_name", "r_comment"),
+                    List.of(
+                        TypeCreator.REQUIRED.I64,
+                        TypeCreator.NULLABLE.fixedChar(25),
+                        TypeCreator.NULLABLE.varChar(152)))));
+
+    Expression.ScalarSubquery expr =
+        Expression.ScalarSubquery.builder()
+            .type(TypeCreator.REQUIRED.I64)
+            .input(subQueryRel)
+            .build();
+
+    RexNode calciteExpr = expr.accept(converter);
+    assertEquals(SqlKind.SCALAR_QUERY, calciteExpr.getKind());
+    assertInstanceOf(RexSubQuery.class, calciteExpr);
+
+    RelNode subQueryRelNode = subQueryRel.accept(relNodeConverter);
+    assertTrue(subQueryRelNode.deepEquals(((RexSubQuery) calciteExpr).rel));
+  }
+
   void assertTypeMatch(RelDataType actual, Type expected) {
     Type type = TypeConverter.DEFAULT.toSubstrait(actual);
     assertEquals(expected, type);