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

Example 6 with BeamRelNode

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

the class CalciteQueryPlanner method convertToBeamRel.

/**
 * It parses and validate the input query, then convert into a {@link BeamRelNode} tree. Note that
 * query parameters are not yet supported.
 */
@Override
public BeamRelNode convertToBeamRel(String sqlStatement, QueryParameters queryParameters) throws ParseException, SqlConversionException {
    Preconditions.checkArgument(queryParameters.getKind() == Kind.NONE, "Beam SQL Calcite dialect does not yet support query parameters.");
    BeamRelNode beamRelNode;
    try {
        SqlNode parsed = planner.parse(sqlStatement);
        TableResolutionUtils.setupCustomTableResolution(connection, parsed);
        SqlNode validated = planner.validate(parsed);
        LOG.info("SQL:\n" + validated);
        // root of original logical plan
        RelRoot root = planner.rel(validated);
        LOG.info("SQLPlan>\n" + RelOptUtil.toString(root.rel));
        RelTraitSet desiredTraits = root.rel.getTraitSet().replace(BeamLogicalConvention.INSTANCE).replace(root.collation).simplify();
        // beam physical plan
        root.rel.getCluster().setMetadataProvider(ChainedRelMetadataProvider.of(ImmutableList.of(NonCumulativeCostImpl.SOURCE, RelMdNodeStats.SOURCE, root.rel.getCluster().getMetadataProvider())));
        root.rel.getCluster().setMetadataQuerySupplier(BeamRelMetadataQuery::instance);
        RelMetadataQuery.THREAD_PROVIDERS.set(JaninoRelMetadataProvider.of(root.rel.getCluster().getMetadataProvider()));
        root.rel.getCluster().invalidateMetadataQuery();
        beamRelNode = (BeamRelNode) planner.transform(0, desiredTraits, root.rel);
        LOG.info("BEAMPlan>\n" + RelOptUtil.toString(beamRelNode));
    } catch (RelConversionException | CannotPlanException e) {
        throw new SqlConversionException(String.format("Unable to convert query %s", sqlStatement), e);
    } catch (SqlParseException | ValidationException e) {
        throw new ParseException(String.format("Unable to parse query %s", sqlStatement), e);
    } finally {
        planner.close();
    }
    return beamRelNode;
}

Example 7 with BeamRelNode

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

the class BigQueryIOPushDownIT method readUsingDefaultMethod.

@Test
public void readUsingDefaultMethod() {
    sqlEnv.executeDdl(String.format(CREATE_TABLE_STATEMENT, Method.DEFAULT.toString()));
    BeamRelNode beamRelNode = sqlEnv.parseQuery(SELECT_STATEMENT);
    BeamSqlRelUtils.toPCollection(pipeline, beamRelNode).apply(ParDo.of(new TimeMonitor<>(NAMESPACE, READ_TIME_METRIC)));
    PipelineResult result = pipeline.run();
    result.waitUntilFinish();
    collectAndPublishMetrics(result, "_default");
}

Example 8 with BeamRelNode

use of org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode 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 9 with BeamRelNode

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

the class ThreeTablesSchema method testSystemNotReorderingMediumSmallLarge.

@Test
public void testSystemNotReorderingMediumSmallLarge() {
    TestTableProvider tableProvider = new TestTableProvider();
    createThreeTables(tableProvider);
    BeamSqlEnv env = BeamSqlEnv.withTableProvider(tableProvider);
    // This is a correct ordered join because large table is on the top. It should not change that.
    BeamRelNode parsedQuery = env.parseQuery("select * from medium_table " + " JOIN small_table on medium_table.small_key = small_table.medium_key " + " JOIN large_table on large_table.medium_key = medium_table.large_key ");
    assertTopTableInJoins(parsedQuery, "large_table");
}

Example 10 with BeamRelNode

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

the class CalciteQueryPlannerTest method testclusterCostHandlerUsesBeamCost.

@Test
public void testclusterCostHandlerUsesBeamCost() {
    String sql = "select * from medium_table";
    BeamRelNode root = env.parseQuery(sql);
    Assert.assertTrue(root.getCluster().getPlanner().getCost(root, root.getCluster().getMetadataQuery()) instanceof BeamCostModel);
}