Java Code Examples for org.apache.calcite.plan.RelTraitSet#plus()

@Test void testRelDistributionSatisfy() {
  RelDistribution distribution1 = RelDistributions.hash(ImmutableList.of(0));
  RelDistribution distribution2 = RelDistributions.hash(ImmutableList.of(1));

  RelTraitSet traitSet = RelTraitSet.createEmpty();
  RelTraitSet simpleTrait1 = traitSet.plus(distribution1);
  RelTraitSet simpleTrait2 = traitSet.plus(distribution2);
  RelTraitSet compositeTrait =
      traitSet.replace(RelDistributionTraitDef.INSTANCE,
          ImmutableList.of(distribution1, distribution2));

  assertThat(compositeTrait.satisfies(simpleTrait1), is(true));
  assertThat(compositeTrait.satisfies(simpleTrait2), is(true));

  assertThat(distribution1.compareTo(distribution2), is(-1));
  assertThat(distribution2.compareTo(distribution1), is(1));
  //noinspection EqualsWithItself
  assertThat(distribution2.compareTo(distribution2), is(0));
}
 

private static RelTraitSet traits(
    RelOptCluster cluster,
    double rowCount,
    int splitCount,
    RelTraitSet traitSet) {
  PlannerSettings settings = PrelUtil.getPlannerSettings(cluster.getPlanner());
  boolean smallInput = rowCount < settings.getSliceTarget();

  DistributionTrait distribution;
  if(settings.isMultiPhaseAggEnabled() && !settings.isSingleMode() && !smallInput && splitCount > 1) {
    distribution = DistributionTrait.ANY;
  } else {
    distribution = DistributionTrait.SINGLETON;
  }

  return traitSet.plus(distribution);
}
 

public RelTraitSet inferTraits(final RelTraitSet inputTraitSet, final RelDataType inputRowType) {
  final RelTraitSet relTraits = inputTraitSet.plus(Prel.PHYSICAL);

  if (hasDistributions()) {
    return relTraits.plus(hashDistributedOn(distributionColumns, inputRowType));
  }

  if (hasPartitions()) {
    switch (partitionDistributionStrategy) {

    case HASH:
      return relTraits.plus(hashDistributedOn(partitionColumns, inputRowType));

    case ROUND_ROBIN:
      return relTraits.plus(DistributionTrait.ROUND_ROBIN);

    case UNSPECIFIED:
    case STRIPED:
      // fall through ..
    }
  }

  return relTraits;
}
 

@Override
public void onMatch(RelOptRuleCall call) {
  final DrillSortRel sort = call.rel(0);
  final RelNode input = sort.getInput();

  // Keep the collation in logical sort. Convert input into a RelNode with 1) this collation, 2) Physical, 3) hash distributed on

  DrillDistributionTrait hashDistribution =
          new DrillDistributionTrait(DrillDistributionTrait.DistributionType.HASH_DISTRIBUTED, ImmutableList.copyOf(getDistributionField(sort)));

  final RelTraitSet traits = RelTraitSet.createEmpty().plus(Prel.DRILL_PHYSICAL).plus(hashDistribution);
  SortPrel child = new SortPrel(sort.getCluster(), traits.plus(sort.getCollation()),
          convert(sort.getInput(), traits), sort.getCollation(), false);

  if(isSingleMode(call)){
    call.transformTo(child);
  }else{
    RelNode exch = new SingleMergeExchangePrel(sort.getCluster(), sort.getTraitSet().plus(Prel.DRILL_PHYSICAL).plus(DrillDistributionTrait.SINGLETON), child, sort.getCollation());
    call.transformTo(exch);  // transform logical "sort" into "SingleMergeExchange".

  }
}
 

@Override
public void onMatch(RelOptRuleCall call) {
  final DrillLateralJoinRel lateralJoinRel = call.rel(0);
  final RelNode left = lateralJoinRel.getLeft();
  final RelNode right = lateralJoinRel.getRight();
  RelTraitSet traitsLeft = left.getTraitSet().plus(Prel.DRILL_PHYSICAL);
  RelTraitSet traitsRight = right.getTraitSet().plus(Prel.DRILL_PHYSICAL);

  RelTraitSet corrTraits = traitsLeft.plus(DrillDistributionTrait.RANDOM_DISTRIBUTED);

  final RelNode convertedLeft = convert(left, traitsLeft);
  final RelNode convertedRight = convert(right, traitsRight);

  final LateralJoinPrel lateralJoinPrel = new LateralJoinPrel(lateralJoinRel.getCluster(),
                                corrTraits,
                                convertedLeft, convertedRight, lateralJoinRel.excludeCorrelateColumn, lateralJoinRel.getCorrelationId(),
                                lateralJoinRel.getRequiredColumns(),lateralJoinRel.getJoinType());
  call.transformTo(lateralJoinPrel);
}
 

