Java Code Examples for org.apache.calcite.rex.RexUtil#apply()

public void rewriteRel(Sort rel) {
    RelCollation oldCollation = rel.getCollation();
    final RelNode oldChild = rel.getInput();
    final RelNode newChild = getNewForOldRel(oldChild);
    final Mappings.TargetMapping mapping = getNewForOldInputMapping(oldChild);

    // validate
    for (RelFieldCollation field : oldCollation.getFieldCollations()) {
        int oldInput = field.getFieldIndex();
        RelDataType sortFieldType = oldChild.getRowType().getFieldList().get(oldInput).getType();
        if (sortFieldType.isStruct()) {
            // TODO jvs 10-Feb-2005
            throw Util.needToImplement("sorting on structured types");
        }
    }
    RelCollation newCollation = RexUtil.apply(mapping, oldCollation);
    Sort newRel = LogicalSort.create(newChild, newCollation, rel.offset, rel.fetch);
    setNewForOldRel(rel, newRel);
}
 

public void rewriteRel(Sort rel) {
  RelCollation oldCollation = rel.getCollation();
  final RelNode oldChild = rel.getInput();
  final RelNode newChild = getNewForOldRel(oldChild);
  final Mappings.TargetMapping mapping =
      getNewForOldInputMapping(oldChild);

  // validate
  for (RelFieldCollation field : oldCollation.getFieldCollations()) {
    int oldInput = field.getFieldIndex();
    RelDataType sortFieldType =
        oldChild.getRowType().getFieldList().get(oldInput).getType();
    if (sortFieldType.isStruct()) {
      // TODO jvs 10-Feb-2005
      throw Util.needToImplement("sorting on structured types");
    }
  }
  RelCollation newCollation = RexUtil.apply(mapping, oldCollation);
  Sort newRel =
      LogicalSort.create(newChild, newCollation, rel.offset, rel.fetch);
  setNewForOldRel(rel, newRel);
}
 

/**
 * Returns a bit set describing the inputs used by a collection of
 * project expressions and an optional condition.
 */
public static ImmutableBitSet bits(List<RexNode> exprs, RexNode expr) {
    final InputFinder inputFinder = new InputFinder();
    RexUtil.apply(inputFinder, exprs, expr);
    return inputFinder.inputBitSet.build();
}
 

@Override public void onMatch(RelOptRuleCall call) {
  final Union union = call.rel(0);

  final RexBuilder rexBuilder = union.getCluster().getRexBuilder();
  final RelMetadataQuery mq = call.getMetadataQuery();
  final RelOptPredicateList predicates = mq.getPulledUpPredicates(union);
  if (predicates == null) {
    return;
  }

  final Map<Integer, RexNode> constants = new HashMap<>();
  for (Map.Entry<RexNode, RexNode> e : predicates.constantMap.entrySet()) {
    if (e.getKey() instanceof RexInputRef) {
      constants.put(((RexInputRef) e.getKey()).getIndex(), e.getValue());
    }
  }

  // None of the expressions are constant. Nothing to do.
  if (constants.isEmpty()) {
    return;
  }

  // Create expressions for Project operators before and after the Union
  List<RelDataTypeField> fields = union.getRowType().getFieldList();
  List<RexNode> topChildExprs = new ArrayList<>();
  List<String> topChildExprsFields = new ArrayList<>();
  List<RexNode> refs = new ArrayList<>();
  ImmutableBitSet.Builder refsIndexBuilder = ImmutableBitSet.builder();
  for (RelDataTypeField field : fields) {
    final RexNode constant = constants.get(field.getIndex());
    if (constant != null) {
      topChildExprs.add(constant);
      topChildExprsFields.add(field.getName());
    } else {
      final RexNode expr = rexBuilder.makeInputRef(union, field.getIndex());
      topChildExprs.add(expr);
      topChildExprsFields.add(field.getName());
      refs.add(expr);
      refsIndexBuilder.set(field.getIndex());
    }
  }
  ImmutableBitSet refsIndex = refsIndexBuilder.build();

  // Update top Project positions
  final Mappings.TargetMapping mapping =
      RelOptUtil.permutation(refs, union.getInput(0).getRowType()).inverse();
  topChildExprs = RexUtil.apply(mapping, topChildExprs);

  // Create new Project-Union-Project sequences
  final RelBuilder relBuilder = call.builder();
  for (RelNode input : union.getInputs()) {
    List<Pair<RexNode, String>> newChildExprs = new ArrayList<>();
    for (int j : refsIndex) {
      newChildExprs.add(
          Pair.of(rexBuilder.makeInputRef(input, j),
              input.getRowType().getFieldList().get(j).getName()));
    }
    if (newChildExprs.isEmpty()) {
      // At least a single item in project is required.
      newChildExprs.add(
          Pair.of(topChildExprs.get(0), topChildExprsFields.get(0)));
    }
    // Add the input with project on top
    relBuilder.push(input);
    relBuilder.project(Pair.left(newChildExprs), Pair.right(newChildExprs));
  }
  relBuilder.union(union.all, union.getInputs().size());
  // Create top Project fixing nullability of fields
  relBuilder.project(topChildExprs, topChildExprsFields);
  relBuilder.convert(union.getRowType(), false);

  call.transformTo(relBuilder.build());
}
 

