Search in sources :

Example 21 with Calc

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Calc 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();
}
Also used : RuleSet(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.tools.RuleSet) ProjectCalcMergeRule(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.rules.ProjectCalcMergeRule) ImmutableList(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList) BeamUnnestRule(org.apache.beam.sdk.extensions.sql.impl.rule.BeamUnnestRule) BeamUncollectRule(org.apache.beam.sdk.extensions.sql.impl.rule.BeamUncollectRule) FilterCalcMergeRule(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.rules.FilterCalcMergeRule) JoinCommuteRule(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.rules.JoinCommuteRule) BeamCalcRule(org.apache.beam.sdk.extensions.sql.impl.rule.BeamCalcRule) RelOptRule(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.plan.RelOptRule)

Example 22 with Calc

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

the class IOPushDownRuleTest method testFindUtilisedInputRefs.

@Test
public void testFindUtilisedInputRefs() {
    String sqlQuery = "select id+10 from TEST where name='one'";
    BeamRelNode basicRel = sqlEnv.parseQuery(sqlQuery);
    assertThat(basicRel, instanceOf(Calc.class));
    Calc calc = (Calc) basicRel;
    final Pair<ImmutableList<RexNode>, ImmutableList<RexNode>> projectFilter = calc.getProgram().split();
    final ImmutableList<RexNode> projects = projectFilter.left;
    final ImmutableList<RexNode> filters = projectFilter.right;
    Set<String> usedFields = new HashSet<>();
    BeamIOPushDownRule.INSTANCE.findUtilizedInputRefs(calc.getProgram().getInputRowType(), projects.get(0), usedFields);
    assertThat(usedFields, containsInAnyOrder("id"));
    BeamIOPushDownRule.INSTANCE.findUtilizedInputRefs(calc.getProgram().getInputRowType(), filters.get(0), usedFields);
    assertThat(usedFields, containsInAnyOrder("id", "name"));
}
Also used : BeamRelNode(org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode) ImmutableList(org.apache.beam.vendor.calcite.v1_28_0.com.google.common.collect.ImmutableList) Calc(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Calc) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode) HashSet(java.util.HashSet) Test(org.junit.Test)

Example 23 with Calc

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

the class IOPushDownRuleTest method testReMapRexNodeToNewInputs.

@Test
public void testReMapRexNodeToNewInputs() {
    String sqlQuery = "select id+10 from TEST where name='one'";
    BeamRelNode basicRel = sqlEnv.parseQuery(sqlQuery);
    assertThat(basicRel, instanceOf(Calc.class));
    Calc calc = (Calc) basicRel;
    final Pair<ImmutableList<RexNode>, ImmutableList<RexNode>> projectFilter = calc.getProgram().split();
    final ImmutableList<RexNode> projects = projectFilter.left;
    final ImmutableList<RexNode> filters = projectFilter.right;
    List<Integer> mapping = ImmutableList.of(1, 2);
    RexNode newProject = BeamIOPushDownRule.INSTANCE.reMapRexNodeToNewInputs(projects.get(0), mapping);
    assertThat(newProject.toString(), equalTo("+($0, 10)"));
    RexNode newFilter = BeamIOPushDownRule.INSTANCE.reMapRexNodeToNewInputs(filters.get(0), mapping);
    assertThat(newFilter.toString(), equalTo("=($1, 'one')"));
}
Also used : BeamRelNode(org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode) ImmutableList(org.apache.beam.vendor.calcite.v1_28_0.com.google.common.collect.ImmutableList) Calc(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Calc) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode) Test(org.junit.Test)

Example 24 with Calc

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

the class CalcRelSplitter method createProgramForLevel.

/**
 * Creates a program containing the expressions for a given level.
 *
 * <p>The expression list of the program will consist of all entries in the expression list <code>
 * allExprs[i]</code> for which the corresponding level ordinal <code>exprLevels[i]</code> is
 * equal to <code>level</code>. Expressions are mapped according to <code>inputExprOrdinals</code>
 * .
 *
 * @param level Level ordinal
 * @param levelCount Number of levels
 * @param inputRowType Input row type
 * @param allExprs Array of all expressions
 * @param exprLevels Array of the level ordinal of each expression
 * @param inputExprOrdinals Ordinals in the expression list of input expressions. Input expression
 *     <code>i</code> will be found at position <code>inputExprOrdinals[i]</code>.
 * @param projectExprOrdinals Ordinals of the expressions to be output this level.
 * @param conditionExprOrdinal Ordinal of the expression to form the condition for this level, or
 *     -1 if there is no condition.
 * @param outputRowType Output row type
 * @return Relational expression
 */