public static RelTraitSet removeCollation(RelTraitSet traitSet, RelOptRuleCall call) {
    RelTraitSet newTraitSet = call.getPlanner().emptyTraitSet();

    for (RelTrait trait : traitSet) {
        if (!trait.getTraitDef().getTraitClass().equals(RelCollation.class)) {
            newTraitSet = newTraitSet.plus(trait);
        }
    }

    return newTraitSet;
}
 

@Override public RelTraitSet emptyTraitSet() {
  RelTraitSet traitSet = super.emptyTraitSet();
  for (RelTraitDef traitDef : traitDefs) {
    if (traitDef.multiple()) {
      // TODO: restructure RelTraitSet to allow a list of entries
      //  for any given trait
    }
    traitSet = traitSet.plus(traitDef.getDefault());
  }
  return traitSet;
}
 

public RelTraitSet newTraitSet(RelTrait... traits) {
  RelTraitSet set = call.getPlanner().emptyTraitSet();
  for (RelTrait t : traits) {
    set = set.plus(t);
  }
  return set;

}
 

@Override public RelTraitSet emptyTraitSet() {
  RelTraitSet traitSet = super.emptyTraitSet();
  for (RelTraitDef traitDef : traitDefs) {
    if (traitDef.multiple()) {
      // TODO: restructure RelTraitSet to allow a list of entries
      //  for any given trait
    }
    traitSet = traitSet.plus(traitDef.getDefault());
  }
  return traitSet;
}
 

public RelTraitSet newTraitSet(RelTrait... traits) {
  RelTraitSet set = call.getPlanner().emptyTraitSet();
  for (RelTrait t : traits) {
    set = set.plus(t);
  }
  return set;

}
 

@Override
public RelNode convert(RelNode r) {
  Sort rel = (Sort) r;
  RelTraitSet traits = rel.getInput().getTraitSet().replace(Prel.DRILL_PHYSICAL);
  return new SortPrel(rel.getCluster(),
                      traits.plus(rel.getCollation()),
                      convert(rel.getInput(), traits.simplify()),
                      rel.getCollation());
}
 

public RelTraitSet newTraitSet(RelTrait... traits) {
  RelTraitSet set = indexContext.getCall().getPlanner().emptyTraitSet();
  for (RelTrait t : traits) {
    if(t != null) {
      set = set.plus(t);
    }
  }
  return set;
}
 

public static ScanPrel buildCoveringIndexScan(DrillScanRelBase origScan,
    IndexGroupScan indexGroupScan,
    IndexCallContext indexContext,
    IndexDescriptor indexDesc) {

  FunctionalIndexInfo functionInfo = indexDesc.getFunctionalInfo();
  //to record the new (renamed)paths added
  List<SchemaPath> rewrittenPaths = Lists.newArrayList();
  DbGroupScan dbGroupScan = (DbGroupScan) getGroupScan(origScan);
  indexGroupScan.setColumns(
      rewriteFunctionColumn(dbGroupScan.getColumns(),
          functionInfo, rewrittenPaths));

  DrillDistributionTrait partition = scanIsPartition(getGroupScan(origScan))?
      DrillDistributionTrait.RANDOM_DISTRIBUTED : DrillDistributionTrait.SINGLETON;
  RelDataType newRowType = FunctionalIndexHelper.rewriteFunctionalRowType(origScan, indexContext, functionInfo, rewrittenPaths);

  // add a default collation trait otherwise Calcite runs into a ClassCastException, which at first glance
  // seems like a Calcite bug
  RelTraitSet indexScanTraitSet = origScan.getTraitSet().plus(Prel.DRILL_PHYSICAL).
      plus(RelCollationTraitDef.INSTANCE.getDefault()).plus(partition);

  // Create the collation traits for index scan based on the index columns under the
  // condition that the index actually has collation property (e.g hash indexes don't)
  if (indexDesc.getCollation() != null) {
    RelCollation collationTrait = buildCollationCoveringIndexScan(indexDesc, indexContext);
    indexScanTraitSet = indexScanTraitSet.plus(collationTrait);
  }

  ScanPrel indexScanPrel = new ScanPrel(origScan.getCluster(),
      indexScanTraitSet, indexGroupScan,
      newRowType, origScan.getTable());

  return indexScanPrel;
}
 