public TrimResult trimFields(
    SortExchange sortExchange,
    ImmutableBitSet fieldsUsed,
    Set<RelDataTypeField> extraFields) {
  final RelDataType rowType = sortExchange.getRowType();
  final int fieldCount = rowType.getFieldCount();
  final RelCollation collation = sortExchange.getCollation();
  final RelDistribution distribution = sortExchange.getDistribution();
  final RelNode input = sortExchange.getInput();

  // We use the fields used by the consumer, plus any fields used as sortExchange
  // keys.
  final ImmutableBitSet.Builder inputFieldsUsed = fieldsUsed.rebuild();
  for (RelFieldCollation field : collation.getFieldCollations()) {
    inputFieldsUsed.set(field.getFieldIndex());
  }
  for (int keyIndex : distribution.getKeys()) {
    inputFieldsUsed.set(keyIndex);
  }

  // Create input with trimmed columns.
  final Set<RelDataTypeField> inputExtraFields = Collections.emptySet();
  TrimResult trimResult =
      trimChild(sortExchange, input, inputFieldsUsed.build(), inputExtraFields);
  RelNode newInput = trimResult.left;
  final Mapping inputMapping = trimResult.right;

  // If the input is unchanged, and we need to project all columns,
  // there's nothing we can do.
  if (newInput == input
      && inputMapping.isIdentity()
      && fieldsUsed.cardinality() == fieldCount) {
    return result(sortExchange, Mappings.createIdentity(fieldCount));
  }

  relBuilder.push(newInput);
  RelCollation newCollation = RexUtil.apply(inputMapping, collation);
  RelDistribution newDistribution = distribution.apply(inputMapping);
  relBuilder.sortExchange(newDistribution, newCollation);

  return result(relBuilder.build(), inputMapping);
}
 

/**
 * Returns a bit set describing the inputs used by a collection of
 * project expressions and an optional condition.
 */
public static ImmutableBitSet bits(List<RexNode> exprs, RexNode expr) {
  final InputFinder inputFinder = new InputFinder();
  RexUtil.apply(inputFinder, exprs, expr);
  return inputFinder.build();
}
 

