Examples with RelTraitSet org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.plan.RelTraitSet used on opensource projects

Example 46 with RelTraitSet

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.plan.RelTraitSet in project drill by apache.

the class WindowPrule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    final DrillWindowRel window = call.rel(0);
    RelNode input = call.rel(1);
    // TODO: Order window based on existing partition by
    // input.getTraitSet().subsumes()
    boolean partitionby = false;
    boolean addMerge = false;
    // The start index of the constant fields of DrillWindowRel
    final int startConstantsIndex = window.getInput().getRowType().getFieldCount();
    int constantShiftIndex = 0;
    for (final Ord<Window.Group> w : Ord.zip(window.groups)) {
        Window.Group windowBase = w.getValue();
        RelTraitSet traits = call.getPlanner().emptyTraitSet().plus(Prel.DRILL_PHYSICAL);
        // For empty Over-Clause
        if (windowBase.keys.isEmpty() && windowBase.orderKeys.getFieldCollations().isEmpty()) {
            DrillDistributionTrait distEmptyKeys = new DrillDistributionTrait(DrillDistributionTrait.DistributionType.SINGLETON);
            traits = traits.plus(distEmptyKeys);
        } else if (windowBase.keys.size() > 0) {
            DrillDistributionTrait distOnAllKeys = new DrillDistributionTrait(DrillDistributionTrait.DistributionType.HASH_DISTRIBUTED, ImmutableList.copyOf(getDistributionFields(windowBase)));
            partitionby = true;
            traits = traits.plus(distOnAllKeys);
        } else if (windowBase.orderKeys.getFieldCollations().size() > 0) {
            // if only the order-by clause is specified, there is a single partition
            // consisting of all the rows, so we do a distributed sort followed by a
            // single merge as the input of the window operator
            DrillDistributionTrait distKeys = new DrillDistributionTrait(DrillDistributionTrait.DistributionType.HASH_DISTRIBUTED, ImmutableList.copyOf(getDistributionFieldsFromCollation(windowBase)));
            traits = traits.plus(distKeys);
            if (!isSingleMode(call)) {
                addMerge = true;
            }
        }
        // Add collation trait if either partition-by or order-by is specified.
        if (partitionby || windowBase.orderKeys.getFieldCollations().size() > 0) {
            RelCollation collation = getCollation(windowBase);
            traits = traits.plus(collation);
        }
        RelNode convertedInput = convert(input, traits);
        if (addMerge) {
            traits = traits.plus(DrillDistributionTrait.SINGLETON);
            convertedInput = new SingleMergeExchangePrel(window.getCluster(), traits, convertedInput, windowBase.collation());
        }
        List<RelDataTypeField> newRowFields = Lists.newArrayList();
        newRowFields.addAll(convertedInput.getRowType().getFieldList());
        Iterable<RelDataTypeField> newWindowFields = Iterables.filter(window.getRowType().getFieldList(), new Predicate<RelDataTypeField>() {

            @Override
            public boolean apply(RelDataTypeField relDataTypeField) {
                return relDataTypeField.getName().startsWith("w" + w.i + "$");
            }
        });
        for (RelDataTypeField newField : newWindowFields) {
            newRowFields.add(newField);
        }
        RelDataType rowType = new RelRecordType(newRowFields);
        List<Window.RexWinAggCall> newWinAggCalls = Lists.newArrayList();
        for (Ord<Window.RexWinAggCall> aggOrd : Ord.zip(windowBase.aggCalls)) {
            Window.RexWinAggCall aggCall = aggOrd.getValue();
            // If the argument points at the constant and
            // additional fields have been generated by the Window below,
            // the index of constants will be shifted
            final List<RexNode> newOperandsOfWindowFunction = Lists.newArrayList();
            for (RexNode operand : aggCall.getOperands()) {
                if (operand instanceof RexInputRef) {
                    final RexInputRef rexInputRef = (RexInputRef) operand;
                    final int refIndex = rexInputRef.getIndex();
                    // Check if this RexInputRef points at the constants
                    if (rexInputRef.getIndex() >= startConstantsIndex) {
                        operand = new RexInputRef(refIndex + constantShiftIndex, window.constants.get(refIndex - startConstantsIndex).getType());
                    }
                }
                newOperandsOfWindowFunction.add(operand);
            }
            aggCall = new Window.RexWinAggCall((SqlAggFunction) aggCall.getOperator(), aggCall.getType(), newOperandsOfWindowFunction, aggCall.ordinal, aggCall.distinct);
            newWinAggCalls.add(new Window.RexWinAggCall((SqlAggFunction) aggCall.getOperator(), aggCall.getType(), aggCall.getOperands(), aggOrd.i, aggCall.distinct));
        }
        windowBase = new Window.Group(windowBase.keys, windowBase.isRows, windowBase.lowerBound, windowBase.upperBound, windowBase.orderKeys, newWinAggCalls);
        input = new WindowPrel(window.getCluster(), window.getTraitSet().merge(traits), convertedInput, window.getConstants(), rowType, windowBase);
        constantShiftIndex += windowBase.aggCalls.size();
    }
    call.transformTo(input);
}

