Java Code Examples for org.apache.calcite.rel.RelNode#getTraitSet()

@Override
public void onMatch(RelOptRuleCall call) {
  final DrillWriterRel writer = call.rel(0);
  final RelNode input = call.rel(1);

  final List<Integer> keys = writer.getPartitionKeys();
  final RelCollation collation = getCollation(keys);
  final boolean hashDistribute = PrelUtil.getPlannerSettings(call.getPlanner()).getOptions().getOption(ExecConstants.CTAS_PARTITIONING_HASH_DISTRIBUTE_VALIDATOR);
  final RelTraitSet traits = hashDistribute ?
      input.getTraitSet().plus(Prel.DRILL_PHYSICAL).plus(collation).plus(getDistribution(keys)) :
      input.getTraitSet().plus(Prel.DRILL_PHYSICAL).plus(collation);

  final RelNode convertedInput = convert(input, traits);

  if (!new WriteTraitPull(call).go(writer, convertedInput)) {
    DrillWriterRelBase newWriter = new WriterPrel(writer.getCluster(), convertedInput.getTraitSet(),
        convertedInput, writer.getCreateTableEntry());

    call.transformTo(newWriter);
  }
}
 

@Override
public RelNode createSort(RelTraitSet traits, RelNode input, RelCollation collation, RexNode offset, RexNode fetch) {
  RelNode newInput;
  if (!collation.getFieldCollations().isEmpty()) {
    collation = RelCollationTraitDef.INSTANCE.canonize(collation);
    newInput = SortRel.create(input.getCluster(), traits, input, collation, null, null);
    traits = newInput.getTraitSet();
  } else {
    newInput = input;
  }

  if (!isOffsetEmpty(offset) || !isFetchEmpty(fetch)) {
    return LimitRel.create(newInput.getCluster(), traits, newInput, offset, fetch);
  }
  return newInput;
}
 

@Override
public void onMatch(RelOptRuleCall call) {
  final LimitPrel limit = (LimitPrel) call.rel(0);
  final UnionExchangePrel unionExchangePrel = (UnionExchangePrel) call.rel(1);

  RelNode child = unionExchangePrel.getInput();

  final int offset = limit.getOffset() != null ? Math.max(0, RexLiteral.intValue(limit.getOffset())) : 0;
  final int fetch = Math.max(0, RexLiteral.intValue(limit.getFetch()));

  // child Limit uses conservative approach:  use offset 0 and fetch = parent limit offset + parent limit fetch.
  final RexNode childFetch = limit.getCluster().getRexBuilder().makeExactLiteral(BigDecimal.valueOf(offset + fetch));

  final RelNode limitUnderExchange = new LimitPrel(child.getCluster(), child.getTraitSet(), child, null, childFetch);
  final RelNode newUnionExch = new UnionExchangePrel(unionExchangePrel.getCluster(), unionExchangePrel.getTraitSet(), limitUnderExchange);
  final RelNode limitAboveExchange = new LimitPrel(limit.getCluster(), limit.getTraitSet(), newUnionExch, limit.getOffset(), limit.getFetch(), true);

  call.transformTo(limitAboveExchange);
}
 

@Override
public void onMatch(RelOptRuleCall call) {
  final LimitPrel limit = (LimitPrel) call.rel(0);
  final ExchangePrel exchangePrel = (ExchangePrel) call.rel(1);

  RelNode child = exchangePrel.getInput();

  final int offset = limit.getOffset() != null ? Math.max(0, RexLiteral.intValue(limit.getOffset())) : 0;
  final int fetch = Math.max(0, RexLiteral.intValue(limit.getFetch()));

  // child Limit uses conservative approach:  use offset 0 and fetch = parent limit offset + parent limit fetch.
  final RexNode childFetch = limit.getCluster().getRexBuilder().makeExactLiteral(BigDecimal.valueOf(offset + fetch));

  final RelNode limitUnderExchange = new LimitPrel(child.getCluster(), child.getTraitSet(), child, null, childFetch);
  final RelNode newExch = exchangePrel.copy(exchangePrel.getTraitSet(), ImmutableList.of(limitUnderExchange));
  final RelNode limitAboveExchange = new LimitPrel(limit.getCluster(), limit.getTraitSet(), newExch, limit.getOffset(), limit.getFetch(), true);

  call.transformTo(limitAboveExchange);
}
 

/**
 * Converts a relation expression to a given set of traits, if it does not
 * already have those traits.
 *
 * @param rel      Relational expression to convert
 * @param toTraits desired traits
 * @return a relational expression with the desired traits; never null
 */
