Java Code Examples for org.apache.calcite.plan.RelOptUtil#InputFinder

/**
 * Add the Filter condition to the list obtained from the input.
 */
public RelOptPredicateList getAllPredicates(Filter filter, RelMetadataQuery mq) {
  final RelNode input = filter.getInput();
  final RexBuilder rexBuilder = filter.getCluster().getRexBuilder();
  final RexNode pred = filter.getCondition();

  final RelOptPredicateList predsBelow = mq.getAllPredicates(input);
  if (predsBelow == null) {
    // Safety check
    return null;
  }

  // Extract input fields referenced by Filter condition
  final Set<RelDataTypeField> inputExtraFields = new LinkedHashSet<>();
  final RelOptUtil.InputFinder inputFinder = new RelOptUtil.InputFinder(inputExtraFields);
  pred.accept(inputFinder);
  final ImmutableBitSet inputFieldsUsed = inputFinder.inputBitSet.build();

  // Infer column origin expressions for given references
  final Map<RexInputRef, Set<RexNode>> mapping = new LinkedHashMap<>();
  for (int idx : inputFieldsUsed) {
    final RexInputRef ref = RexInputRef.of(idx, filter.getRowType().getFieldList());
    final Set<RexNode> originalExprs = mq.getExpressionLineage(filter, ref);
    if (originalExprs == null) {
      // Bail out
      return null;
    }
    mapping.put(ref, originalExprs);
  }

  // Replace with new expressions and return union of predicates
  final Set<RexNode> allExprs =
      RelMdExpressionLineage.createAllPossibleExpressions(rexBuilder, pred, mapping);
  if (allExprs == null) {
    return null;
  }
  return predsBelow.union(rexBuilder, RelOptPredicateList.of(rexBuilder, allExprs));
}
 

/**
 * Variant of {@link #trimFields(RelNode, ImmutableBitSet, Set)} for
 * {@link org.apache.calcite.rel.logical.LogicalFilter}.
 */
public TrimResult trimFields(Filter filter, ImmutableBitSet fieldsUsed, Set<RelDataTypeField> extraFields) {
    final RelDataType rowType = filter.getRowType();
    final int fieldCount = rowType.getFieldCount();
    final RexNode conditionExpr = filter.getCondition();
    final RelNode input = filter.getInput();

    // We use the fields used by the consumer, plus any fields used in the
    // filter.
    final Set<RelDataTypeField> inputExtraFields = new LinkedHashSet<>(extraFields);
    RelOptUtil.InputFinder inputFinder = new RelOptUtil.InputFinder(inputExtraFields);
    inputFinder.inputBitSet.addAll(fieldsUsed);
    conditionExpr.accept(inputFinder);
    final ImmutableBitSet inputFieldsUsed = inputFinder.inputBitSet.build();

    // Create input with trimmed columns.
    TrimResult trimResult = trimChild(filter, input, inputFieldsUsed, 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 && fieldsUsed.cardinality() == fieldCount) {
        return result(filter, Mappings.createIdentity(fieldCount));
    }

    // Build new project expressions, and populate the mapping.
    final RexVisitor<RexNode> shuttle = new RexPermuteInputsShuttle(inputMapping, newInput);
    RexNode newConditionExpr = conditionExpr.accept(shuttle);

    // Use copy rather than relBuilder so that correlating variables get set.
    relBuilder.push(filter.copy(filter.getTraitSet(), newInput, newConditionExpr));

    // The result has the same mapping as the input gave us. Sometimes we
    // return fields that the consumer didn't ask for, because the filter
    // needs them for its condition.
    return result(relBuilder.build(), inputMapping);
}
 

/**
 * Add the Filter condition to the list obtained from the input.
 */
public RelOptPredicateList getAllPredicates(Filter filter, RelMetadataQuery mq) {
  final RelNode input = filter.getInput();
  final RexBuilder rexBuilder = filter.getCluster().getRexBuilder();
  final RexNode pred = filter.getCondition();

  final RelOptPredicateList predsBelow = mq.getAllPredicates(input);
  if (predsBelow == null) {
    // Safety check
    return null;
  }

  // Extract input fields referenced by Filter condition
  final Set<RelDataTypeField> inputExtraFields = new LinkedHashSet<>();
  final RelOptUtil.InputFinder inputFinder = new RelOptUtil.InputFinder(inputExtraFields);
  pred.accept(inputFinder);
  final ImmutableBitSet inputFieldsUsed = inputFinder.build();

  // Infer column origin expressions for given references
  final Map<RexInputRef, Set<RexNode>> mapping = new LinkedHashMap<>();
  for (int idx : inputFieldsUsed) {
    final RexInputRef ref = RexInputRef.of(idx, filter.getRowType().getFieldList());
    final Set<RexNode> originalExprs = mq.getExpressionLineage(filter, ref);
    if (originalExprs == null) {
      // Bail out
      return null;
    }
    mapping.put(ref, originalExprs);
  }

  // Replace with new expressions and return union of predicates
  final Set<RexNode> allExprs =
      RelMdExpressionLineage.createAllPossibleExpressions(rexBuilder, pred, mapping);
  if (allExprs == null) {
    return null;
  }
  return predsBelow.union(rexBuilder, RelOptPredicateList.of(rexBuilder, allExprs));
}
 

