Examples with RuleSet org.apache.calcite.tools.RuleSet used on opensource projects

Example 6 with RuleSet

use of org.apache.calcite.tools.RuleSet in project beam by apache.

the class BigQueryIOPushDownIT method readUsingDirectReadMethod.

@Test
public void readUsingDirectReadMethod() {
    List<RelOptRule> ruleList = new ArrayList<>();
    for (RuleSet x : getRuleSets()) {
        x.iterator().forEachRemaining(ruleList::add);
    }
    // Remove push-down rule
    ruleList.remove(BeamIOPushDownRule.INSTANCE);
    InMemoryMetaStore inMemoryMetaStore = new InMemoryMetaStore();
    inMemoryMetaStore.registerProvider(new BigQueryPerfTableProvider(NAMESPACE, FIELDS_READ_METRIC));
    sqlEnv = BeamSqlEnv.builder(inMemoryMetaStore).setPipelineOptions(PipelineOptionsFactory.create()).setRuleSets(ImmutableList.of(RuleSets.ofList(ruleList))).build();
    sqlEnv.executeDdl(String.format(CREATE_TABLE_STATEMENT, Method.DIRECT_READ.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, "_directread");
}

Example 7 with RuleSet

use of org.apache.calcite.tools.RuleSet in project beam by apache.

the class CalciteQueryPlanner method defaultConfig.

public FrameworkConfig defaultConfig(JdbcConnection connection, Collection<RuleSet> ruleSets) {
    final CalciteConnectionConfig config = connection.config();
    final SqlParser.ConfigBuilder parserConfig = SqlParser.configBuilder().setQuotedCasing(config.quotedCasing()).setUnquotedCasing(config.unquotedCasing()).setQuoting(config.quoting()).setConformance(config.conformance()).setCaseSensitive(config.caseSensitive());
    final SqlParserImplFactory parserFactory = config.parserFactory(SqlParserImplFactory.class, null);
    if (parserFactory != null) {
        parserConfig.setParserFactory(parserFactory);
    }
    final SchemaPlus schema = connection.getRootSchema();
    final SchemaPlus defaultSchema = connection.getCurrentSchemaPlus();
    final ImmutableList<RelTraitDef> traitDefs = ImmutableList.of(ConventionTraitDef.INSTANCE);
    final CalciteCatalogReader catalogReader = new CalciteCatalogReader(CalciteSchema.from(schema), ImmutableList.of(defaultSchema.getName()), connection.getTypeFactory(), connection.config());
    final SqlOperatorTable opTab0 = connection.config().fun(SqlOperatorTable.class, SqlStdOperatorTable.instance());
    return Frameworks.newConfigBuilder().parserConfig(parserConfig.build()).defaultSchema(defaultSchema).traitDefs(traitDefs).context(Contexts.of(connection.config())).ruleSets(ruleSets.toArray(new RuleSet[0])).costFactory(BeamCostModel.FACTORY).typeSystem(connection.getTypeFactory().getTypeSystem()).operatorTable(SqlOperatorTables.chain(opTab0, catalogReader)).build();
}

Example 8 with RuleSet

use of org.apache.calcite.tools.RuleSet in project beam by apache.

the class ZetaSQLPushDownTest method initializeCalciteEnvironmentWithContext.

private static void initializeCalciteEnvironmentWithContext(Context... extraContext) {
    JdbcConnection jdbcConnection = JdbcDriver.connect(tableProvider, PipelineOptionsFactory.create());
    SchemaPlus defaultSchemaPlus = jdbcConnection.getCurrentSchemaPlus();
    final ImmutableList<RelTraitDef> traitDefs = ImmutableList.of(ConventionTraitDef.INSTANCE);
    Object[] contexts = ImmutableList.<Context>builder().add(Contexts.of(jdbcConnection.config())).add(extraContext).build().toArray();
    config = Frameworks.newConfigBuilder().defaultSchema(defaultSchemaPlus).traitDefs(traitDefs).context(Contexts.of(contexts)).ruleSets(ZetaSQLQueryPlanner.getZetaSqlRuleSets().toArray(new RuleSet[0])).costFactory(BeamCostModel.FACTORY).typeSystem(jdbcConnection.getTypeFactory().getTypeSystem()).build();
}

Example 9 with RuleSet

use of org.apache.calcite.tools.RuleSet in project beam by apache.

the class ZetaSQLQueryPlanner method modifyRuleSetsForZetaSql.

private static Collection<RuleSet> modifyRuleSetsForZetaSql(Collection<RuleSet> ruleSets, Collection<RelOptRule> calc) {
    ImmutableList.Builder<RuleSet> ret = ImmutableList.builder();
    for (RuleSet ruleSet : ruleSets) {
        ImmutableList.Builder<RelOptRule> bd = ImmutableList.builder();
        for (RelOptRule rule : ruleSet) {
            // requires the JoinCommuteRule, which doesn't work without struct flattening.
            if (rule instanceof JoinCommuteRule) {
                continue;
            } else if (rule instanceof FilterCalcMergeRule || rule instanceof ProjectCalcMergeRule) {
                // planning result eventually.
                continue;
            } else if (rule instanceof BeamCalcRule) {
                bd.addAll(calc);
            } else if (rule instanceof BeamUnnestRule) {
                bd.add(BeamZetaSqlUnnestRule.INSTANCE);
            } else if (rule instanceof BeamUncollectRule) {
                bd.add(BeamZetaSqlUncollectRule.INSTANCE);
            } else {
                bd.add(rule);
            }
        }
        ret.add(RuleSets.ofList(bd.build()));
    }
    return ret.build();
}

Example 10 with RuleSet

use of org.apache.calcite.tools.RuleSet in project beam by apache.

the class ThreeTablesSchema method testBeamJoinPushThroughJoinRuleLeft.

@Test
public void testBeamJoinPushThroughJoinRuleLeft() throws Exception {
    RuleSet prepareRules = RuleSets.ofList(CoreRules.SORT_PROJECT_TRANSPOSE, EnumerableRules.ENUMERABLE_JOIN_RULE, EnumerableRules.ENUMERABLE_PROJECT_RULE, EnumerableRules.ENUMERABLE_SORT_RULE, EnumerableRules.ENUMERABLE_TABLE_SCAN_RULE);
    String sqlQuery = "select * from \"tt\".\"large_table\" as large_table " + " JOIN \"tt\".\"medium_table\" as medium_table on large_table.\"medium_key\" = medium_table.\"large_key\" " + " JOIN \"tt\".\"small_table\" as small_table on medium_table.\"small_key\" = small_table.\"medium_key\" ";
    RelNode originalPlan = transform(sqlQuery, prepareRules);
    RelNode optimizedPlan = transform(sqlQuery, RuleSets.ofList(ImmutableList.<RelOptRule>builder().addAll(prepareRules).add(BeamJoinPushThroughJoinRule.LEFT).build()));
    assertTopTableInJoins(originalPlan, "small_table");
    assertTopTableInJoins(optimizedPlan, "large_table");
}