private RexProgram createProgramForLevel(int level, int levelCount, RelDataType inputRowType, RexNode[] allExprs, int[] exprLevels, int[] inputExprOrdinals, final int[] projectExprOrdinals, int conditionExprOrdinal, @Nullable RelDataType outputRowType) {
    // Build a list of expressions to form the calc.
    List<RexNode> exprs = new ArrayList<>();
    // exprInverseOrdinals describes where an expression in allExprs comes
    // from -- from an input, from a calculated expression, or -1 if not
    // available at this level.
    int[] exprInverseOrdinals = new int[allExprs.length];
    Arrays.fill(exprInverseOrdinals, -1);
    int j = 0;
    // and are used here.
    for (int i = 0; i < inputExprOrdinals.length; i++) {
        final int inputExprOrdinal = inputExprOrdinals[i];
        exprs.add(new RexInputRef(i, allExprs[inputExprOrdinal].getType()));
        exprInverseOrdinals[inputExprOrdinal] = j;
        ++j;
    }
    // Next populate the computed expressions.
    final RexShuttle shuttle = new InputToCommonExprConverter(exprInverseOrdinals, exprLevels, level, inputExprOrdinals, allExprs);
    for (int i = 0; i < allExprs.length; i++) {
        if (exprLevels[i] == level || exprLevels[i] == -1 && level == (levelCount - 1) && allExprs[i] instanceof RexLiteral) {
            RexNode expr = allExprs[i];
            final RexNode translatedExpr = expr.accept(shuttle);
            exprs.add(translatedExpr);
            assert exprInverseOrdinals[i] == -1;
            exprInverseOrdinals[i] = j;
            ++j;
        }
    }
    // Form the projection and condition list. Project and condition
    // ordinals are offsets into allExprs, so we need to map them into
    // exprs.
    final List<RexLocalRef> projectRefs = new ArrayList<>(projectExprOrdinals.length);
    final List<String> fieldNames = new ArrayList<>(projectExprOrdinals.length);
    for (int i = 0; i < projectExprOrdinals.length; i++) {
        final int projectExprOrdinal = projectExprOrdinals[i];
        final int index = exprInverseOrdinals[projectExprOrdinal];
        assert index >= 0;
        RexNode expr = allExprs[projectExprOrdinal];
        projectRefs.add(new RexLocalRef(index, expr.getType()));
        // Inherit meaningful field name if possible.
        fieldNames.add(deriveFieldName(expr, i));
    }
    RexLocalRef conditionRef;
    if (conditionExprOrdinal >= 0) {
        final int index = exprInverseOrdinals[conditionExprOrdinal];
        conditionRef = new RexLocalRef(index, allExprs[conditionExprOrdinal].getType());
    } else {
        conditionRef = null;
    }
    if (outputRowType == null) {
        outputRowType = RexUtil.createStructType(typeFactory, projectRefs, fieldNames, null);
    }
    final RexProgram program = new RexProgram(inputRowType, exprs, projectRefs, conditionRef, outputRowType);
    // call operands), since literals should be inlined.
    return program;
}
Also used : RexLiteral(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexLiteral) RexShuttle(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexShuttle) RexProgram(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexProgram) ArrayList(java.util.ArrayList) RexInputRef(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexInputRef) RexLocalRef(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexLocalRef) RexNode(org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)

Aggregations

Calc (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Calc)10 RexNode (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexNode)8 BeamRelNode (org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode)6 Calc (org.apache.calcite.rel.core.Calc)6 Test (org.junit.Test)6 BeamIOSourceRel (org.apache.beam.sdk.extensions.sql.impl.rel.BeamIOSourceRel)5 RelNode (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelNode)5 BeamCalcRel (org.apache.beam.sdk.extensions.sql.impl.rel.BeamCalcRel)4 Row (org.apache.beam.sdk.values.Row)3 ImmutableList (org.apache.beam.vendor.calcite.v1_28_0.com.google.common.collect.ImmutableList)3 LogicalCalc (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.logical.LogicalCalc)3 RelDataType (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataType)3 RexLocalRef (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexLocalRef)3 LogicalCalc (org.apache.calcite.rel.logical.LogicalCalc)3 ArrayList (java.util.ArrayList)2 HashSet (java.util.HashSet)2 LinkedHashSet (java.util.LinkedHashSet)2 RelDataTypeField (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.type.RelDataTypeField)2 RexInputRef (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexInputRef)2 RexLiteral (org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexLiteral)2