fix: potential fix for decimal template error

Signed-off-by: Andrew Coleman <[email protected]>
substrait-io · Sep 19, 2024 · 2807c2b · 2807c2b
1 parent 7a4ced9
commit 2807c2b
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/core/src/main/java/io/substrait/type/parser/ParseToPojo.java b/core/src/main/java/io/substrait/type/parser/ParseToPojo.java
@@ -261,6 +261,16 @@ public TypeExpression visitDecimal(final SubstraitTypeParser.DecimalContext ctx)
         return withNullP(nullable).decimalE((String) precision, (String) scale);
       }
 
+      if (precision instanceof String && scale instanceof Integer) {
+        checkParameterizedOrExpression();
+        return withNullP(nullable).decimalE((String) precision, String.valueOf(scale));
+      }
+
+      if (precision instanceof Integer && scale instanceof String) {
+        checkParameterizedOrExpression();
+        return withNullP(nullable).decimalE(String.valueOf(precision), (String) scale);
+      }
+
       checkExpression();
       return withNullE(nullable).decimalE(ctx.precision.accept(this), ctx.scale.accept(this));
     }

diff --git a/core/src/test/java/io/substrait/type/parser/TestTypeParser.java b/core/src/test/java/io/substrait/type/parser/TestTypeParser.java
@@ -119,6 +119,9 @@ private <T> void parameterizedTests(ParseToPojo.Visitor v) {
     test(v, pn.listE(pr.parameter("any")), "list?<any>");
     test(v, pn.listE(pn.parameter("any")), "list?<any?>");
     test(v, pn.structE(r.I8, r.I16, n.I8, pr.parameter("K")), "STRUCT?<i8, i16, i8?, K>");
+    test(v, pr.decimalE("P", "S"), "DECIMAL<P, S>");
+    test(v, pr.decimalE("P", "0"), "DECIMAL<P, 0>");
+    test(v, pr.decimalE("14", "S"), "DECIMAL<14, S>");
   }
 
   private static void test(ParseToPojo.Visitor visitor, TypeExpression expected, String toParse) {