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



##########
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:
       +1 for rewrite in go. :)




-- 
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