Search in sources :

Example 1 with ResolvedAggregateScan

use of com.google.zetasql.resolvedast.ResolvedNodes.ResolvedAggregateScan 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)

Aggregations

ResolvedColumn (com.google.zetasql.resolvedast.ResolvedColumn)1 ResolvedAggregateScan (com.google.zetasql.resolvedast.ResolvedNodes.ResolvedAggregateScan)1 ResolvedComputedColumn (com.google.zetasql.resolvedast.ResolvedNodes.ResolvedComputedColumn)1 ArrayList (java.util.ArrayList)1 RexNode (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)1