public LimitRelBase(RelOptCluster cluster, RelTraitSet traitSet, RelNode child, RexNode offset, RexNode fetch, boolean pushDown) {
  super(cluster, traitSet.plus(RelCollations.EMPTY), child, RelCollations.EMPTY, offset, fetch);
  this.traitSet = traitSet;
  this.pushDown = pushDown;
}
 

protected static RelTraitSet adjustTraits(RelTraitSet traits, RelCollation collation) {
  return traits.plus(collation);
}
 

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;
}
 

@Override
public RelNode convertChild(final RelNode filter, final RelNode input) throws InvalidRelException {

    if (indexGroupScan == null) {
        logger.error("Null indexgroupScan in CoveringIndexPlanGenerator.convertChild");
        return null;
    }

    RexNode coveringCondition;
    ScanPrel indexScanPrel = IndexPlanUtils.buildCoveringIndexScan(origScan, indexGroupScan, indexContext,
            indexDesc);

    // If remainder condition, then combine the index and remainder conditions. This is a covering plan so we can
    // pushed the entire condition into the index.
    coveringCondition = IndexPlanUtils.getTotalFilter(indexCondition, remainderCondition,
            indexScanPrel.getCluster().getRexBuilder());
    RexNode newIndexCondition = rewriteFunctionalCondition(coveringCondition, indexScanPrel.getRowType(),
            functionInfo);

    // build collation for filter
    RelTraitSet indexFilterTraitSet = indexScanPrel.getTraitSet();

    FilterPrel indexFilterPrel = new FilterPrel(indexScanPrel.getCluster(), indexFilterTraitSet, indexScanPrel,
            newIndexCondition);

    ProjectPrel indexProjectPrel = null;
    if (origProject != null) {
        RelCollation collation = IndexPlanUtils.buildCollationProject(IndexPlanUtils.getProjects(origProject), null,
                origScan, functionInfo, indexContext);
        indexProjectPrel = new ProjectPrel(origScan.getCluster(), indexFilterTraitSet.plus(collation),
                indexFilterPrel, IndexPlanUtils.getProjects(origProject), origProject.getRowType());
    }

    RelNode finalRel;
    if (indexProjectPrel != null) {
        finalRel = indexProjectPrel;
    } else {
        finalRel = indexFilterPrel;
    }

    if (upperProject != null) {
        RelCollation newCollation = IndexPlanUtils.buildCollationProject(IndexPlanUtils.getProjects(upperProject),
                origProject, origScan, functionInfo, indexContext);

        ProjectPrel cap = new ProjectPrel(upperProject.getCluster(),
                newCollation == null ? finalRel.getTraitSet() : finalRel.getTraitSet().plus(newCollation), finalRel,
                IndexPlanUtils.getProjects(upperProject), upperProject.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 = cap;
            if (indexProjectPrel != null) {
                newProject = (ProjectPrel) DrillMergeProjectRule.replace(newProject, indexProjectPrel);
            }
            // 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(),
                    newCollation == null ? newProject.getTraitSet() : newProject.getTraitSet().plus(newCollation),
                    indexFilterPrel, newProjects, newProject.getRowType());

            cap = rewrittenProject;
        }

        finalRel = cap;
    }

    if (indexContext.getSort() != null) {
        finalRel = getSortNode(indexContext, finalRel, false, true, true);
        Preconditions.checkArgument(finalRel != null);
    }

    finalRel = Prule.convert(finalRel, finalRel.getTraitSet().plus(Prel.DRILL_PHYSICAL));

    logger.debug("CoveringIndexPlanGenerator got finalRel {} from origScan {}, original digest {}, new digest {}.",
            finalRel.toString(), origScan.toString(),
            upperProject == null ? indexContext.getFilter().getDigest() : upperProject.getDigest(),
            finalRel.getDigest());
    return finalRel;
}