private static ImmutableBitSet extractInputRefs(RexNode expr) {
  final Set<RelDataTypeField> inputExtraFields = new LinkedHashSet<>();
  final RelOptUtil.InputFinder inputFinder = new RelOptUtil.InputFinder(inputExtraFields);
  expr.accept(inputFinder);
  return inputFinder.inputBitSet.build();
}
 

/**
 * Add the Join condition to the list obtained from the input.
 */
public RelOptPredicateList getAllPredicates(Join join, RelMetadataQuery mq) {
  if (join.getJoinType() != JoinRelType.INNER) {
    // We cannot map origin of this expression.
    return null;
  }

  final RexBuilder rexBuilder = join.getCluster().getRexBuilder();
  final RexNode pred = join.getCondition();

  final Multimap<List<String>, RelTableRef> qualifiedNamesToRefs = HashMultimap.create();
  RelOptPredicateList newPreds = RelOptPredicateList.EMPTY;
  for (RelNode input : join.getInputs()) {
    final RelOptPredicateList inputPreds = mq.getAllPredicates(input);
    if (inputPreds == null) {
      // Bail out
      return null;
    }
    // Gather table references
    final Set<RelTableRef> tableRefs = mq.getTableReferences(input);
    if (input == join.getLeft()) {
      // Left input references remain unchanged
      for (RelTableRef leftRef : tableRefs) {
        qualifiedNamesToRefs.put(leftRef.getQualifiedName(), leftRef);
      }
      newPreds = newPreds.union(rexBuilder, inputPreds);
    } else {
      // Right input references might need to be updated if there are table name
      // clashes with left input
      final Map<RelTableRef, RelTableRef> currentTablesMapping = new HashMap<>();
      for (RelTableRef rightRef : tableRefs) {
        int shift = 0;
        Collection<RelTableRef> lRefs = qualifiedNamesToRefs.get(
            rightRef.getQualifiedName());
        if (lRefs != null) {
          shift = lRefs.size();
        }
        currentTablesMapping.put(rightRef,
            RelTableRef.of(rightRef.getTable(), shift + rightRef.getEntityNumber()));
      }
      final List<RexNode> updatedPreds = Lists.newArrayList(
          Iterables.transform(inputPreds.pulledUpPredicates,
              e -> RexUtil.swapTableReferences(rexBuilder, e,
                  currentTablesMapping)));
      newPreds = newPreds.union(rexBuilder,
          RelOptPredicateList.of(rexBuilder, updatedPreds));
    }
  }

  // Extract input fields referenced by Join condition
  final Set<RelDataTypeField> inputExtraFields = new LinkedHashSet<>();
  final RelOptUtil.InputFinder inputFinder = new RelOptUtil.InputFinder(inputExtraFields);
  pred.accept(inputFinder);
  final ImmutableBitSet inputFieldsUsed = inputFinder.inputBitSet.build();

  // Infer column origin expressions for given references
  final Map<RexInputRef, Set<RexNode>> mapping = new LinkedHashMap<>();
  for (int idx : inputFieldsUsed) {
    final RexInputRef inputRef = RexInputRef.of(idx, join.getRowType().getFieldList());
    final Set<RexNode> originalExprs = mq.getExpressionLineage(join, inputRef);
    if (originalExprs == null) {
      // Bail out
      return null;
    }
    final RexInputRef ref = RexInputRef.of(idx, join.getRowType().getFieldList());
    mapping.put(ref, originalExprs);
  }

  // Replace with new expressions and return union of predicates
  final Set<RexNode> allExprs =
      RelMdExpressionLineage.createAllPossibleExpressions(rexBuilder, pred, mapping);
  if (allExprs == null) {
    return null;
  }
  return newPreds.union(rexBuilder, RelOptPredicateList.of(rexBuilder, allExprs));
}
 

private ImmutableBitSet fieldBitmap(RexNode joinFilter) {
  final RelOptUtil.InputFinder inputFinder = new RelOptUtil.InputFinder();
  joinFilter.accept(inputFinder);
  return inputFinder.inputBitSet.build();
}
 

