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



##########
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:
       It actually appears there is a block of code common with scalar 
functions here that could be shared... Something like:
   ```
   FunctionArgumentType resultType =
       new FunctionArgumentType(
             ZetaSqlCalciteTranslationUtils.toZetaSqlType(
                 function.getReturnType(typeFactory)));
   FunctionSignature functionSignature =
       new FunctionSignature(resultType, getArgumentTypes(function), 0L);
   ```

##########
File path: 
sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/translation/AggregateScanConverter.java
##########
@@ -200,14 +203,14 @@ private AggregateCall convertAggCall(
               + " aggregation.");
     }
 
-    SqlAggFunction sqlAggFunction =
-        (SqlAggFunction)
-            
SqlOperatorMappingTable.ZETASQL_FUNCTION_TO_CALCITE_SQL_OPERATOR.get(
-                aggregateFunctionCall.getFunction().getName());
-    if (sqlAggFunction == null) {
-      throw new UnsupportedOperationException(
-          "Does not support ZetaSQL aggregate function: "
-              + aggregateFunctionCall.getFunction().getName());
+    SqlAggFunction sqlAggFunction;

Review comment:
       nit: consider `final`?

##########
File path: 
sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/translation/AggregateScanConverter.java
##########
@@ -234,4 +237,41 @@ private AggregateCall convertAggCall(
     return AggregateCall.create(
         sqlAggFunction, false, false, false, argList, -1, RelCollations.EMPTY, 
returnType, aggName);
   }
+
+  private SqlAggFunction getBuiltinFunction(ResolvedAggregateFunctionCall 
aggregateFunctionCall) {

Review comment:
       nit: In my opinion moving these into functions reduces the readability 
of this code. These are small functions called in exactly one place, consider 
inlining and adding a comment if you think it is needed.

##########
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:
       nit: Looks like you could just move this to before the `if` block like 
you did with path rather than making it a function? I have no strong opinion 
here but that might make this a bit simpler.

##########
File path: 
sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/BeamZetaSqlCatalog.java
##########
@@ -103,6 +107,7 @@
 
   private final Map<List<String>, UserFunctionDefinitions.JavaScalarFunction> 
javaScalarUdfs =
       new HashMap<>();
+  private final Map<List<String>, Combine.CombineFn<?, ?, ?>> javaUdafs = new 
HashMap<>();

Review comment:
       I'm suspicious of the `Combine.CombineFn` value. Is this something you 
were going to change in a followup?




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