Example 47 with RelTraitSet

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.plan.RelTraitSet in project drill by apache.

the class StreamAggPrule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    final DrillAggregateRel aggregate = call.rel(0);
    RelNode input = aggregate.getInput();
    final RelCollation collation = getCollation(aggregate);
    RelTraitSet traits;
    if (aggregate.containsDistinctCall()) {
        // currently, don't use StreamingAggregate if any of the logical aggrs contains DISTINCT
        return;
    }
    try {
        if (aggregate.getGroupSet().isEmpty()) {
            DrillDistributionTrait singleDist = DrillDistributionTrait.SINGLETON;
            final RelTraitSet singleDistTrait = call.getPlanner().emptyTraitSet().plus(Prel.DRILL_PHYSICAL).plus(singleDist);
            if (create2PhasePlan(call, aggregate)) {
                traits = call.getPlanner().emptyTraitSet().plus(Prel.DRILL_PHYSICAL);
                RelNode convertedInput = convert(input, traits);
                new SubsetTransformer<DrillAggregateRel, InvalidRelException>(call) {

                    @Override
                    public RelNode convertChild(final DrillAggregateRel join, final RelNode rel) throws InvalidRelException {
                        DrillDistributionTrait toDist = rel.getTraitSet().getTrait(DrillDistributionTraitDef.INSTANCE);
                        RelTraitSet traits = newTraitSet(Prel.DRILL_PHYSICAL, toDist);
                        RelNode newInput = convert(rel, traits);
                        StreamAggPrel phase1Agg = new StreamAggPrel(aggregate.getCluster(), traits, newInput, aggregate.getGroupSet(), aggregate.getGroupSets(), aggregate.getAggCallList(), OperatorPhase.PHASE_1of2);
                        UnionExchangePrel exch = new UnionExchangePrel(phase1Agg.getCluster(), singleDistTrait, phase1Agg);
                        ImmutableBitSet newGroupSet = remapGroupSet(aggregate.getGroupSet());
                        List<ImmutableBitSet> newGroupSets = Lists.newArrayList();
                        for (ImmutableBitSet groupSet : aggregate.getGroupSets()) {
                            newGroupSets.add(remapGroupSet(groupSet));
                        }
                        return new StreamAggPrel(aggregate.getCluster(), singleDistTrait, exch, newGroupSet, newGroupSets, phase1Agg.getPhase2AggCalls(), OperatorPhase.PHASE_2of2);
                    }
                }.go(aggregate, convertedInput);
            } else {
                createTransformRequest(call, aggregate, input, singleDistTrait);
            }
        } else {
            // hash distribute on all grouping keys
            final DrillDistributionTrait distOnAllKeys = new DrillDistributionTrait(DrillDistributionTrait.DistributionType.HASH_DISTRIBUTED, ImmutableList.copyOf(getDistributionField(aggregate, true)));
            traits = call.getPlanner().emptyTraitSet().plus(Prel.DRILL_PHYSICAL).plus(collation).plus(distOnAllKeys);
            createTransformRequest(call, aggregate, input, traits);
            // hash distribute on one grouping key
            DrillDistributionTrait distOnOneKey = new DrillDistributionTrait(DrillDistributionTrait.DistributionType.HASH_DISTRIBUTED, ImmutableList.copyOf(getDistributionField(aggregate, false)));
            traits = call.getPlanner().emptyTraitSet().plus(Prel.DRILL_PHYSICAL).plus(collation).plus(distOnOneKey);
            if (create2PhasePlan(call, aggregate)) {
                traits = call.getPlanner().emptyTraitSet().plus(Prel.DRILL_PHYSICAL);
                RelNode convertedInput = convert(input, traits);
                new SubsetTransformer<DrillAggregateRel, InvalidRelException>(call) {

                    @Override
                    public RelNode convertChild(final DrillAggregateRel aggregate, final RelNode rel) throws InvalidRelException {
                        DrillDistributionTrait toDist = rel.getTraitSet().getTrait(DrillDistributionTraitDef.INSTANCE);
                        RelTraitSet traits = newTraitSet(Prel.DRILL_PHYSICAL, collation, toDist);
                        RelNode newInput = convert(rel, traits);
                        StreamAggPrel phase1Agg = new StreamAggPrel(aggregate.getCluster(), traits, newInput, aggregate.getGroupSet(), aggregate.getGroupSets(), aggregate.getAggCallList(), OperatorPhase.PHASE_1of2);
                        int numEndPoints = PrelUtil.getSettings(phase1Agg.getCluster()).numEndPoints();
                        HashToMergeExchangePrel exch = new HashToMergeExchangePrel(phase1Agg.getCluster(), phase1Agg.getTraitSet().plus(Prel.DRILL_PHYSICAL).plus(distOnAllKeys), phase1Agg, ImmutableList.copyOf(getDistributionField(aggregate, true)), collation, numEndPoints);
                        ImmutableBitSet newGroupSet = remapGroupSet(aggregate.getGroupSet());
                        List<ImmutableBitSet> newGroupSets = Lists.newArrayList();
                        for (ImmutableBitSet groupSet : aggregate.getGroupSets()) {
                            newGroupSets.add(remapGroupSet(groupSet));
                        }
                        return new StreamAggPrel(aggregate.getCluster(), exch.getTraitSet(), exch, newGroupSet, newGroupSets, phase1Agg.getPhase2AggCalls(), OperatorPhase.PHASE_2of2);
                    }
                }.go(aggregate, convertedInput);
            }
        }
    } catch (InvalidRelException e) {
        tracer.warn(e.toString());
    }
}