public static RelNode convert(RelNode rel, RelTraitSet toTraits) {
  RelOptPlanner planner = rel.getCluster().getPlanner();

  RelTraitSet outTraits = rel.getTraitSet();
  for (int i = 0; i < toTraits.size(); i++) {
    RelTrait toTrait = toTraits.getTrait(i);
    if (toTrait != null) {
      outTraits = outTraits.replace(i, toTrait);
    }
  }

  if (rel.getTraitSet().matches(outTraits)) {
    return rel;
  }

  return planner.changeTraits(rel, outTraits);
}
 

public static RelNode removeColumns(RelNode node, Predicate<RelDataTypeField> predicate) {
  if (node.getTraitSet() == null) {
    // for test purposes.
    return node;
  }

  // identify all fields that match passed predicate
  Set<RelDataTypeField> toRemove = FluentIterable.from(node.getRowType().getFieldList()).filter(predicate).toSet();

  if (toRemove.isEmpty()) {
    return node;
  }

  final RexBuilder rexBuilder = node.getCluster().getRexBuilder();
  final RelDataTypeFactory.FieldInfoBuilder rowTypeBuilder = new RelDataTypeFactory.FieldInfoBuilder(node.getCluster().getTypeFactory());
  final List<RexNode> projects = FluentIterable.from(node.getRowType().getFieldList())
    .filter(Predicates.not(toRemove::contains))
    .transform((RelDataTypeField field) -> {
      rowTypeBuilder.add(field);
      return (RexNode) rexBuilder.makeInputRef(field.getType(), field.getIndex());
    }).toList();

  return new LogicalProject(node.getCluster(), node.getTraitSet(), node, projects, rowTypeBuilder.build());
}
 

public RelNode buildOriginalProject (RelNode newRel) {
  RelDataType origRowType = origProject == null ? origScan.getRowType() : origProject.getRowType();

  final RelDataTypeFactory.FieldInfoBuilder finalFieldTypeBuilder =
      origScan.getCluster().getTypeFactory().builder();

  List<RelDataTypeField> hjRowFields = newRel.getRowType().getFieldList();
  int toRemoveRowKeyCount = 1;
  if (getRowKeyIndex(origRowType, origScan)  < 0 ) {
    toRemoveRowKeyCount = 2;
  }
  finalFieldTypeBuilder.addAll(hjRowFields.subList(0, hjRowFields.size()-toRemoveRowKeyCount));
  final RelDataType finalProjectRowType = finalFieldTypeBuilder.build();

  List<RexNode> resetExprs = Lists.newArrayList();
  for (int idx=0; idx<hjRowFields.size()-toRemoveRowKeyCount; ++idx) {
    resetExprs.add(RexInputRef.of(idx, newRel.getRowType()));
  }

  final ProjectPrel resetProjectPrel = new ProjectPrel(newRel.getCluster(), newRel.getTraitSet(),
      newRel, resetExprs, finalProjectRowType);
  newRel = resetProjectPrel;

  RelNode finalRel = Prule.convert(newRel, newRel.getTraitSet());
  return finalRel;
}
 

public RelNode buildRowKeyProject(RelNode inputRel, int fieldIndex) {
  List<RelDataTypeField> inputFields = inputRel.getRowType().getFieldList();
  final RelDataTypeField rowKeyField = inputFields.get(fieldIndex);
    RexNode expr = builder.makeInputRef(rowKeyField.getType(), rowKeyField.getIndex());
  List<RexNode> exprs = Lists.newArrayList();
  exprs.add(expr);

  final RelDataTypeFactory.FieldInfoBuilder rightFieldTypeBuilder =
      inputRel.getCluster().getTypeFactory().builder();

  rightFieldTypeBuilder.add(rowKeyField);
  final RelDataType projectRowType = rightFieldTypeBuilder.build();

  ProjectPrel proj = new ProjectPrel(inputRel.getCluster(), inputRel.getTraitSet(), inputRel, exprs, projectRowType);

  return proj;
}
 

/**
 * Creates a ConverterImpl.
 *
 * @param cluster  planner's cluster
 * @param traitDef the RelTraitDef this converter converts
 * @param traits   the output traits of this converter
 * @param child    child rel (provides input traits)
 */
protected ConverterImpl(
    RelOptCluster cluster,
    RelTraitDef traitDef,
    RelTraitSet traits,
    RelNode child) {
  super(cluster, traits, child);
  this.inTraits = child.getTraitSet();
  this.traitDef = traitDef;
}
 

