Search in sources :

Example 11 with RexCorrelVariable

use of org.apache.calcite.rex.RexCorrelVariable in project calcite by apache.

the class RelFieldTrimmer method result.

protected TrimResult result(RelNode r, final Mapping mapping) {
    final RexBuilder rexBuilder = relBuilder.getRexBuilder();
    for (final CorrelationId correlation : r.getVariablesSet()) {
        r = r.accept(new CorrelationReferenceFinder() {

            protected RexNode handle(RexFieldAccess fieldAccess) {
                final RexCorrelVariable v = (RexCorrelVariable) fieldAccess.getReferenceExpr();
                if (v.id.equals(correlation) && v.getType().getFieldCount() == mapping.getSourceCount()) {
                    final int old = fieldAccess.getField().getIndex();
                    final int new_ = mapping.getTarget(old);
                    final RelDataTypeFactory.Builder typeBuilder = relBuilder.getTypeFactory().builder();
                    for (int target : Util.range(mapping.getTargetCount())) {
                        typeBuilder.add(v.getType().getFieldList().get(mapping.getSource(target)));
                    }
                    final RexNode newV = rexBuilder.makeCorrel(typeBuilder.build(), v.id);
                    if (old != new_) {
                        return rexBuilder.makeFieldAccess(newV, new_);
                    }
                }
                return fieldAccess;
            }
        });
    }
    return new TrimResult(r, mapping);
}
Also used : RexCorrelVariable(org.apache.calcite.rex.RexCorrelVariable) RelDataTypeFactory(org.apache.calcite.rel.type.RelDataTypeFactory) RexBuilder(org.apache.calcite.rex.RexBuilder) CorrelationId(org.apache.calcite.rel.core.CorrelationId) RexFieldAccess(org.apache.calcite.rex.RexFieldAccess) RexNode(org.apache.calcite.rex.RexNode)

Example 12 with RexCorrelVariable

use of org.apache.calcite.rex.RexCorrelVariable in project drill by apache.

the class LateralUnnestRowIDVisitor method visitLateral.

@Override
public Prel visitLateral(LateralJoinPrel prel, Boolean isRightOfLateral) throws RuntimeException {
    List<RelNode> children = Lists.newArrayList();
    children.add(((Prel) prel.getInput(0)).accept(this, isRightOfLateral));
    children.add(((Prel) prel.getInput(1)).accept(this, true));
    if (!isRightOfLateral) {
        return (Prel) prel.copy(prel.getTraitSet(), children);
    } else {
        // Adjust the column numbering due to an additional column "$drill_implicit_field$" is added to the inputs.
        Map<Integer, Integer> requiredColsMap = new HashMap<>();
        for (Integer corrColIndex : prel.getRequiredColumns()) {
            requiredColsMap.put(corrColIndex, corrColIndex + 1);
        }
        ImmutableBitSet requiredColumns = prel.getRequiredColumns().shift(1);
        CorrelationId corrId = prel.getCluster().createCorrel();
        RexCorrelVariable updatedCorrel = (RexCorrelVariable) prel.getCluster().getRexBuilder().makeCorrel(children.get(0).getRowType(), corrId);
        RelNode rightChild = children.get(1).accept(new CorrelateVarReplacer(new ProjectCorrelateTransposeRule.RexFieldAccessReplacer(prel.getCorrelationId(), updatedCorrel, prel.getCluster().getRexBuilder(), requiredColsMap)));
        return (Prel) prel.copy(prel.getTraitSet(), children.get(0), rightChild, corrId, requiredColumns, prel.getJoinType());
    }
}
Also used : RexCorrelVariable(org.apache.calcite.rex.RexCorrelVariable) RelNode(org.apache.calcite.rel.RelNode) ImmutableBitSet(org.apache.calcite.util.ImmutableBitSet) HashMap(java.util.HashMap) CorrelationId(org.apache.calcite.rel.core.CorrelationId) LateralJoinPrel(org.apache.drill.exec.planner.physical.LateralJoinPrel) Prel(org.apache.drill.exec.planner.physical.Prel) UnnestPrel(org.apache.drill.exec.planner.physical.UnnestPrel)

Aggregations

RexCorrelVariable (org.apache.calcite.rex.RexCorrelVariable)12 RexFieldAccess (org.apache.calcite.rex.RexFieldAccess)9 RexNode (org.apache.calcite.rex.RexNode)9 CorrelationId (org.apache.calcite.rel.core.CorrelationId)5 RexCall (org.apache.calcite.rex.RexCall)4 RexInputRef (org.apache.calcite.rex.RexInputRef)4 RelNode (org.apache.calcite.rel.RelNode)3 RexBuilder (org.apache.calcite.rex.RexBuilder)3 RelBuilder (org.apache.calcite.tools.RelBuilder)3 ImmutableBitSet (org.apache.calcite.util.ImmutableBitSet)3 ArrayList (java.util.ArrayList)2 HashMap (java.util.HashMap)2 List (java.util.List)2 Map (java.util.Map)2 RelCollation (org.apache.calcite.rel.RelCollation)2 RelFieldCollation (org.apache.calcite.rel.RelFieldCollation)2 RelMetadataQuery (org.apache.calcite.rel.metadata.RelMetadataQuery)2 RelDataTypeFactory (org.apache.calcite.rel.type.RelDataTypeFactory)2 CorrelationReferenceFinder (org.apache.calcite.sql2rel.CorrelationReferenceFinder)2 Test (org.junit.Test)2