@Override protected UnifyResult apply(UnifyRuleCall call) {
  final MutableAggregate query = (MutableAggregate) call.query;
  final MutableCalc qInput = (MutableCalc) query.getInput();
  final Pair<RexNode, List<RexNode>> qInputExplained = explainCalc(qInput);
  final RexNode qInputCond = qInputExplained.left;
  final List<RexNode> qInputProjs = qInputExplained.right;

  final MutableAggregate target = (MutableAggregate) call.target;

  final RexBuilder rexBuilder = call.getCluster().getRexBuilder();

  final Mappings.TargetMapping mapping =
      Project.getMapping(fieldCnt(qInput.getInput()), qInputProjs);
  if (mapping == null) {
    return null;
  }

  if (!qInputCond.isAlwaysTrue()) {
    try {
      // Fail the matching when filtering condition references
      // non-grouping columns in target.
      qInputCond.accept(new RexVisitorImpl<Void>(true) {
        @Override public Void visitInputRef(RexInputRef inputRef) {
          if (!target.groupSets.stream()
              .allMatch(groupSet -> groupSet.get(inputRef.getIndex()))) {
            throw Util.FoundOne.NULL;
          }
          return super.visitInputRef(inputRef);
        }
      });
    } catch (Util.FoundOne one) {
      return null;
    }
  }

  final Mapping inverseMapping = mapping.inverse();
  final MutableAggregate aggregate2 =
      permute(query, qInput.getInput(), inverseMapping);

  final Mappings.TargetMapping mappingForQueryCond = Mappings.target(
      target.groupSet::indexOf,
      target.getInput().rowType.getFieldCount(),
      target.groupSet.cardinality());
  final RexNode targetCond = RexUtil.apply(mappingForQueryCond, qInputCond);

  final MutableRel unifiedAggregate =
      unifyAggregates(aggregate2, targetCond, target);
  if (unifiedAggregate == null) {
    return null;
  }
  // Add Project if the mapping breaks order of fields in GroupSet
  if (!Mappings.keepsOrdering(mapping)) {
    final List<Integer> posList = new ArrayList<>();
    final int fieldCount = aggregate2.rowType.getFieldCount();
    final List<Pair<Integer, Integer>> pairs = new ArrayList<>();
    final List<Integer> groupings = aggregate2.groupSet.toList();
    for (int i = 0; i < groupings.size(); i++) {
      pairs.add(Pair.of(mapping.getTarget(groupings.get(i)), i));
    }
    Collections.sort(pairs);
    pairs.forEach(pair -> posList.add(pair.right));
    for (int i = posList.size(); i < fieldCount; i++) {
      posList.add(i);
    }
    final List<RexNode> compenProjs =
        MutableRels.createProjectExprs(unifiedAggregate, posList);
    final RexProgram compensatingRexProgram = RexProgram.create(
        unifiedAggregate.rowType, compenProjs, null,
        query.rowType, rexBuilder);
    final MutableCalc compenCalc =
        MutableCalc.of(unifiedAggregate, compensatingRexProgram);
    if (unifiedAggregate instanceof MutableCalc) {
      final MutableCalc newCompenCalc =
          mergeCalc(rexBuilder, compenCalc, (MutableCalc) unifiedAggregate);
      return tryMergeParentCalcAndGenResult(call, newCompenCalc);
    } else {
      return tryMergeParentCalcAndGenResult(call, compenCalc);
    }
  } else {
    return tryMergeParentCalcAndGenResult(call, unifiedAggregate);
  }
}
 

/**
 * Applies mapping to a given collation.
 *
 * If mapping destroys the collation prefix, this method returns an empty collation.
 * Examples of applying mappings to collation [0, 1]:
 * <ul>
 *   <li>mapping(0, 1) =&gt; [0, 1]</li>
 *   <li>mapping(1, 0) =&gt; [1, 0]</li>
 *   <li>mapping(0) =&gt; [0]</li>
 *   <li>mapping(1) =&gt; []</li>
 *   <li>mapping(2, 0) =&gt; [1]</li>
 *   <li>mapping(2, 1, 0) =&gt; [2, 1]</li>
 *   <li>mapping(2, 1) =&gt; []</li>
 * </ul>
 *
 * @param mapping   Mapping
 * @return Collation with applied mapping.
 */
@Override public RelCollationImpl apply(
    final Mappings.TargetMapping mapping) {
  return (RelCollationImpl) RexUtil.apply(mapping, this);
}