Examples with Project org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Project used on opensource projects

Example 46 with Project

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Project in project beam by apache.

the class BeamZetaSqlUnnestRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    LogicalCorrelate correlate = call.rel(0);
    RelNode outer = call.rel(1);
    RelNode uncollect = call.rel(2);
    if (correlate.getRequiredColumns().cardinality() != 1) {
        // can only unnest a single column
        return;
    }
    if (correlate.getJoinType() != JoinRelType.INNER) {
        return;
    }
    if (!(uncollect instanceof ZetaSqlUnnest)) {
        // Drop projection
        uncollect = ((SingleRel) uncollect).getInput();
        if (uncollect instanceof RelSubset) {
            uncollect = ((RelSubset) uncollect).getOriginal();
        }
        if (!(uncollect instanceof ZetaSqlUnnest)) {
            return;
        }
    }
    RelNode project = ((ZetaSqlUnnest) uncollect).getInput();
    if (project instanceof RelSubset) {
        project = ((RelSubset) project).getOriginal();
    }
    if (!(project instanceof LogicalProject)) {
        return;
    }
    if (((LogicalProject) project).getProjects().size() != 1) {
        // can only unnest a single column
        return;
    }
    RexNode exp = ((LogicalProject) project).getProjects().get(0);
    if (!(exp instanceof RexFieldAccess)) {
        return;
    }
    RexFieldAccess fieldAccess = (RexFieldAccess) exp;
    // Innermost field index comes first (e.g. struct.field1.field2 => [2, 1])
    ImmutableList.Builder<Integer> fieldAccessIndices = ImmutableList.builder();
    while (true) {
        fieldAccessIndices.add(fieldAccess.getField().getIndex());
        if (!(fieldAccess.getReferenceExpr() instanceof RexFieldAccess)) {
            break;
        }
        fieldAccess = (RexFieldAccess) fieldAccess.getReferenceExpr();
    }
    call.transformTo(new BeamZetaSqlUnnestRel(correlate.getCluster(), correlate.getTraitSet().replace(BeamLogicalConvention.INSTANCE), convert(outer, outer.getTraitSet().replace(BeamLogicalConvention.INSTANCE)), call.rel(2).getRowType(), fieldAccessIndices.build()));
}

Example 47 with Project

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Project in project beam by apache.

the class TestTableProviderWithFilterPushDown method testIOSourceRel_withFilter_shouldProjectAllFields.

@Test
public void testIOSourceRel_withFilter_shouldProjectAllFields() {
    String selectTableStatement = "SELECT name FROM TEST where name='two'";
    BeamRelNode beamRelNode = sqlEnv.parseQuery(selectTableStatement);
    PCollection<Row> result = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
    assertThat(beamRelNode, instanceOf(BeamCalcRel.class));
    // Condition should be pushed-down to IO level
    assertNull(((Calc) beamRelNode).getProgram().getCondition());
    assertThat(beamRelNode.getInput(0), instanceOf(BeamIOSourceRel.class));
    List<String> projects = beamRelNode.getInput(0).getRowType().getFieldNames();
    // When performing standalone filter push-down IO should project all fields.
    assertThat(projects, containsInAnyOrder("unused1", "id", "name", "unused2", "b"));
    assertEquals(Schema.builder().addStringField("name").build(), result.getSchema());
    PAssert.that(result).containsInAnyOrder(row(result.getSchema(), "two"));
    pipeline.run().waitUntilFinish(Duration.standardMinutes(2));
}

Example 48 with Project

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Project in project beam by apache.

the class TestTableProviderWithFilterPushDown method testIOSourceRel_withUnsupportedFilter_calcPreservesCondition.

@Test
public void testIOSourceRel_withUnsupportedFilter_calcPreservesCondition() {
    String selectTableStatement = "SELECT name FROM TEST where id+1=2";
    BeamRelNode beamRelNode = sqlEnv.parseQuery(selectTableStatement);
    PCollection<Row> result = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
    assertThat(beamRelNode, instanceOf(BeamCalcRel.class));
    // Unsupported condition should be preserved in a Calc
    assertNotNull(((Calc) beamRelNode).getProgram().getCondition());
    assertThat(beamRelNode.getInput(0), instanceOf(BeamIOSourceRel.class));
    List<String> projects = beamRelNode.getInput(0).getRowType().getFieldNames();
    // When performing standalone filter push-down IO should project all fields.
    assertThat(projects, containsInAnyOrder("unused1", "id", "name", "unused2", "b"));
    assertEquals(Schema.builder().addStringField("name").build(), result.getSchema());
    PAssert.that(result).containsInAnyOrder(row(result.getSchema(), "one"));
    pipeline.run().waitUntilFinish(Duration.standardMinutes(2));
}

Example 49 with Project

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Project in project beam by apache.

the class TestTableProviderWithFilterPushDown method testIOSourceRel_withSupportedFilter_selectInRandomOrder.

@Test
public void testIOSourceRel_withSupportedFilter_selectInRandomOrder() {
    String selectTableStatement = "SELECT unused2, id, name FROM TEST where b";
    BeamRelNode beamRelNode = sqlEnv.parseQuery(selectTableStatement);
    PCollection<Row> result = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
    assertThat(beamRelNode, instanceOf(BeamCalcRel.class));
    // Condition should be pushed-down to IO level
    assertNull(((Calc) beamRelNode).getProgram().getCondition());
    assertThat(beamRelNode.getInput(0), instanceOf(BeamIOSourceRel.class));
    List<String> projects = beamRelNode.getInput(0).getRowType().getFieldNames();
    // When performing standalone filter push-down IO should project all fields.
    assertThat(projects, containsInAnyOrder("unused1", "id", "name", "unused2", "b"));
    assertEquals(Schema.builder().addInt16Field("unused2").addInt32Field("id").addStringField("name").build(), result.getSchema());
    PAssert.that(result).containsInAnyOrder(row(result.getSchema(), (short) 100, 1, "one"));
    pipeline.run().waitUntilFinish(Duration.standardMinutes(2));
}

Example 50 with Project

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Project in project beam by apache.

the class BeamAggregateProjectMergeRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    final Project project = call.rel(1);
    BeamIOSourceRel io = getUnderlyingIO(new HashSet<>(), project);
    // supported.
    if (io == null || !io.getBeamSqlTable().supportsProjects().isSupported()) {
        super.onMatch(call);
    }
}