chore(isthmus): handle nullability and EnumArgument in SimplExtensionToSqlOperator

Author: ZorinAnton

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

@@ -6,6 +6,7 @@
 import io.substrait.function.TypeExpression;
 import io.substrait.type.Type;
 import io.substrait.type.TypeExpressionEvaluator;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -46,17 +47,19 @@ private static SqlFunction toSqlFunction(
       SimpleExtension.Function function,
       RelDataTypeFactory typeFactory,
       TypeConverter typeConverter) {
-    List<SimpleExtension.Argument> requiredArgs =
-        function.args().stream()
-            .filter(SimpleExtension.Argument::required)
-            .filter(t -> t instanceof SimpleExtension.ValueArgument || t instanceof SimpleExtension.EnumArgument)
-            .map(t -> (SimpleExtension.Argument) t)
-            .collect(Collectors.toList());
-
-    List<SqlTypeFamily> argFamilies =
-        requiredArgs.stream()
-            .map(arg -> arg.value().accept(new CalciteTypeVisitor()).getFamily())
-            .collect(Collectors.toList());
+
+    List<SqlTypeFamily> argFamilies = new ArrayList<>();
+
+    for (SimpleExtension.Argument arg : function.requiredArguments()) {
+      if (arg instanceof SimpleExtension.ValueArgument) {
+        SimpleExtension.ValueArgument valueArg = (SimpleExtension.ValueArgument) arg;
+        SqlTypeName typeName = valueArg.value().accept(new CalciteTypeVisitor());
+        argFamilies.add(typeName.getFamily());
+      } else if (arg instanceof SimpleExtension.EnumArgument) {
+        // Treat an EnumArgument as a required string literal.
+        argFamilies.add(SqlTypeFamily.STRING);
+      }
+    }
 
     SqlReturnTypeInference returnTypeInference =
         new SubstraitReturnTypeInference(function, typeFactory, typeConverter);
@@ -97,7 +100,27 @@ public RelDataType inferReturnType(SqlOperatorBinding opBinding) {
           TypeExpressionEvaluator.evaluateExpression(
               returnExpression, function.args(), substraitArgTypes);
 
-      return typeConverter.toCalcite(typeFactory, resolvedSubstraitType);
+      boolean finalIsNullable;
+      switch (function.nullability()) {
+        case MIRROR:
+          // If any input is nullable, the output is nullable.
+          finalIsNullable =
+              opBinding.collectOperandTypes().stream().anyMatch(RelDataType::isNullable);
+          break;
+        case DISCRETE:
+          // The function can return null even if inputs are not null.
+          finalIsNullable = true;
+          break;
+        case DECLARED_OUTPUT:
+        default:
+          // Use the nullability declared on the resolved Substrait type.
+          finalIsNullable = resolvedSubstraitType.nullable();
+          break;
+      }
+
+      RelDataType baseCalciteType = typeConverter.toCalcite(typeFactory, resolvedSubstraitType);
+
+      return typeFactory.createTypeWithNullability(baseCalciteType, finalIsNullable);
     }
   }
 

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