private static ImmutableBitSet extractInputRefs(RexNode expr) {
  final Set<RelDataTypeField> inputExtraFields = new LinkedHashSet<>();
  final RelOptUtil.InputFinder inputFinder = new RelOptUtil.InputFinder(inputExtraFields);
  expr.accept(inputFinder);
  return inputFinder.build();
}
 

/**
 * Add the Join condition to the list obtained from the input.
 */
public RelOptPredicateList getAllPredicates(Join join, RelMetadataQuery mq) {
  if (join.getJoinType().isOuterJoin()) {
    // We cannot map origin of this expression.
    return null;
  }

  final RexBuilder rexBuilder = join.getCluster().getRexBuilder();
  final RexNode pred = join.getCondition();

  final Multimap<List<String>, RelTableRef> qualifiedNamesToRefs = HashMultimap.create();
  RelOptPredicateList newPreds = RelOptPredicateList.EMPTY;
  for (RelNode input : join.getInputs()) {
    final RelOptPredicateList inputPreds = mq.getAllPredicates(input);
    if (inputPreds == null) {
      // Bail out
      return null;
    }
    // Gather table references
    final Set<RelTableRef> tableRefs = mq.getTableReferences(input);
    if (input == join.getLeft()) {
      // Left input references remain unchanged
      for (RelTableRef leftRef : tableRefs) {
        qualifiedNamesToRefs.put(leftRef.getQualifiedName(), leftRef);
      }
      newPreds = newPreds.union(rexBuilder, inputPreds);
    } else {
      // Right input references might need to be updated if there are table name
      // clashes with left input
      final Map<RelTableRef, RelTableRef> currentTablesMapping = new HashMap<>();
      for (RelTableRef rightRef : tableRefs) {
        int shift = 0;
        Collection<RelTableRef> lRefs = qualifiedNamesToRefs.get(
            rightRef.getQualifiedName());
        if (lRefs != null) {
          shift = lRefs.size();
        }
        currentTablesMapping.put(rightRef,
            RelTableRef.of(rightRef.getTable(), shift + rightRef.getEntityNumber()));
      }
      final List<RexNode> updatedPreds = Lists.newArrayList(
          Iterables.transform(inputPreds.pulledUpPredicates,
              e -> RexUtil.swapTableReferences(rexBuilder, e,
                  currentTablesMapping)));
      newPreds = newPreds.union(rexBuilder,
          RelOptPredicateList.of(rexBuilder, updatedPreds));
    }
  }

  // Extract input fields referenced by Join condition
  final Set<RelDataTypeField> inputExtraFields = new LinkedHashSet<>();
  final RelOptUtil.InputFinder inputFinder = new RelOptUtil.InputFinder(inputExtraFields);
  pred.accept(inputFinder);
  final ImmutableBitSet inputFieldsUsed = inputFinder.build();

  // Infer column origin expressions for given references
  final Map<RexInputRef, Set<RexNode>> mapping = new LinkedHashMap<>();
  final RelDataType fullRowType = SqlValidatorUtil.createJoinType(
      rexBuilder.getTypeFactory(),
      join.getLeft().getRowType(),
      join.getRight().getRowType(),
      null,
      ImmutableList.of());
  for (int idx : inputFieldsUsed) {
    final RexInputRef inputRef = RexInputRef.of(idx, fullRowType.getFieldList());
    final Set<RexNode> originalExprs = mq.getExpressionLineage(join, inputRef);
    if (originalExprs == null) {
      // Bail out
      return null;
    }
    final RexInputRef ref = RexInputRef.of(idx, fullRowType.getFieldList());
    mapping.put(ref, originalExprs);
  }

  // Replace with new expressions and return union of predicates
  final Set<RexNode> allExprs =
      RelMdExpressionLineage.createAllPossibleExpressions(rexBuilder, pred, mapping);
  if (allExprs == null) {
    return null;
  }
  return newPreds.union(rexBuilder, RelOptPredicateList.of(rexBuilder, allExprs));
}
 

private ImmutableBitSet fieldBitmap(RexNode joinFilter) {
  final RelOptUtil.InputFinder inputFinder = new RelOptUtil.InputFinder();
  joinFilter.accept(inputFinder);
  return inputFinder.build();
}
 

