Search in sources :

Example 1 with ResolvedColumn

use of com.google.zetasql.resolvedast.ResolvedColumn in project beam by apache.

the class ExpressionConverter method extractWatermarkColumnFromDescriptor.

private ResolvedColumn extractWatermarkColumnFromDescriptor(ResolvedNodes.ResolvedDescriptor descriptor) {
    ResolvedColumn wmCol = descriptor.getDescriptorColumnList().get(0);
    checkArgument(wmCol.getType().getKind() == TYPE_TIMESTAMP, "Watermarked column should be TIMESTAMP type: %s", descriptor.getDescriptorColumnNameList().get(0));
    return wmCol;
}
Also used : ResolvedColumn(com.google.zetasql.resolvedast.ResolvedColumn)

Example 2 with ResolvedColumn

use of com.google.zetasql.resolvedast.ResolvedColumn in project beam by apache.

the class ExpressionConverter method retrieveRexNode.

/**
 * Extract expressions from a project scan node.
 */
public List<RexNode> retrieveRexNode(ResolvedProjectScan node, List<RelDataTypeField> fieldList) {
    List<RexNode> ret = new ArrayList<>();
    for (ResolvedColumn column : node.getColumnList()) {
        int index = -1;
        if ((index = indexOfResolvedColumnInExprList(node.getExprList(), column)) != -1) {
            ResolvedComputedColumn computedColumn = node.getExprList().get(index);
            int windowFieldIndex = -1;
            if (computedColumn.getExpr().nodeKind() == RESOLVED_FUNCTION_CALL) {
                String functionName = ((ResolvedFunctionCall) computedColumn.getExpr()).getFunction().getName();
                if (WINDOW_START_END_FUNCTION_SET.contains(functionName)) {
                    ResolvedAggregateScan resolvedAggregateScan = (ResolvedAggregateScan) node.getInputScan();
                    windowFieldIndex = indexOfWindowField(resolvedAggregateScan.getGroupByList(), resolvedAggregateScan.getColumnList(), WINDOW_START_END_TO_WINDOW_MAP.get(functionName));
                }
            }
            ret.add(convertRexNodeFromComputedColumnWithFieldList(computedColumn, node.getInputScan().getColumnList(), fieldList, windowFieldIndex));
        } else {
            // ResolvedColumn is not a expression, which means it has to be an input column reference.
            index = indexOfProjectionColumnRef(column.getId(), node.getInputScan().getColumnList());
            if (index < 0 || index >= node.getInputScan().getColumnList().size()) {
                throw new IllegalStateException(String.format("Cannot find %s in fieldList %s", column, fieldList));
            }
            ret.add(rexBuilder().makeInputRef(fieldList.get(index).getType(), index));
        }
    }
    return ret;
}
Also used : ResolvedAggregateScan(com.google.zetasql.resolvedast.ResolvedNodes.ResolvedAggregateScan) ArrayList(java.util.ArrayList) ResolvedColumn(com.google.zetasql.resolvedast.ResolvedColumn) ResolvedComputedColumn(com.google.zetasql.resolvedast.ResolvedNodes.ResolvedComputedColumn) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)

Example 3 with ResolvedColumn

use of com.google.zetasql.resolvedast.ResolvedColumn in project beam by apache.

the class ExpressionConverter method retrieveRexNodeFromOrderByScan.

/**
 * Extract expressions from order by scan node.
 */
public List<RexNode> retrieveRexNodeFromOrderByScan(RelOptCluster cluster, ResolvedOrderByScan node, List<RelDataTypeField> fieldList) {
    final RexBuilder rexBuilder = cluster.getRexBuilder();
    List<RexNode> ret = new ArrayList<>();
    for (ResolvedColumn column : node.getColumnList()) {
        int index = indexOfProjectionColumnRef(column.getId(), node.getInputScan().getColumnList());
        ret.add(rexBuilder.makeInputRef(fieldList.get(index).getType(), index));
    }
    return ret;
}
Also used : ArrayList(java.util.ArrayList) RexBuilder(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexBuilder) ResolvedColumn(com.google.zetasql.resolvedast.ResolvedColumn) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)

Example 4 with ResolvedColumn

use of com.google.zetasql.resolvedast.ResolvedColumn in project beam by apache.

the class ExpressionConverter method convertTableValuedFunction.

/**
 * Convert a TableValuedFunction in ZetaSQL to a RexCall in Calcite.
 */
