ibzib commented on a change in pull request #14588:
URL: https://github.com/apache/beam/pull/14588#discussion_r617734308



##########
File path: 
sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/BeamZetaSqlCatalog.java
##########
@@ -379,6 +391,16 @@ private void addUdfsFromSchema() {
     }
   }
 
+  private List<FunctionArgumentType> getArgumentTypes(

Review comment:
       done

##########
File path: 
sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/BeamZetaSqlCatalog.java
##########
@@ -341,34 +347,40 @@ private void addUdfsFromSchema() {
       }
       for 
(org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.schema.Function 
function :
           functions) {
+        List<String> path = Arrays.asList(functionName.split("\\."));
         if (function instanceof ScalarFunctionImpl) {
           ScalarFunctionImpl scalarFunction = (ScalarFunctionImpl) function;
           // Validate types before converting from Calcite to ZetaSQL, since 
the conversion may fail
           // for unsupported types.
           validateScalarFunctionImpl(scalarFunction);
-          List<String> path = Arrays.asList(functionName.split("\\."));
           Method method = scalarFunction.method;
           javaScalarUdfs.put(path, 
UserFunctionDefinitions.JavaScalarFunction.create(method, ""));
           FunctionArgumentType resultType =
               new FunctionArgumentType(
                   ZetaSqlCalciteTranslationUtils.toZetaSqlType(
                       scalarFunction.getReturnType(typeFactory)));
-          List<FunctionArgumentType> argumentTypes =
-              scalarFunction.getParameters().stream()
-                  .map(
-                      (arg) ->
-                          new FunctionArgumentType(
-                              ZetaSqlCalciteTranslationUtils.toZetaSqlType(
-                                  arg.getType(typeFactory))))
-                  .collect(Collectors.toList());
           FunctionSignature functionSignature =
-              new FunctionSignature(resultType, argumentTypes, 0L);
+              new FunctionSignature(resultType, 
getArgumentTypes(scalarFunction), 0L);
           zetaSqlCatalog.addFunction(
               new Function(
                   path,
                   USER_DEFINED_JAVA_SCALAR_FUNCTIONS,
                   ZetaSQLFunctions.FunctionEnums.Mode.SCALAR,
                   ImmutableList.of(functionSignature)));
+        } else if (function instanceof UdafImpl) {
+          UdafImpl<?, ?, ?> udaf = (UdafImpl) function;
+          javaUdafs.put(path, udaf.getCombineFn());
+          FunctionArgumentType resultType =

Review comment:
       `getReturnType` is defined separately for 
[ScalarFunction](https://github.com/apache/calcite/blob/df95257f357e67eac330ba4c6d6bc749851e99fb/core/src/main/java/org/apache/calcite/schema/ScalarFunction.java#L32)
 and 
[AggregateFunction](https://github.com/apache/calcite/blob/df95257f357e67eac330ba4c6d6bc749851e99fb/core/src/main/java/org/apache/calcite/schema/AggregateFunction.java#L32).
 Maybe if Java had Go-style implicit interfaces...




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
[email protected]


Reply via email to