| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| package org.apache.beam.sdk.extensions.sql.zetasql.translation; |
| |
| import static com.google.zetasql.ZetaSQLResolvedNodeKind.ResolvedNodeKind.RESOLVED_CAST; |
| import static com.google.zetasql.ZetaSQLResolvedNodeKind.ResolvedNodeKind.RESOLVED_COLUMN_REF; |
| import static com.google.zetasql.ZetaSQLResolvedNodeKind.ResolvedNodeKind.RESOLVED_GET_STRUCT_FIELD; |
| import static org.apache.beam.sdk.extensions.sql.zetasql.TypeUtils.toSimpleRelDataType; |
| |
| import com.google.zetasql.FunctionSignature; |
| import com.google.zetasql.ZetaSQLType.TypeKind; |
| import com.google.zetasql.resolvedast.ResolvedNode; |
| 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.ResolvedExpr; |
| import java.util.ArrayList; |
| import java.util.Collections; |
| import java.util.List; |
| import java.util.stream.Collectors; |
| import java.util.stream.IntStream; |
| import org.apache.beam.sdk.extensions.sql.zetasql.SqlStdOperatorMappingTable; |
| import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.RelNode; |
| import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.core.AggregateCall; |
| import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.logical.LogicalAggregate; |
| import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.logical.LogicalProject; |
| import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.type.RelDataType; |
| import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rex.RexNode; |
| import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.SqlAggFunction; |
| import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.sql.fun.SqlStdOperatorTable; |
| import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.util.ImmutableBitSet; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList; |
| |
| /** Converts aggregate calls. */ |
| class AggregateScanConverter extends RelConverter<ResolvedAggregateScan> { |
| private static final String AVG_ILLEGAL_LONG_INPUT_TYPE = |
| "AVG(LONG) is not supported. You might want to use AVG(CAST(expression AS DOUBLE)."; |
| |
| AggregateScanConverter(ConversionContext context) { |
| super(context); |
| } |
| |
| @Override |
| public List<ResolvedNode> getInputs(ResolvedAggregateScan zetaNode) { |
| return Collections.singletonList(zetaNode.getInputScan()); |
| } |
| |
| @Override |
| public RelNode convert(ResolvedAggregateScan zetaNode, List<RelNode> inputs) { |
| RelNode input = convertAggregateScanInputScanToLogicalProject(zetaNode, inputs.get(0)); |
| |
| // Calcite LogicalAggregate's GroupSet is indexes of group fields starting from 0. |
| int groupFieldsListSize = zetaNode.getGroupByList().size(); |
| ImmutableBitSet groupSet; |
| if (groupFieldsListSize != 0) { |
| groupSet = |
| ImmutableBitSet.of( |
| IntStream.rangeClosed(0, groupFieldsListSize - 1) |
| .boxed() |
| .collect(Collectors.toList())); |
| } else { |
| groupSet = ImmutableBitSet.of(); |
| } |
| |
| // TODO: add support for indicator |
| |
| List<AggregateCall> aggregateCalls; |
| if (zetaNode.getAggregateList().isEmpty()) { |
| aggregateCalls = ImmutableList.of(); |
| } else { |
| aggregateCalls = new ArrayList<>(); |
| // For aggregate calls, their input ref follow after GROUP BY input ref. |
| int columnRefoff = groupFieldsListSize; |
| for (ResolvedComputedColumn computedColumn : zetaNode.getAggregateList()) { |
| aggregateCalls.add(convertAggCall(computedColumn, columnRefoff)); |
| columnRefoff++; |
| } |
| } |
| |
| LogicalAggregate logicalAggregate = |
| new LogicalAggregate( |
| getCluster(), |
| input.getTraitSet(), |
| input, |
| false, |
| groupSet, |
| ImmutableList.of(groupSet), |
| aggregateCalls); |
| |
| return logicalAggregate; |
| } |
| |
| private RelNode convertAggregateScanInputScanToLogicalProject( |
| ResolvedAggregateScan node, RelNode input) { |
| // AggregateScan's input is the source of data (e.g. TableScan), which is different from the |
| // design of CalciteSQL, in which the LogicalAggregate's input is a LogicalProject, whose input |
| // is a LogicalTableScan. When AggregateScan's input is WithRefScan, the WithRefScan is |
| // ebullient to a LogicalTableScan. So it's still required to build another LogicalProject as |
| // the input of LogicalAggregate. |
| List<RexNode> projects = new ArrayList<>(); |
| List<String> fieldNames = new ArrayList<>(); |
| |
| // LogicalProject has a list of expr, which including UDF in GROUP BY clause for |
| // LogicalAggregate. |
| for (ResolvedComputedColumn computedColumn : node.getGroupByList()) { |
| projects.add( |
| getExpressionConverter() |
| .convertRexNodeFromResolvedExpr( |
| computedColumn.getExpr(), |
| node.getInputScan().getColumnList(), |
| input.getRowType().getFieldList())); |
| fieldNames.add(getTrait().resolveAlias(computedColumn.getColumn())); |
| } |
| |
| // 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()) { |
| // 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 |
| // aggregation? |
| ResolvedAggregateFunctionCall aggregateFunctionCall = |
| ((ResolvedAggregateFunctionCall) resolvedComputedColumn.getExpr()); |
| if (aggregateFunctionCall.getArgumentList() != null |
| && aggregateFunctionCall.getArgumentList().size() == 1) { |
| ResolvedExpr resolvedExpr = aggregateFunctionCall.getArgumentList().get(0); |
| |
| // TODO: assume aggregate function's input is either a ColumnRef or a cast(ColumnRef). |
| // TODO: user might use multiple CAST so we need to handle this rare case. |
| projects.add( |
| getExpressionConverter() |
| .convertRexNodeFromResolvedExpr( |
| resolvedExpr, |
| node.getInputScan().getColumnList(), |
| input.getRowType().getFieldList())); |
| fieldNames.add(getTrait().resolveAlias(resolvedComputedColumn.getColumn())); |
| } else if (aggregateFunctionCall.getArgumentList() != null |
| && aggregateFunctionCall.getArgumentList().size() > 1) { |
| throw new RuntimeException( |
| aggregateFunctionCall.getFunction().getName() + " has more than one argument."); |
| } |
| } |
| |
| return LogicalProject.create(input, projects, fieldNames); |
| } |
| |
| private AggregateCall convertAggCall(ResolvedComputedColumn computedColumn, int columnRefOff) { |
| ResolvedAggregateFunctionCall aggregateFunctionCall = |
| (ResolvedAggregateFunctionCall) computedColumn.getExpr(); |
| |
| // Reject AVG(INT64) |
| if (aggregateFunctionCall.getFunction().getName().equals("avg")) { |
| FunctionSignature signature = aggregateFunctionCall.getSignature(); |
| if (signature |
| .getFunctionArgumentList() |
| .get(0) |
| .getType() |
| .getKind() |
| .equals(TypeKind.TYPE_INT64)) { |
| throw new RuntimeException(AVG_ILLEGAL_LONG_INPUT_TYPE); |
| } |
| } |
| |
| // Reject aggregation DISTINCT |
| if (aggregateFunctionCall.getDistinct()) { |
| throw new RuntimeException( |
| "Does not support " |
| + aggregateFunctionCall.getFunction().getSqlName() |
| + " DISTINCT. 'SELECT DISTINCT' syntax could be used to deduplicate before" |
| + " aggregation."); |
| } |
| |
| SqlAggFunction sqlAggFunction = |
| (SqlAggFunction) |
| SqlStdOperatorMappingTable.ZETASQL_FUNCTION_TO_CALCITE_SQL_OPERATOR.get( |
| aggregateFunctionCall.getFunction().getName()); |
| if (sqlAggFunction == null) { |
| throw new RuntimeException( |
| "Does not support ZetaSQL aggregate function: " |
| + aggregateFunctionCall.getFunction().getName()); |
| } |
| |
| List<Integer> argList = new ArrayList<>(); |
| for (ResolvedExpr expr : |
| ((ResolvedAggregateFunctionCall) computedColumn.getExpr()).getArgumentList()) { |
| // Throw an error if aggregate function's input isn't either a ColumnRef or a cast(ColumnRef). |
| // TODO: is there a general way to handle aggregation calls conversion? |
| if (expr.nodeKind() == RESOLVED_CAST |
| || expr.nodeKind() == RESOLVED_COLUMN_REF |
| || expr.nodeKind() == RESOLVED_GET_STRUCT_FIELD) { |
| argList.add(columnRefOff); |
| } else { |
| throw new RuntimeException( |
| "Aggregate function only accepts Column Reference or CAST(Column Reference) as its" |
| + " input."); |
| } |
| } |
| |
| // TODO: there should be a general way to decide if a return type of a aggcall is nullable. |
| RelDataType returnType; |
| if (sqlAggFunction.equals(SqlStdOperatorTable.ANY_VALUE)) { |
| returnType = |
| toSimpleRelDataType( |
| computedColumn.getColumn().getType().getKind(), getCluster().getRexBuilder(), true); |
| } else { |
| returnType = |
| toSimpleRelDataType( |
| computedColumn.getColumn().getType().getKind(), getCluster().getRexBuilder(), false); |
| } |
| |
| String aggName = getTrait().resolveAlias(computedColumn.getColumn()); |
| return AggregateCall.create(sqlAggFunction, false, false, argList, -1, returnType, aggName); |
| } |
| } |