Example 48 with RelTraitSet

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.plan.RelTraitSet in project drill by apache.

the class DrillCassandraLimitRule method convert.

public RelNode convert(RelNode relNode) {
    DrillLimitRelBase limit = (DrillLimitRelBase) relNode;
    final RelTraitSet traitSet = limit.getTraitSet().replace(CassandraRel.CONVENTION);
    return new CassandraLimit(limit.getCluster(), traitSet, convert(limit.getInput(), CassandraRel.CONVENTION), limit.getOffset(), limit.getFetch());
}

Example 49 with RelTraitSet

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.plan.RelTraitSet in project drill by apache.

the class CoveringPlanNoFilterGenerator method convertChild.

public RelNode convertChild() throws InvalidRelException {
    Preconditions.checkNotNull(indexContext.getSort());
    if (indexGroupScan == null) {
        logger.error("Null indexgroupScan in CoveringIndexPlanGenerator.convertChild");
        return null;
    }
    // update sort expressions in context
    IndexPlanUtils.updateSortExpression(indexContext, indexContext.getSort() != null ? indexContext.getCollation().getFieldCollations() : null);
    ScanPrel indexScanPrel = IndexPlanUtils.buildCoveringIndexScan(origScan, indexGroupScan, indexContext, indexDesc);
    ((IndexGroupScan) indexScanPrel.getGroupScan()).setStatistics(((DbGroupScan) IndexPlanUtils.getGroupScan(origScan)).getStatistics());
    RelTraitSet indexScanTraitSet = indexScanPrel.getTraitSet();
    RelNode finalRel = indexScanPrel;
    if (indexContext.getLowerProject() != null) {
        RelCollation collation = IndexPlanUtils.buildCollationProject(indexContext.getLowerProject().getProjects(), null, indexContext.getScan(), functionInfo, indexContext);
        finalRel = new ProjectPrel(indexContext.getScan().getCluster(), indexScanTraitSet.plus(collation), indexScanPrel, indexContext.getLowerProject().getProjects(), indexContext.getLowerProject().getRowType());
        if (functionInfo.hasFunctional()) {
            // if there is functional index field, then a rewrite may be needed in upperProject/indexProject
            // merge upperProject with indexProjectPrel(from origProject) if both exist,
            ProjectPrel newProject = (ProjectPrel) finalRel;
            // then rewrite functional expressions in new project.
            List<RexNode> newProjects = Lists.newArrayList();
            DrillParseContext parseContxt = new DrillParseContext(PrelUtil.getPlannerSettings(newProject.getCluster()));
            for (RexNode projectRex : newProject.getProjects()) {
                RexNode newRex = IndexPlanUtils.rewriteFunctionalRex(indexContext, parseContxt, null, origScan, projectRex, indexScanPrel.getRowType(), functionInfo);
                newProjects.add(newRex);
            }
            ProjectPrel rewrittenProject = new ProjectPrel(newProject.getCluster(), collation == null ? newProject.getTraitSet() : newProject.getTraitSet().plus(collation), indexScanPrel, newProjects, newProject.getRowType());
            finalRel = rewrittenProject;
        }
    }
    finalRel = getSortNode(indexContext, finalRel, true, isSingletonSortedStream, indexContext.getExchange() != null);
    if (finalRel == null) {
        return null;
    }
    finalRel = Prule.convert(finalRel, finalRel.getTraitSet().plus(Prel.DRILL_PHYSICAL));
    logger.debug("CoveringPlanNoFilterGenerator got finalRel {} from origScan {}, original digest {}, new digest {}.", finalRel.toString(), indexContext.getScan().toString(), indexContext.getLowerProject() != null ? indexContext.getLowerProject().getDigest() : indexContext.getScan().getDigest(), finalRel.getDigest());
    return finalRel;
}

Example 50 with RelTraitSet

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.plan.RelTraitSet in project drill by apache.

the class DirPrunedEnumerableTableScan method create.

/**
 * Creates an DirPrunedEnumerableTableScan.
 */
public static EnumerableTableScan create(RelOptCluster cluster, RelOptTable relOptTable, String digestFromSelection) {
    final Table table = relOptTable.unwrap(Table.class);
    Class elementType = EnumerableTableScan.deduceElementType(table);
    final RelTraitSet traitSet = cluster.traitSetOf(EnumerableConvention.INSTANCE).replaceIfs(RelCollationTraitDef.INSTANCE, () -> {
        if (table != null) {
            return table.getStatistic().getCollations();
        }
        return ImmutableList.of();
    });
    return new DirPrunedEnumerableTableScan(cluster, traitSet, relOptTable, elementType, digestFromSelection);
}