@Test void testTableReferencesUnionUnknownNode() {
  final String sql = "select * from emp limit 10";
  final RelNode node = convertSql(sql);
  final RelNode nodeWithUnknown = new DummyRelNode(
      node.getCluster(), node.getTraitSet(), node);
  // Union
  final LogicalUnion union =
      LogicalUnion.create(ImmutableList.of(nodeWithUnknown, node),
          true);
  final RelMetadataQuery mq = node.getCluster().getMetadataQuery();
  final Set<RelTableRef> tableReferences = mq.getTableReferences(union);
  assertNull(tableReferences);
}
 

private RelNode rename(RelNode input, List<RelDataTypeField> inputFields, List<String> outputFieldNames) {
  List<RexNode> exprs = Lists.newArrayList();

  for (RelDataTypeField field : inputFields) {
    RexNode expr = input.getCluster().getRexBuilder().makeInputRef(field.getType(), field.getIndex());
    exprs.add(expr);
  }

  RelDataType rowType = RexUtil.createStructType(input.getCluster().getTypeFactory(),
      exprs, outputFieldNames, null);

  ProjectPrel proj = new ProjectPrel(input.getCluster(), input.getTraitSet(), input, exprs, rowType);

  return proj;
}
 

/**
 * Converts one trait of a relational expression, if it does not
 * already have that trait.
 *
 * @param rel      Relational expression to convert
 * @param toTrait  Desired trait
 * @return a relational expression with the desired trait; never null
 */
public static RelNode convert(RelNode rel, RelTrait toTrait) {
  RelOptPlanner planner = rel.getCluster().getPlanner();
  RelTraitSet outTraits = rel.getTraitSet();
  if (toTrait != null) {
    outTraits = outTraits.replace(toTrait);
  }

  if (rel.getTraitSet().matches(outTraits)) {
    return rel;
  }

  return planner.changeTraits(rel, outTraits.simplify());
}
 

@Override public RelNode createJoin(RelNode left, RelNode right, List<RelHint> hints,
    RexNode condition, Set<CorrelationId> variablesSet, JoinRelType joinType,
    boolean semiJoinDone) {
  Util.discard(hints);
  Util.discard(variablesSet);
  Util.discard(semiJoinDone);
  return new PigJoin(left.getCluster(), left.getTraitSet(), left, right, condition, joinType);
}
 

/**
 * Merges this limit with an ElasticsearchSample operator.
 * @param sample
 * @param treeWithoutLimit
 * @return
 */
public ElasticsearchLimit merge(ElasticsearchSample sample, RelNode treeWithoutLimit){
  if(sample == null){
    return this;
  }

  long sampleSize = SampleCrel.getSampleSizeAndSetMinSampleSize(PrelUtil.getPlannerSettings(getCluster().getPlanner()), ElasticSampleRule.SAMPLE_SIZE_DENOMINATOR);
  int limitAmount = RexLiteral.intValue(getFetch());
  int finalLimit = Math.min((int) sampleSize,  limitAmount);
  RexNode offset = getCluster().getRexBuilder().makeExactLiteral(BigDecimal.valueOf(0), getCluster().getTypeFactory().createSqlType(SqlTypeName.INTEGER));
  RexNode fetch = getCluster().getRexBuilder().makeExactLiteral(BigDecimal.valueOf(finalLimit), getCluster().getTypeFactory().createSqlType(SqlTypeName.INTEGER));
  return new ElasticsearchLimit(getCluster(), treeWithoutLimit.getTraitSet(), treeWithoutLimit, offset, fetch, isPushDown(), getPluginId());
}
 

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

@Override
public RelNode convertChild(DrillFilterRel filter, RelNode rel) {
  return new FilterPrel(filter.getCluster(), rel.getTraitSet(), rel, filter.getCondition());
}
 

@Override
public RelNode createCorrelate(RelNode left, RelNode right, CorrelationId correlationId, ImmutableBitSet requiredColumns, SemiJoinType joinType) {
  return new CorrelateRel(left.getCluster(), left.getTraitSet(), left, right, correlationId, requiredColumns, joinType);
}
 

@Override
public RelNode convertChild(FilterRel filter, RelNode rel) {
  return new FilterPrel(filter.getCluster(), rel.getTraitSet(), rel, filter.getCondition());
}
 

public static FilterRel create(RelNode child, RexNode condition) {
  return new FilterRel(child.getCluster(), child.getTraitSet(), child, condition)  ;
}
 

private InvalidViewRel(ViewTable viewTable, RelNode input) {
  super(input.getCluster(), input.getTraitSet(), input);
  this.rowType = viewTable.getRowType(getCluster().getTypeFactory());
  this.viewTable = viewTable;
}