From 6c207d1a698570bb4a74c3a3f8f04ce1748c4529 Mon Sep 17 00:00:00 2001 From: Yi Hu Date: Tue, 19 Nov 2024 15:21:22 -0500 Subject: [PATCH] Bump zetasql version to 2024.11.1 --- sdks/java/extensions/sql/zetasql/build.gradle | 2 +- .../zetasql/translation/AggregateScanConverter.java | 7 ++++--- .../zetasql/translation/ExpressionConverter.java | 13 +++++++++++-- 3 files changed, 16 insertions(+), 6 deletions(-) diff --git a/sdks/java/extensions/sql/zetasql/build.gradle b/sdks/java/extensions/sql/zetasql/build.gradle index 8d6e2aac0bf4..29a3f95402b0 100644 --- 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/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/translation/AggregateScanConverter.java b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/translation/AggregateScanConverter.java index 2dfc7fe372f5..412cd46001f8 100644 --- a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/translation/AggregateScanConverter.java +++ b/sdks/java/extensions/sql/zetasql/src/main/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 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/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/translation/ExpressionConverter.java b/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/translation/ExpressionConverter.java index 14db554d6f0b..0f32451504b3 100644 --- a/sdks/java/extensions/sql/zetasql/src/main/java/org/apache/beam/sdk/extensions/sql/zetasql/translation/ExpressionConverter.java +++ b/sdks/java/extensions/sql/zetasql/src/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 cols; + try { + cols = (List) FieldUtils.readField(rel, "columns", true); + } catch (IllegalAccessException e) { + throw new RuntimeException(e); + } List relDataTypeFields = - convertTVFRelationColumnsToRelDataTypeFields( - fixedOutputSchemaTVF.getOutputSchema().getColumns()); + convertTVFRelationColumnsToRelDataTypeFields(cols); return new RelRecordType(relDataTypeFields); }, null,