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();
}
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"));
}
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')"));
}
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;
}
Aggregations