Examples with BeamRelMetadataQuery org.apache.beam.sdk.extensions.sql.impl.planner.BeamRelMetadataQuery used on opensource projects

Example 1 with BeamRelMetadataQuery

use of org.apache.beam.sdk.extensions.sql.impl.planner.BeamRelMetadataQuery in project beam by apache.

the class ZetaSQLQueryPlanner method convertToBeamRelInternal.

private BeamRelNode convertToBeamRelInternal(String sql, QueryParameters queryParams) {
    RelRoot root = plannerImpl.rel(sql, queryParams);
    RelTraitSet desiredTraits = root.rel.getTraitSet().replace(BeamLogicalConvention.INSTANCE).replace(root.collation).simplify();
    // beam physical plan
    root.rel.getCluster().setMetadataProvider(ChainedRelMetadataProvider.of(org.apache.beam.vendor.calcite.v1_28_0.com.google.common.collect.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();
    try {
        BeamRelNode beamRelNode = (BeamRelNode) plannerImpl.transform(0, desiredTraits, root.rel);
        LOG.info("BEAMPlan>\n" + RelOptUtil.toString(beamRelNode));
        return beamRelNode;
    } catch (RelOptPlanner.CannotPlanException e) {
        throw new SqlConversionException("Failed to produce plan for query " + sql, e);
    }
}

Example 2 with BeamRelMetadataQuery

use of org.apache.beam.sdk.extensions.sql.impl.planner.BeamRelMetadataQuery 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 3 with BeamRelMetadataQuery

use of org.apache.beam.sdk.extensions.sql.impl.planner.BeamRelMetadataQuery in project beam by apache.

the class BeamSideInputJoinRelTest method testNodeStatsEstimation.

@Test
public void testNodeStatsEstimation() {
    String sql = "SELECT o1.order_id, o1.sum_site_id, o2.buyer FROM " + "(select order_id, sum(site_id) as sum_site_id FROM ORDER_DETAILS " + "          GROUP BY order_id, TUMBLE(order_time, INTERVAL '1' HOUR)) o1 " + " JOIN " + " ORDER_DETAILS1 o2 " + " on " + " o1.order_id=o2.order_id";
    RelNode root = env.parseQuery(sql);
    while (!(root instanceof BeamSideInputJoinRel)) {
        root = root.getInput(0);
    }
    NodeStats estimate = BeamSqlRelUtils.getNodeStats(root, ((BeamRelMetadataQuery) root.getCluster().getMetadataQuery()));
    NodeStats leftEstimate = BeamSqlRelUtils.getNodeStats(((BeamSideInputJoinRel) root).getLeft(), ((BeamRelMetadataQuery) root.getCluster().getMetadataQuery()));
    NodeStats rightEstimate = BeamSqlRelUtils.getNodeStats(((BeamSideInputJoinRel) root).getRight(), ((BeamRelMetadataQuery) root.getCluster().getMetadataQuery()));
    Assert.assertFalse(estimate.isUnknown());
    Assert.assertEquals(0d, estimate.getRowCount(), 0.01);
    Assert.assertNotEquals(0d, estimate.getRate(), 0.001);
    Assert.assertTrue(estimate.getRate() < leftEstimate.getRowCount() * rightEstimate.getWindow() + rightEstimate.getRowCount() * leftEstimate.getWindow());
    Assert.assertNotEquals(0d, estimate.getWindow(), 0.001);
    Assert.assertTrue(estimate.getWindow() < leftEstimate.getWindow() * rightEstimate.getWindow());
}

Example 4 with BeamRelMetadataQuery

use of org.apache.beam.sdk.extensions.sql.impl.planner.BeamRelMetadataQuery in project beam by apache.

the class BeamUnionRelTest method testNodeStatsEstimation.

@Test
public void testNodeStatsEstimation() {
    String sql = "SELECT " + " order_id, site_id, price " + "FROM ORDER_DETAILS " + " UNION SELECT " + " order_id, site_id, price " + "FROM ORDER_DETAILS ";
    RelNode root = env.parseQuery(sql);
    while (!(root instanceof BeamUnionRel)) {
        root = root.getInput(0);
    }
    NodeStats estimate = BeamSqlRelUtils.getNodeStats(root, ((BeamRelMetadataQuery) root.getCluster().getMetadataQuery()));
    Assert.assertFalse(estimate.isUnknown());
    Assert.assertEquals(0d, estimate.getRate(), 0.01);
    Assert.assertEquals(2., estimate.getRowCount(), 0.01);
    Assert.assertEquals(2., estimate.getWindow(), 0.01);
}

Example 5 with BeamRelMetadataQuery

use of org.apache.beam.sdk.extensions.sql.impl.planner.BeamRelMetadataQuery in project beam by apache.

the class BeamCalcRelTest method testNodeStatsNumberOfConditions.

@Test
public void testNodeStatsNumberOfConditions() {
    String equalSql = "SELECT * FROM ORDER_DETAILS_BOUNDED where order_id=1";
    String doubleEqualSql = "SELECT * FROM ORDER_DETAILS_BOUNDED WHERE order_id=1 AND site_id=2 ";
    RelNode equalRoot = env.parseQuery(equalSql);
    RelNode doubleEqualRoot = env.parseQuery(doubleEqualSql);
    NodeStats equalEstimate = BeamSqlRelUtils.getNodeStats(equalRoot, ((BeamRelMetadataQuery) equalRoot.getCluster().getMetadataQuery()));
    NodeStats doubleEqualEstimate = BeamSqlRelUtils.getNodeStats(doubleEqualRoot, ((BeamRelMetadataQuery) doubleEqualRoot.getCluster().getMetadataQuery()));
    Assert.assertTrue(doubleEqualEstimate.getRowCount() < equalEstimate.getRowCount());
    Assert.assertTrue(doubleEqualEstimate.getWindow() < equalEstimate.getWindow());
}