Search in sources :

Example 31 with AND

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlKind.AND in project beam by apache.

the class AggregateScanConverter method convertAggregateScanInputScanToLogicalProject.

private LogicalProject 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<>();
    // LogicalAggregate.
    for (ResolvedComputedColumn computedColumn : node.getGroupByList()) {
        projects.add(getExpressionConverter().convertRexNodeFromResolvedExpr(computedColumn.getExpr(), node.getInputScan().getColumnList(), input.getRowType().getFieldList(), ImmutableMap.of()));
        fieldNames.add(getTrait().resolveAlias(computedColumn.getColumn()));
    }
    // 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);
            for (int i = 0; i < aggregateFunctionCall.getArgumentList().size(); i++) {
                if (i == 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(), ImmutableMap.of()));
                } else {
                    projects.add(getExpressionConverter().convertRexNodeFromResolvedExpr(aggregateFunctionCall.getArgumentList().get(i)));
                }
                fieldNames.add(getTrait().resolveAlias(resolvedComputedColumn.getColumn()));
            }
        }
    }
    return LogicalProject.create(input, ImmutableList.of(), projects, fieldNames);
}
Also used : ResolvedAggregateFunctionCall(com.google.zetasql.resolvedast.ResolvedNodes.ResolvedAggregateFunctionCall) ArrayList(java.util.ArrayList) ResolvedExpr(com.google.zetasql.resolvedast.ResolvedNodes.ResolvedExpr) ResolvedComputedColumn(com.google.zetasql.resolvedast.ResolvedNodes.ResolvedComputedColumn) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)

Example 32 with AND

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlKind.AND in project beam by apache.

the class LimitOffsetScanToLimitConverter method convert.

@Override
public RelNode convert(ResolvedLimitOffsetScan zetaNode, List<RelNode> inputs) {
    RelNode input = inputs.get(0);
    RelCollation relCollation = RelCollations.of(ImmutableList.of());
    RexNode offset = zetaNode.getOffset() == null ? null : getExpressionConverter().convertRexNodeFromResolvedExpr(zetaNode.getOffset());
    RexNode fetch = getExpressionConverter().convertRexNodeFromResolvedExpr(zetaNode.getLimit(), zetaNode.getColumnList(), input.getRowType().getFieldList(), ImmutableMap.of());
    // offset or fetch being RexDynamicParam means it is NULL (the only param supported currently)
    if (offset instanceof RexDynamicParam || RexLiteral.isNullLiteral(offset) || fetch instanceof RexDynamicParam || RexLiteral.isNullLiteral(fetch)) {
        throw new UnsupportedOperationException("Limit requires non-null count and offset");
    }
    return LogicalSort.create(input, relCollation, offset, fetch);
}
Also used : RelCollation(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelCollation) RelNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelNode) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode) RexDynamicParam(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexDynamicParam)

Example 33 with AND

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlKind.AND in project beam by apache.

the class TableResolutionTest method testResolveNestedWithDots.

/**
 * Unit test for resolving a table with dots in the subschema names and the table name.
 */
@Test
public void testResolveNestedWithDots() {
    String subSchema = "fake.schema";
    String tableName = "fake.table";
    when(mockSchemaPlus.getSubSchema(subSchema)).thenReturn(innerSchemaPlus);
    when(innerSchemaPlus.getTable(tableName)).thenReturn(mockTable);
    Table table = TableResolution.resolveCalciteTable(mockSchemaPlus, ImmutableList.of(subSchema, tableName));
    assertThat(table, Matchers.is(mockTable));
}
Also used : Table(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.Table) Test(org.junit.Test)

Example 34 with AND

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlKind.AND in project beam by apache.

the class TableResolutionTest method testMissingTableInSubschema.

/**
 * Unit test for resolving a table with some hierarchy and the table is missing.
 */
@Test
public void testMissingTableInSubschema() {
    String subSchema = "fake_schema";
    String tableName = "fake_table";
    when(mockSchemaPlus.getSubSchema(subSchema)).thenReturn(innerSchemaPlus);
    when(innerSchemaPlus.getTable(tableName)).thenReturn(null);
    Table table = TableResolution.resolveCalciteTable(mockSchemaPlus, ImmutableList.of(subSchema, tableName));
    assertThat(table, Matchers.nullValue());
}
Also used : Table(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.schema.Table) Test(org.junit.Test)

Example 35 with AND

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlKind.AND in project beam by apache.

the class LimitOffsetScanToOrderByLimitConverter method convert.

@Override
public RelNode convert(ResolvedLimitOffsetScan zetaNode, List<RelNode> inputs) {
    ResolvedOrderByScan inputOrderByScan = (ResolvedOrderByScan) zetaNode.getInputScan();
    RelNode input = inputs.get(0);
    RelCollation relCollation = getRelCollation(inputOrderByScan);
    RexNode offset = zetaNode.getOffset() == null ? null : getExpressionConverter().convertRexNodeFromResolvedExpr(zetaNode.getOffset());
    RexNode fetch = getExpressionConverter().convertRexNodeFromResolvedExpr(zetaNode.getLimit(), zetaNode.getColumnList(), input.getRowType().getFieldList(), ImmutableMap.of());
    if (RexLiteral.isNullLiteral(offset) || RexLiteral.isNullLiteral(fetch)) {
        throw new UnsupportedOperationException("Limit requires non-null count and offset");
    }
    RelNode sorted = LogicalSort.create(input, relCollation, offset, fetch);
    return convertOrderByScanToLogicalScan(inputOrderByScan, sorted);
}
Also used : RelCollation(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelCollation) RelNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelNode) ResolvedOrderByScan(com.google.zetasql.resolvedast.ResolvedNodes.ResolvedOrderByScan) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)

Aggregations

RexNode (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)16 ArrayList (java.util.ArrayList)11 Test (org.junit.Test)10 RelNode (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelNode)9 List (java.util.List)8 Collectors (java.util.stream.Collectors)8 Serializable (java.io.Serializable)7 Arrays (java.util.Arrays)7 RelDataType (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataType)7 IOException (java.io.IOException)6 Statement (java.sql.Statement)6 HashMap (java.util.HashMap)6 Map (java.util.Map)6 Schema (org.apache.beam.sdk.schemas.Schema)6 PCollection (org.apache.beam.sdk.values.PCollection)6 JsonProcessingException (com.fasterxml.jackson.core.JsonProcessingException)5 ObjectMapper (com.fasterxml.jackson.databind.ObjectMapper)5 ByteArrayOutputStream (java.io.ByteArrayOutputStream)5 StandardCharsets (java.nio.charset.StandardCharsets)5 UTF_8 (java.nio.charset.StandardCharsets.UTF_8)5