Bump zetasql version to 2024.11.1 (apache#32902)

Polber · Nov 20, 2024 · 5d3088f · 5d3088f
1 parent d2e9928
commit 5d3088f
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 6 deletions.
diff --git a/sdks/java/extensions/sql/zetasql/build.gradle b/sdks/java/extensions/sql/zetasql/build.gradle
@@ -27,7 +27,7 @@ applyJavaNature(
 description = "Apache Beam :: SDKs :: Java :: Extensions :: SQL :: ZetaSQL"
 ext.summary = "ZetaSQL to Calcite translator"
 
-def zetasql_version = "2022.04.1"
+def zetasql_version = "2024.11.1"
 
 dependencies {
   // TODO(https://github.com/apache/beam/issues/21156): Determine how to build without this dependency

diff --git a/...n/java/org/apache/beam/sdk/extensions/sql/zetasql/translation/AggregateScanConverter.java b/...n/java/org/apache/beam/sdk/extensions/sql/zetasql/translation/AggregateScanConverter.java
@@ -29,6 +29,7 @@
 import com.google.zetasql.resolvedast.ResolvedNodes.ResolvedAggregateFunctionCall;
 import com.google.zetasql.resolvedast.ResolvedNodes.ResolvedAggregateScan;
 import com.google.zetasql.resolvedast.ResolvedNodes.ResolvedComputedColumn;
+import com.google.zetasql.resolvedast.ResolvedNodes.ResolvedComputedColumnBase;
 import com.google.zetasql.resolvedast.ResolvedNodes.ResolvedExpr;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -94,7 +95,7 @@ public RelNode convert(ResolvedAggregateScan zetaNode, List<RelNode> inputs) {
       aggregateCalls = new ArrayList<>();
       // For aggregate calls, their input ref follow after GROUP BY input ref.
       int columnRefoff = groupFieldsListSize;
-      for (ResolvedComputedColumn computedColumn : zetaNode.getAggregateList()) {
+      for (ResolvedComputedColumnBase computedColumn : zetaNode.getAggregateList()) {
         AggregateCall aggCall =
             convertAggCall(computedColumn, columnRefoff, groupSet.size(), input);
         aggregateCalls.add(aggCall);
@@ -144,7 +145,7 @@ private LogicalProject convertAggregateScanInputScanToLogicalProject(
     // LogicalProject should also include columns used by aggregate functions. These columns should
     // follow after GROUP BY columns.
     // TODO: remove duplicate columns in projects.
-    for (ResolvedComputedColumn resolvedComputedColumn : node.getAggregateList()) {
+    for (ResolvedComputedColumnBase resolvedComputedColumn : node.getAggregateList()) {
       // Should create Calcite's RexInputRef from ResolvedColumn from ResolvedComputedColumn.
       // TODO: handle aggregate function with more than one argument and handle OVER
       // TODO: is there is general way for column reference tracking and deduplication for
@@ -180,7 +181,7 @@ private LogicalProject convertAggregateScanInputScanToLogicalProject(
   }
 
   private AggregateCall convertAggCall(
-      ResolvedComputedColumn computedColumn, int columnRefOff, int groupCount, RelNode input) {
+      ResolvedComputedColumnBase computedColumn, int columnRefOff, int groupCount, RelNode input) {
     ResolvedAggregateFunctionCall aggregateFunctionCall =
         (ResolvedAggregateFunctionCall) computedColumn.getExpr();
 

diff --git a/...main/java/org/apache/beam/sdk/extensions/sql/zetasql/translation/ExpressionConverter.java b/...main/java/org/apache/beam/sdk/extensions/sql/zetasql/translation/ExpressionConverter.java
@@ -37,6 +37,7 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.zetasql.TVFRelation;
+import com.google.zetasql.TVFRelation.Column;
 import com.google.zetasql.TableValuedFunction;
 import com.google.zetasql.TableValuedFunction.FixedOutputSchemaTVF;
 import com.google.zetasql.Type;
@@ -65,6 +66,7 @@
 import java.util.Map;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
+import org.apache.beam.repackaged.core.org.apache.commons.lang3.reflect.FieldUtils;
 import org.apache.beam.sdk.annotations.Internal;
 import org.apache.beam.sdk.extensions.sql.impl.QueryPlanner.QueryParameters;
 import org.apache.beam.sdk.extensions.sql.impl.ZetaSqlUserDefinedSQLNativeTableValuedFunction;
@@ -495,9 +497,16 @@ public RexCall convertTableValuedFunction(
                   new ZetaSqlUserDefinedSQLNativeTableValuedFunction(
                       new SqlIdentifier(tvf.getName(), SqlParserPos.ZERO),
                       opBinding -> {
+                        TVFRelation rel = fixedOutputSchemaTVF.getOutputSchema();
+                        // TODO(yathu) revert this workaround when ZetaSQL adds back this API.
+                        List<Column> cols;
+                        try {
+                          cols = (List<Column>) FieldUtils.readField(rel, "columns", true);
+                        } catch (IllegalAccessException e) {
+                          throw new RuntimeException(e);
+                        }
                         List<RelDataTypeField> relDataTypeFields =
-                            convertTVFRelationColumnsToRelDataTypeFields(
-                                fixedOutputSchemaTVF.getOutputSchema().getColumns());
+                            convertTVFRelationColumnsToRelDataTypeFields(cols);
                         return new RelRecordType(relDataTypeFields);
                       },
                       null,