Examples with BeamAggregationRel org.apache.beam.sdk.extensions.sql.impl.rel.BeamAggregationRel used on opensource projects

Example 1 with BeamAggregationRel

use of org.apache.beam.sdk.extensions.sql.impl.rel.BeamAggregationRel in project beam by apache.

the class BeamAggregateProjectMergeRuleTest method testBeamAggregateProjectMergeRule_withNoneTable.

@Test
public void testBeamAggregateProjectMergeRule_withNoneTable() {
    // When an IO does not supports project push-down, Projects should be merged with an aggregate.
    String sqlQuery = "select SUM(id) as id_sum from TEST_NONE group by name";
    BeamRelNode beamRel = sqlEnv.parseQuery(sqlQuery);
    BeamAggregationRel aggregate = (BeamAggregationRel) beamRel.getInput(0);
    BeamIOSourceRel ioSourceRel = (BeamIOSourceRel) aggregate.getInput();
    // Make sure project merged with an aggregate.
    assertThat(aggregate.getRowType().getFieldNames(), containsInAnyOrder("id_sum", "name"));
    // IO projects al fields.
    assertThat(ioSourceRel, instanceOf(BeamIOSourceRel.class));
    assertThat(ioSourceRel.getRowType().getFieldNames(), containsInAnyOrder("unused1", "name", "id", "unused2"));
}

Example 2 with BeamAggregationRel

use of org.apache.beam.sdk.extensions.sql.impl.rel.BeamAggregationRel in project beam by apache.

the class BeamBasicAggregationRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    Aggregate aggregate = call.rel(0);
    RelNode relNode = call.rel(1);
    if (aggregate.getGroupType() != Aggregate.Group.SIMPLE) {
        return;
    }
    if (relNode instanceof Project || relNode instanceof Calc || relNode instanceof Filter) {
        if (isWindowed(relNode) || hasWindowedParents(relNode)) {
            // This case is expected to get handled by the 'BeamAggregationRule'
            return;
        }
    }
    RelNode newTableScan = relNode.copy(relNode.getTraitSet(), relNode.getInputs());
    call.transformTo(new BeamAggregationRel(aggregate.getCluster(), aggregate.getTraitSet().replace(BeamLogicalConvention.INSTANCE), convert(newTableScan, newTableScan.getTraitSet().replace(BeamLogicalConvention.INSTANCE)), aggregate.getGroupSet(), aggregate.getGroupSets(), aggregate.getAggCallList(), null, -1));
}

Example 3 with BeamAggregationRel

use of org.apache.beam.sdk.extensions.sql.impl.rel.BeamAggregationRel in project beam by apache.

the class BeamAggregationRule method updateWindow.

private static RelNode updateWindow(Aggregate aggregate, Project project) {
    ImmutableBitSet groupByFields = aggregate.getGroupSet();
    ArrayList<RexNode> projects = new ArrayList(project.getProjects());
    WindowFn windowFn = null;
    int windowFieldIndex = -1;
    for (int groupFieldIndex : groupByFields.asList()) {
        RexNode projNode = projects.get(groupFieldIndex);
        if (!(projNode instanceof RexCall)) {
            continue;
        }
        RexCall rexCall = (RexCall) projNode;
        WindowFn fn = createWindowFn(rexCall.getOperands(), rexCall.op.kind);
        if (fn != null) {
            windowFn = fn;
            windowFieldIndex = groupFieldIndex;
            projects.set(groupFieldIndex, rexCall.getOperands().get(0));
        }
    }
    if (windowFn == null) {
        return null;
    }
    final Project newProject = project.copy(project.getTraitSet(), project.getInput(), projects, project.getRowType());
    return new BeamAggregationRel(aggregate.getCluster(), aggregate.getTraitSet().replace(BeamLogicalConvention.INSTANCE), convert(newProject, newProject.getTraitSet().replace(BeamLogicalConvention.INSTANCE)), aggregate.getGroupSet(), aggregate.getGroupSets(), aggregate.getAggCallList(), windowFn, windowFieldIndex);
}

Example 4 with BeamAggregationRel

use of org.apache.beam.sdk.extensions.sql.impl.rel.BeamAggregationRel in project beam by apache.

the class BeamAggregateProjectMergeRuleTest method testBeamAggregateProjectMergeRule_withProjectTable_withPredicate.

@Test
public void testBeamAggregateProjectMergeRule_withProjectTable_withPredicate() {
    // When an IO supports project push-down, Projects should be merged with an IO.
    String sqlQuery = "select SUM(id) as id_sum from TEST_PROJECT where unused1=1 group by name";
    BeamRelNode beamRel = sqlEnv.parseQuery(sqlQuery);
    BeamAggregationRel aggregate = (BeamAggregationRel) beamRel.getInput(0);
    BeamCalcRel calc = (BeamCalcRel) aggregate.getInput();
    BeamIOSourceRel ioSourceRel = (BeamIOSourceRel) calc.getInput();
    // Make sure project push-down took place.
    assertThat(ioSourceRel, instanceOf(BeamPushDownIOSourceRel.class));
    assertThat(ioSourceRel.getRowType().getFieldNames(), containsInAnyOrder("name", "id", "unused1"));
}

Example 5 with BeamAggregationRel

use of org.apache.beam.sdk.extensions.sql.impl.rel.BeamAggregationRel in project beam by apache.

the class BeamAggregateProjectMergeRuleTest method testBeamAggregateProjectMergeRule_withFilterTable.

@Test
public void testBeamAggregateProjectMergeRule_withFilterTable() {
    // When an IO does not supports project push-down, Projects should be merged with an aggregate.
    String sqlQuery = "select SUM(id) as id_sum from TEST_FILTER group by name";
    BeamRelNode beamRel = sqlEnv.parseQuery(sqlQuery);
    BeamAggregationRel aggregate = (BeamAggregationRel) beamRel.getInput(0);
    BeamIOSourceRel ioSourceRel = (BeamIOSourceRel) aggregate.getInput();
    // Make sure project merged with an aggregate.
    assertThat(aggregate.getRowType().getFieldNames(), containsInAnyOrder("id_sum", "name"));
    // IO projects al fields.
    assertThat(ioSourceRel, instanceOf(BeamIOSourceRel.class));
    assertThat(ioSourceRel.getRowType().getFieldNames(), containsInAnyOrder("unused1", "name", "id", "unused2"));
}