chore(isthmus): handle nullability and EnumArgument in SimplExtensionToSqlOperator

Author: ZorinAnton

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

@@ -1,29 +1,18 @@
 package io.substrait.isthmus;
 
-import io.substrait.isthmus.sql.SubstraitSqlToCalcite;
-import com.google.common.annotations.VisibleForTesting;
 import io.substrait.extension.SimpleExtension;
 import io.substrait.isthmus.calcite.SubstraitOperatorTable;
-import io.substrait.isthmus.sql.SubstraitSqlValidator;
+import io.substrait.isthmus.sql.SubstraitSqlToCalcite;
 import io.substrait.plan.ImmutablePlan.Builder;
 import io.substrait.plan.Plan;
 import io.substrait.plan.Plan.Version;
 import io.substrait.plan.PlanProtoConverter;
+import java.util.List;
 import org.apache.calcite.prepare.Prepare;
-import org.apache.calcite.rel.RelRoot;
-import org.apache.calcite.rel.rules.CoreRules;
-import org.apache.calcite.sql.SqlNode;
-import org.apache.calcite.sql.SqlNodeList;
 import org.apache.calcite.sql.SqlOperator;
 import org.apache.calcite.sql.SqlOperatorTable;
 import org.apache.calcite.sql.parser.SqlParseException;
-import org.apache.calcite.sql.parser.SqlParser;
 import org.apache.calcite.sql.util.SqlOperatorTables;
-import org.apache.calcite.sql.validate.SqlValidator;
-import org.apache.calcite.sql2rel.SqlToRelConverter;
-import org.apache.calcite.sql2rel.StandardConvertletTable;
-
-import java.util.List;
 
 /** Take a SQL statement and a set of table definitions and return a substrait plan. */
 public class SqlToSubstrait extends SqlConverterBase {
@@ -84,9 +73,9 @@ public Plan convert(String sqlStatements, Prepare.CatalogReader catalogReader)
     builder.version(Version.builder().from(Version.DEFAULT_VERSION).producer("isthmus").build());
 
     // TODO: consider case in which one sql passes conversion while others don't
-      SubstraitSqlToCalcite.convertQueries(sqlStatements, catalogReader).stream()
-              .map(root -> SubstraitRelVisitor.convert(root, extensionCollection, featureBoard))
-              .forEach(root -> builder.addRoots(root));
+    SubstraitSqlToCalcite.convertQueries(sqlStatements, catalogReader, operatorTable).stream()
+        .map(root -> SubstraitRelVisitor.convert(root, extensionCollection, featureBoard))
+        .forEach(root -> builder.addRoots(root));
 
     return builder.build();
   }

isthmus/src/main/java/io/substrait/isthmus/sql/SubstraitSqlToCalcite.java

@@ -15,6 +15,7 @@
 import org.apache.calcite.rel.rules.CoreRules;
 import org.apache.calcite.rex.RexBuilder;
 import org.apache.calcite.sql.SqlNode;
+import org.apache.calcite.sql.SqlOperatorTable;
 import org.apache.calcite.sql.parser.SqlParseException;
 import org.apache.calcite.sql.validate.SqlValidator;
 import org.apache.calcite.sql2rel.SqlToRelConverter;
@@ -41,6 +42,23 @@ public static RelRoot convertQuery(String sqlStatement, Prepare.CatalogReader ca
     return convertQuery(sqlStatement, catalogReader, validator, createDefaultRelOptCluster());
   }
 
+  /**
+   * Converts a SQL statement to a Calcite {@link RelRoot}.
+   *
+   * @param sqlStatement a SQL statement string
+   * @param catalogReader the {@link Prepare.CatalogReader} for finding tables/views referenced in
+   *     the SQL statement
+   * @param operatorTable the {@link SqlOperatorTable} for dynamic operators
+   * @return a {@link RelRoot} corresponding to the given SQL statement
+   * @throws SqlParseException if there is an error while parsing the SQL statement
+   */
+  public static RelRoot convertQuery(
+      String sqlStatement, Prepare.CatalogReader catalogReader, SqlOperatorTable operatorTable)
+      throws SqlParseException {
+    SqlValidator validator = new SubstraitSqlValidator(catalogReader, operatorTable);
+    return convertQuery(sqlStatement, catalogReader, validator, createDefaultRelOptCluster());
+  }
+
   /**
    * Converts a SQL statement to a Calcite {@link RelRoot}.
    *
@@ -72,6 +90,24 @@ public static RelRoot convertQuery(
     return relRoots.get(0);
   }
 
+  /**
+   * Converts one or more SQL statements to a List of {@link RelRoot}, with one {@link RelRoot} per
+   * statement.
+   *
+   * @param sqlStatements a string containing one or more SQL statements
+   * @param catalogReader the {@link Prepare.CatalogReader} for finding tables/views referenced in
+   *     the SQL statements
+   * @param operatorTable the {@link SqlOperatorTable} for dynamic operators
+   * @return a list of {@link RelRoot}s corresponding to the given SQL statements
+   * @throws SqlParseException if there is an error while parsing the SQL statements
+   */
+  public static List<RelRoot> convertQueries(
+      String sqlStatements, Prepare.CatalogReader catalogReader, SqlOperatorTable operatorTable)
+      throws SqlParseException {
+    SqlValidator validator = new SubstraitSqlValidator(catalogReader, operatorTable);
+    return convertQueries(sqlStatements, catalogReader, validator, createDefaultRelOptCluster());
+  }
+
   /**
    * Converts one or more SQL statements to a List of {@link RelRoot}, with one {@link RelRoot} per
    * statement.

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

@@ -5,7 +5,6 @@
 import io.substrait.extension.SimpleExtension;
 import io.substrait.isthmus.sql.SubstraitSqlToCalcite;
 import java.io.IOException;
-import java.util.List;
 import org.apache.calcite.plan.hep.HepPlanner;
 import org.apache.calcite.plan.hep.HepProgram;
 import org.apache.calcite.plan.hep.HepProgramBuilder;

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

@@ -167,12 +167,10 @@ protected RelRoot assertSqlSubstraitRelRoundTripWorkaroundOptimizer(
     SqlToSubstrait s = new SqlToSubstrait(extensions, null);
 
     // 1. SQL -> Calcite RelRoot
-    List<RelRoot> relRoots = s.sqlToRelNode(query, catalogReader);
-    assertEquals(1, relRoots.size());
-    RelRoot relRoot1 = relRoots.get(0);
+    Plan plan1 = s.convert(query, catalogReader);
 
     // 2. Calcite RelRoot  -> Substrait Rel
-    Plan.Root pojo1 = SubstraitRelVisitor.convert(relRoot1, extensions);
+    Plan.Root pojo1 = plan1.getRoots().get(0);
 
     // 3. Substrait Rel -> Calcite RelNode
     RelRoot relRoot2 = substraitToCalcite.convert(pojo1);