[GitHub] [beam] amaliujia commented on a change in pull request #12766: [BEAM-10459][BEAM-9514] Fix AssertionError type mismatch from AggregateScanConverter

[GitBox <gi...@apache.org>]

amaliujia commented on a change in pull request #12766:
URL: https://github.com/apache/beam/pull/12766#discussion_r487307937



##########
File path: sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/translation/AggregateScanConverter.java
##########
@@ -227,8 +224,35 @@ private AggregateCall convertAggCall(
         ZetaSqlCalciteTranslationUtils.toCalciteType(
             computedColumn.getColumn().getType(), nullable, getCluster().getRexBuilder());
 
+    SqlAggFunction sqlAggFunction =
+        getSqlAggFunction(aggregateFunctionCall.getFunction().getName(), returnType);
+
     String aggName = getTrait().resolveAlias(computedColumn.getColumn());
     return AggregateCall.create(
         sqlAggFunction, false, false, false, argList, -1, RelCollations.EMPTY, returnType, aggName);
   }
+
+  private static SqlAggFunction getSqlAggFunction(
+      String zetaSqlAggFunctionName, RelDataType returnType) {
+    // ZetaSQL specific aggregation functions, implemented with a user-defined CombineFn
+    if (ZETASQL_UDAF_OPERATORS.containsKey(zetaSqlAggFunctionName)) {
+      return ZETASQL_UDAF_OPERATORS.get(zetaSqlAggFunctionName);
+    }
+
+    if ("$count_star".equals(zetaSqlAggFunctionName)) {
+      zetaSqlAggFunctionName = "COUNT"; // $count_star and count both map to the same implementation
+    } else {
+      // BeamBuiltinAggregations.BUILTIN_AGGREGATOR_FACTORIES uses upper-case function names
+      zetaSqlAggFunctionName = zetaSqlAggFunctionName.toUpperCase();
+    }
+
+    // Beam builtin aggregation functions (available in both ZetaSQL and CalciteSQL), implemented in
+    // {@link org.apache.beam.sdk.extensions.sql.impl.transform.BeamBuiltinAggregations}
+    if (BeamBuiltinAggregations.BUILTIN_AGGREGATOR_FACTORIES.containsKey(zetaSqlAggFunctionName)) {

Review comment:
       Is there a need to use the upper case of `zetaSqlAggFunctionName`?




----------------------------------------------------------------
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:
users@infra.apache.org