public void onMatch(RelOptRuleCall call) {
  final Project origProject;
  final Filter filter;
  if (call.rels.length >= 2) {
    origProject = call.rel(0);
    filter = call.rel(1);
  } else {
    origProject = null;
    filter = call.rel(0);
  }
  final RelNode input = filter.getInput();
  final RexNode origFilter = filter.getCondition();

  if ((origProject != null)
      && RexOver.containsOver(origProject.getProjects(), null)) {
    // Cannot push project through filter if project contains a windowed
    // aggregate -- it will affect row counts. Abort this rule
    // invocation; pushdown will be considered after the windowed
    // aggregate has been implemented. It's OK if the filter contains a
    // windowed aggregate.
    return;
  }

  if ((origProject != null)
      && origProject.getRowType().isStruct()
      && origProject.getRowType().getFieldList().stream()
        .anyMatch(RelDataTypeField::isDynamicStar)) {
    // The PushProjector would change the plan:
    //
    //    prj(**=[$0])
    //    : - filter
    //        : - scan
    //
    // to form like:
    //
    //    prj(**=[$0])                    (1)
    //    : - filter                      (2)
    //        : - prj(**=[$0], ITEM= ...) (3)
    //            :  - scan
    // This new plan has more cost that the old one, because of the new
    // redundant project (3), if we also have FilterProjectTransposeRule in
    // the rule set, it will also trigger infinite match of the ProjectMergeRule
    // for project (1) and (3).
    return;
  }

  final RelBuilder builder = call.builder();
  final RelNode topProject;
  if (origProject != null && (wholeProject || wholeFilter)) {
    builder.push(input);

    final Set<RexNode> set = new LinkedHashSet<>();
    final RelOptUtil.InputFinder refCollector = new RelOptUtil.InputFinder();

    if (wholeFilter) {
      set.add(filter.getCondition());
    } else {
      filter.getCondition().accept(refCollector);
    }
    if (wholeProject) {
      set.addAll(origProject.getProjects());
    } else {
      refCollector.visitEach(origProject.getProjects());
    }

    // Build a list with inputRefs, in order, first, then other expressions.
    final List<RexNode> list = new ArrayList<>();
    final ImmutableBitSet refs = refCollector.build();
    for (RexNode field : builder.fields()) {
      if (refs.get(((RexInputRef) field).getIndex()) || set.contains(field)) {
        list.add(field);
      }
    }
    set.removeAll(list);
    list.addAll(set);
    builder.project(list);
    final Replacer replacer = new Replacer(list, builder);
    builder.filter(replacer.visit(filter.getCondition()));
    builder.project(replacer.visitList(origProject.getProjects()),
        origProject.getRowType().getFieldNames());
    topProject = builder.build();
  } else {
    // The traditional mode of operation of this rule: push down field
    // references. The effect is similar to RelFieldTrimmer.
    final PushProjector pushProjector =
        new PushProjector(origProject, origFilter, input,
            preserveExprCondition, builder);
    topProject = pushProjector.convertProject(null);
  }

  if (topProject != null) {
    call.transformTo(topProject);
  }
}
 

/**
 * Variant of {@link #trimFields(RelNode, ImmutableBitSet, Set)} for
 * {@link org.apache.calcite.rel.logical.LogicalFilter}.
 */
public TrimResult trimFields(
    Filter filter,
    ImmutableBitSet fieldsUsed,
    Set<RelDataTypeField> extraFields) {
  final RelDataType rowType = filter.getRowType();
  final int fieldCount = rowType.getFieldCount();
  final RexNode conditionExpr = filter.getCondition();
  final RelNode input = filter.getInput();

  // We use the fields used by the consumer, plus any fields used in the
  // filter.
  final Set<RelDataTypeField> inputExtraFields =
      new LinkedHashSet<>(extraFields);
  RelOptUtil.InputFinder inputFinder =
      new RelOptUtil.InputFinder(inputExtraFields, fieldsUsed);
  conditionExpr.accept(inputFinder);
  final ImmutableBitSet inputFieldsUsed = inputFinder.build();

  // Create input with trimmed columns.
  TrimResult trimResult =
      trimChild(filter, input, inputFieldsUsed, 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
      && fieldsUsed.cardinality() == fieldCount) {
    return result(filter, Mappings.createIdentity(fieldCount));
  }

  // Build new project expressions, and populate the mapping.
  final RexVisitor<RexNode> shuttle =
      new RexPermuteInputsShuttle(inputMapping, newInput);
  RexNode newConditionExpr =
      conditionExpr.accept(shuttle);

  // Build new filter with trimmed input and condition.
  relBuilder.push(newInput)
      .filter(filter.getVariablesSet(), newConditionExpr);

  // The result has the same mapping as the input gave us. Sometimes we
  // return fields that the consumer didn't ask for, because the filter
  // needs them for its condition.
  return result(relBuilder.build(), inputMapping);
}