add support for multiple argument types

prestodb · Nov 13, 2024 · b773c2e · b773c2e
1 parent 48a8c6d
commit b773c2e
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 6 deletions.
diff --git a/.../src/main/java/com/facebook/presto/expressions/translator/TranslatorAnnotationParser.java b/.../src/main/java/com/facebook/presto/expressions/translator/TranslatorAnnotationParser.java
@@ -145,9 +145,16 @@ private static List<FunctionMetadata> methodToFunctionMetadata(ScalarTranslation
 
         for (SqlSignature signature : signatures) {
             ImmutableList.Builder<TypeSignature> argumentTypes = new ImmutableList.Builder<>();
-            TypeSignature argumentType = parseTypeSignature(signature.argumentType());
-            for (int i = 0; i < method.getParameterCount(); i++) {
-                argumentTypes.add(argumentType);
+            if (signature.argumentTypes().length == method.getParameterCount()) {
+                for (int i = 0; i < method.getParameterCount(); i++) {
+                    argumentTypes.add(parseTypeSignature(signature.argumentTypes()[i]));
+                }
+            }
+            else {
+                TypeSignature argumentType = parseTypeSignature(signature.argumentType());
+                for (int i = 0; i < method.getParameterCount(); i++) {
+                    argumentTypes.add(argumentType);
+                }
             }
             TypeSignature returnType = parseTypeSignature(signature.returnType());
             FunctionMetadata derivedMetadata;

diff --git a/...to-main/src/test/java/com/facebook/presto/sql/relational/TestRowExpressionTranslator.java b/...to-main/src/test/java/com/facebook/presto/sql/relational/TestRowExpressionTranslator.java
@@ -158,7 +158,8 @@ public Object[][] createTestBasicOperatorData()
         return new Object[][] {
                 {BIGINT, BIGINT},
                 {INTEGER, INTEGER},
-                {decimalType, decimalType}
+                {decimalType, decimalType},
+                {BIGINT, INTEGER}
         };
     }
 
@@ -299,7 +300,8 @@ public static String not(String sql)
         @SupportedSignatures({
                 @SqlSignature(argumentType = StandardTypes.INTEGER, returnType = StandardTypes.INTEGER),
                 @SqlSignature(argumentType = StandardTypes.BIGINT, returnType = StandardTypes.BIGINT),
-                @SqlSignature(argumentType = "decimal(38, 0)", returnType = "decimal(38, 0)")})
+                @SqlSignature(argumentType = "decimal(38, 0)", returnType = "decimal(38, 0)"),
+                @SqlSignature(argumentTypes = {StandardTypes.BIGINT, StandardTypes.INTEGER}, returnType = StandardTypes.BIGINT)})
         public static String plus(String left, String right)
         {
             return left + " -|- " + right;

diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/function/SqlSignature.java b/presto-spi/src/main/java/com/facebook/presto/spi/function/SqlSignature.java
@@ -22,8 +22,9 @@
 @Retention(RUNTIME)
 @Target(METHOD)
 public @interface SqlSignature {
-    String argumentType();
+    String argumentType() default "";
     String returnType();
+    String[] argumentTypes() default {};
 
     Class<?> nativeContainerType() default Object.class;
 }