public RexCall convertTableValuedFunction(RelNode input, TableValuedFunction tvf, List<ResolvedNodes.ResolvedFunctionArgument> argumentList, List<ResolvedColumn> inputTableColumns) {
    ResolvedColumn wmCol;
    // Handle builtin windowing TVF.
    switch(tvf.getName()) {
        case TVFStreamingUtils.FIXED_WINDOW_TVF:
            // TUMBLE tvf's second argument is descriptor.
            wmCol = extractWatermarkColumnFromDescriptor(argumentList.get(1).getDescriptorArg());
            return (RexCall) rexBuilder().makeCall(new SqlWindowTableFunction(SqlKind.TUMBLE.name()), convertRelNodeToRexRangeRef(input), convertResolvedColumnToRexInputRef(wmCol, inputTableColumns), convertIntervalToRexIntervalLiteral((ResolvedLiteral) argumentList.get(2).getExpr()));
        case TVFStreamingUtils.SLIDING_WINDOW_TVF:
            // HOP tvf's second argument is descriptor.
            wmCol = extractWatermarkColumnFromDescriptor(argumentList.get(1).getDescriptorArg());
            return (RexCall) rexBuilder().makeCall(new SqlWindowTableFunction(SqlKind.HOP.name()), convertRelNodeToRexRangeRef(input), convertResolvedColumnToRexInputRef(wmCol, inputTableColumns), convertIntervalToRexIntervalLiteral((ResolvedLiteral) argumentList.get(2).getExpr()), convertIntervalToRexIntervalLiteral((ResolvedLiteral) argumentList.get(3).getExpr()));
        case TVFStreamingUtils.SESSION_WINDOW_TVF:
            // SESSION tvf's second argument is descriptor.
            wmCol = extractWatermarkColumnFromDescriptor(argumentList.get(1).getDescriptorArg());
            // SESSION tvf's third argument is descriptor.
            List<ResolvedColumn> keyCol = extractSessionKeyColumnFromDescriptor(argumentList.get(2).getDescriptorArg());
            List<RexNode> operands = new ArrayList<>();
            operands.add(convertRelNodeToRexRangeRef(input));
            operands.add(convertResolvedColumnToRexInputRef(wmCol, inputTableColumns));
            operands.add(convertIntervalToRexIntervalLiteral((ResolvedLiteral) argumentList.get(3).getExpr()));
            operands.addAll(convertResolvedColumnsToRexInputRef(keyCol, inputTableColumns));
            return (RexCall) rexBuilder().makeCall(new SqlWindowTableFunction(SqlKind.SESSION.name()), operands);
    }
    if (tvf instanceof FixedOutputSchemaTVF) {
        FixedOutputSchemaTVF fixedOutputSchemaTVF = (FixedOutputSchemaTVF) tvf;
        return (RexCall) rexBuilder().makeCall(new ZetaSqlUserDefinedSQLNativeTableValuedFunction(new SqlIdentifier(tvf.getName(), SqlParserPos.ZERO), opBinding -> {
            List<RelDataTypeField> relDataTypeFields = convertTVFRelationColumnsToRelDataTypeFields(fixedOutputSchemaTVF.getOutputSchema().getColumns());
            return new RelRecordType(relDataTypeFields);
        }, null, null, null, null));
    }
    throw new UnsupportedOperationException("Does not support table-valued function: " + tvf.getName());
}
Also used : FixedOutputSchemaTVF(com.google.zetasql.TableValuedFunction.FixedOutputSchemaTVF) ZetaSqlUserDefinedSQLNativeTableValuedFunction(org.apache.beam.sdk.extensions.sql.impl.ZetaSqlUserDefinedSQLNativeTableValuedFunction) ArrayList(java.util.ArrayList) SqlIdentifier(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlIdentifier) ResolvedColumn(com.google.zetasql.resolvedast.ResolvedColumn) RelRecordType(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelRecordType) RexCall(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexCall) RelDataTypeField(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataTypeField) ResolvedLiteral(com.google.zetasql.resolvedast.ResolvedNodes.ResolvedLiteral) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)

Example 5 with ResolvedColumn

use of com.google.zetasql.resolvedast.ResolvedColumn in project beam by apache.

the class JoinScanConverter method convert.

@Override
public RelNode convert(ResolvedJoinScan zetaNode, List<RelNode> inputs) {
    RelNode convertedLeftInput = inputs.get(0);
    RelNode convertedRightInput = inputs.get(1);
    List<ResolvedColumn> combinedZetaFieldsList = ImmutableList.<ResolvedColumn>builder().addAll(zetaNode.getLeftScan().getColumnList()).addAll(zetaNode.getRightScan().getColumnList()).build();
    List<RelDataTypeField> combinedCalciteFieldsList = ImmutableList.<RelDataTypeField>builder().addAll(convertedLeftInput.getRowType().getFieldList()).addAll(convertedRightInput.getRowType().getFieldList()).build();
    final RexNode condition;
    if (zetaNode.getJoinExpr() == null) {
        condition = getExpressionConverter().trueLiteral();
    } else {
        condition = getExpressionConverter().convertRexNodeFromResolvedExpr(zetaNode.getJoinExpr(), combinedZetaFieldsList, combinedCalciteFieldsList, ImmutableMap.of());
    }
    return LogicalJoin.create(convertedLeftInput, convertedRightInput, ImmutableList.of(), condition, ImmutableSet.of(), convertResolvedJoinType(zetaNode.getJoinType()));
}
Also used : RelDataTypeField(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataTypeField) RelNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelNode) ResolvedColumn(com.google.zetasql.resolvedast.ResolvedColumn) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)

Aggregations

ResolvedColumn (com.google.zetasql.resolvedast.ResolvedColumn)5 RexNode (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)4 ArrayList (java.util.ArrayList)3 RelDataTypeField (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataTypeField)2 FixedOutputSchemaTVF (com.google.zetasql.TableValuedFunction.FixedOutputSchemaTVF)1 ResolvedAggregateScan (com.google.zetasql.resolvedast.ResolvedNodes.ResolvedAggregateScan)1 ResolvedComputedColumn (com.google.zetasql.resolvedast.ResolvedNodes.ResolvedComputedColumn)1 ResolvedLiteral (com.google.zetasql.resolvedast.ResolvedNodes.ResolvedLiteral)1 ZetaSqlUserDefinedSQLNativeTableValuedFunction (org.apache.beam.sdk.extensions.sql.impl.ZetaSqlUserDefinedSQLNativeTableValuedFunction)1 RelNode (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelNode)1 RelRecordType (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelRecordType)1 RexBuilder (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexBuilder)1 RexCall (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexCall)1 SqlIdentifier (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlIdentifier)1