@@ -47,7 +47,7 @@ public SqlToSubstrait(SimpleExtension.ExtensionCollection extensions, FeatureBoa
 
     this.operatorTable =
         SqlOperatorTables.chain(
-            SqlOperatorTables.of(generatedDynamicOperators), SubstraitOperatorTable.INSTANCE);
+            SubstraitOperatorTable.INSTANCE, SqlOperatorTables.of(generatedDynamicOperators));
   }
 
   /**
@@ -94,7 +94,7 @@ public Plan convert(String sql, Prepare.CatalogReader catalogReader) throws SqlP
   @VisibleForTesting
   List<RelRoot> sqlToRelNode(String sql, Prepare.CatalogReader catalogReader)
       throws SqlParseException {
-    SqlValidator validator = new SubstraitSqlValidator(catalogReader);
+    SqlValidator validator = new SubstraitSqlValidator(catalogReader, operatorTable);
     SqlParser parser = SqlParser.create(sql, parserConfig);
     SqlNodeList parsedList = parser.parseStmtList();
     SqlToRelConverter converter = createSqlToRelConverter(validator, catalogReader);

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

@@ -16,7 +16,7 @@ public class SimpleExtensionToSqlOperatorTest {
 
   @Test
   void test() throws IOException {
-    String customFunctionPath = "/extensions/functions_string_custom.yaml";
+    String customFunctionPath = "/extensions/scalar_functions_custom.yaml";
 
     SimpleExtension.ExtensionCollection customExtensions =
         SimpleExtension.load(
@@ -25,8 +25,6 @@ void test() throws IOException {
 
     List<SqlOperator> operators = SimpleExtensionToSqlOperator.from(customExtensions);
 
-    assertEquals(1, operators.size(), "Should generate one operator from the YAML file.");
-
     Optional<SqlOperator> function =
         operators.stream()
             .filter(op -> op.getName().equalsIgnoreCase("REGEXP_EXTRACT"))

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

@@ -8,7 +8,7 @@
 
 public class UdfSqlSubstraitTest extends PlanTestBase {
 
-  private static final String CUSTOM_FUNCTION_PATH = "/extensions/functions_string_custom.yaml";
+  private static final String CUSTOM_FUNCTION_PATH = "/extensions/scalar_functions_custom.yaml";
 
   UdfSqlSubstraitTest() {
     super(loadExtensions(List.of(CUSTOM_FUNCTION_PATH)));
@@ -17,14 +17,18 @@ public class UdfSqlSubstraitTest extends PlanTestBase {
   @Test
   public void customUdfTest() throws Exception {
 
-    final String[] sql = {
-      "CREATE TABLE t(x VARCHAR NOT NULL)", "SELECT regexp_extract(x, 'ab') from t"
-    };
-
     final Prepare.CatalogReader catalogReader =
-        SubstraitCreateStatementParser.processCreateStatementsToCatalog(sql[0]);
-
-    assertSqlSubstraitRelRoundTripWorkaroundOptimizer(sql[1], catalogReader);
+        SubstraitCreateStatementParser.processCreateStatementsToCatalog(
+            "CREATE TABLE t(x VARCHAR NOT NULL)");
+
+    assertSqlSubstraitRelRoundTripWorkaroundOptimizer(
+        "SELECT regexp_extract(x, 'ab') from t", catalogReader);
+    assertSqlSubstraitRelRoundTripWorkaroundOptimizer(
+        "SELECT format_text('UPPER', x) FROM t", catalogReader);
+    assertSqlSubstraitRelRoundTripWorkaroundOptimizer(
+        "SELECT system_property_get(x) FROM t", catalogReader);
+    assertSqlSubstraitRelRoundTripWorkaroundOptimizer(
+        "SELECT safe_divide(10,0) FROM t", catalogReader);
   }
 
   private static SimpleExtension.ExtensionCollection loadExtensions(

isthmus/src/test/resources/extensions/functions_string_custom.yaml

@@ -0,0 +1,44 @@
+%YAML 1.2
+---
+scalar_functions:
+  - name: "regexp_extract"
+    impls:
+      - args:
+          - name: "text"
+            value: string
+          - name: "pattern"
+            value: string
+        return: string
+
+  - name: "format_text"
+    description: "Formats text based on a mode. The output is nullable if the input is."
+    impls:
+      - args:
+          - name: "mode"
+#            options: ["UPPER", "LOWER"]
+            value: string
+          - name: "input_text"
+#            options: ["UPPER", "LOWER"]
+            value: string
+        return: string
+        nullability: MIRROR
+
+  - name: "system_property_get"
+    description: "Safely gets a system property. Always returns a nullable string."
+    impls:
+      - args:
+          - name: "property_name"
+            value: string
+        return: string?
+        nullability: DECLARED_OUTPUT
+
+  - name: "safe_divide"
+    description: "Performs division, returning NULL if the denominator is zero."
+    impls:
+      - args:
+          - name: "numerator"
+            value: i32
+          - name: "denominator"
+            value: i32
+        return: fp32?
+        nullability: DISCRETE