Java Code Examples for org.apache.calcite.util.ImmutableBitSet#Builder

/**
 * Computes the population size for a set of keys returned from a join
 *
 * @param joinRel  the join rel
 * @param groupKey keys to compute the population for
 * @return computed population size
 */
public static Double getJoinPopulationSize(RelMetadataQuery mq,
    RelNode joinRel, ImmutableBitSet groupKey) {
  Join join = (Join) joinRel;
  if (!join.getJoinType().projectsRight()) {
    return mq.getPopulationSize(join.getLeft(), groupKey);
  }
  ImmutableBitSet.Builder leftMask = ImmutableBitSet.builder();
  ImmutableBitSet.Builder rightMask = ImmutableBitSet.builder();
  RelNode left = joinRel.getInputs().get(0);
  RelNode right = joinRel.getInputs().get(1);

  // separate the mask into masks for the left and right
  RelMdUtil.setLeftRightBitmaps(
      groupKey, leftMask, rightMask, left.getRowType().getFieldCount());

  Double population =
      NumberUtil.multiply(
          mq.getPopulationSize(left, leftMask.build()),
          mq.getPopulationSize(right, rightMask.build()));

  return numDistinctVals(population, mq.getRowCount(joinRel));
}
 

public void rewriteRel(LogicalCorrelate rel) {
  ImmutableBitSet.Builder newPos = ImmutableBitSet.builder();
  for (int pos : rel.getRequiredColumns()) {
    RelDataType corrFieldType =
        rel.getLeft().getRowType().getFieldList().get(pos)
            .getType();
    if (corrFieldType.isStruct()) {
      throw Util.needToImplement("correlation on structured type");
    }
    newPos.set(getNewForOldInput(pos));
  }
  LogicalCorrelate newRel =
      LogicalCorrelate.create(getNewForOldRel(rel.getLeft()),
          getNewForOldRel(rel.getRight()),
          rel.getCorrelationId(),
          newPos.build(),
          rel.getJoinType());
  setNewForOldRel(rel, newRel);
}
 

/**
 * Extract projects from the Aggregate and return the index mapping between the new projects
 * and it's input.
 */
private Mapping extractProjectsAndMapping(
	Aggregate aggregate,
	RelNode input,
	RelBuilder relBuilder) {

	// Compute which input fields are used.
	final ImmutableBitSet.Builder inputFieldsUsed = getInputFieldUsed(aggregate, input);

	final List<RexNode> projects = new ArrayList<>();
	final Mapping mapping =
		Mappings.create(MappingType.INVERSE_SURJECTION,
			aggregate.getInput().getRowType().getFieldCount(),
			inputFieldsUsed.cardinality());
	int j = 0;
	for (int i : inputFieldsUsed.build()) {
		projects.add(relBuilder.field(i));
		mapping.set(i, j++);
	}

	if (input instanceof Project) {
		// this will not create trivial projects
		relBuilder.project(projects);
	} else {
		relBuilder.project(projects, Collections.emptyList(), true);
	}

	return mapping;
}
 

/**
 * Sets the bitmap indicating which factors a filter references based on
 * which fields it references
 *
 * @param fieldRefBitmap bitmap representing fields referenced
 * @return bitmap representing factors referenced that will
 * be set by this method
 */
private ImmutableBitSet factorBitmap(ImmutableBitSet fieldRefBitmap) {
  ImmutableBitSet.Builder factorRefBitmap = ImmutableBitSet.builder();
  for (int field : fieldRefBitmap) {
    int factor = findRef(field);
    factorRefBitmap.set(factor);
  }
  return factorRefBitmap.build();
}
 

/** Splits a column set between left and right sets. */
private static Pair<ImmutableBitSet, ImmutableBitSet>
    splitLeftAndRightColumns(int leftCount, final ImmutableBitSet columns) {
  ImmutableBitSet.Builder leftBuilder = ImmutableBitSet.builder();
  ImmutableBitSet.Builder rightBuilder = ImmutableBitSet.builder();
  for (int bit : columns) {
    if (bit < leftCount) {
      leftBuilder.set(bit);
    } else {
      rightBuilder.set(bit - leftCount);
    }
  }
  return Pair.of(leftBuilder.build(), rightBuilder.build());
}
 

/**
 * Sets the bitmap indicating which factors a filter references based on
 * which fields it references
 *
 * @param fieldRefBitmap bitmap representing fields referenced
 * @return bitmap representing factors referenced that will
 * be set by this method
 */
private ImmutableBitSet factorBitmap(ImmutableBitSet fieldRefBitmap) {
  ImmutableBitSet.Builder factorRefBitmap = ImmutableBitSet.builder();
  for (int field : fieldRefBitmap) {
    int factor = findRef(field);
    factorRefBitmap.set(factor);
  }
  return factorRefBitmap.build();
}
 

public static Boolean areColumnsUnique(RelMetadataQuery mq, RelNode rel,
    List<RexInputRef> columnRefs) {
  ImmutableBitSet.Builder colMask = ImmutableBitSet.builder();
  for (RexInputRef columnRef : columnRefs) {
    colMask.set(columnRef.getIndex());
  }
  return mq.areColumnsUnique(rel, colMask.build());
}
 

public static Boolean areColumnsUnique(RelMetadataQuery mq, RelNode rel,
    List<RexInputRef> columnRefs) {
  ImmutableBitSet.Builder colMask = ImmutableBitSet.builder();
  for (RexInputRef columnRef : columnRefs) {
    colMask.set(columnRef.getIndex());
  }
  return mq.areColumnsUnique(rel, colMask.build());
}
 

public Double getDistinctRowCount(Aggregate rel, RelMetadataQuery mq,
    ImmutableBitSet groupKey, RexNode predicate) {
  if (predicate == null || predicate.isAlwaysTrue()) {
    if (groupKey.isEmpty()) {
      return 1D;
    }
  }
  // determine which predicates can be applied on the child of the
  // aggregate
  final List<RexNode> notPushable = new ArrayList<>();
  final List<RexNode> pushable = new ArrayList<>();
  RelOptUtil.splitFilters(
      rel.getGroupSet(),
      predicate,
      pushable,
      notPushable);
  final RexBuilder rexBuilder = rel.getCluster().getRexBuilder();
  RexNode childPreds =
      RexUtil.composeConjunction(rexBuilder, pushable, true);

  // set the bits as they correspond to the child input
  ImmutableBitSet.Builder childKey = ImmutableBitSet.builder();
  RelMdUtil.setAggChildKeys(groupKey, rel, childKey);

  Double distinctRowCount =
      mq.getDistinctRowCount(rel.getInput(), childKey.build(), childPreds);
  if (distinctRowCount == null) {
    return null;
  } else if (notPushable.isEmpty()) {
    return distinctRowCount;
  } else {
    RexNode preds =
        RexUtil.composeConjunction(rexBuilder, notPushable, true);
    return distinctRowCount * RelMdUtil.guessSelectivity(preds);
  }
}
 

/**
 * Separates a bit-mask representing a join into masks representing the left
 * and right inputs into the join.
 *
 * @param groupKey      original bit-mask
 * @param leftMask      left bit-mask to be set
 * @param rightMask     right bit-mask to be set
 * @param nFieldsOnLeft number of fields in the left input
 */
public static void setLeftRightBitmaps(
    ImmutableBitSet groupKey,
    ImmutableBitSet.Builder leftMask,
    ImmutableBitSet.Builder rightMask,
    int nFieldsOnLeft) {
  for (int bit : groupKey) {
    if (bit < nFieldsOnLeft) {
      leftMask.set(bit);
    } else {
      rightMask.set(bit - nFieldsOnLeft);
    }
  }
}
 

/**
 * Sets a join key if only one of the specified input references corresponds
 * to a specified factor as determined by its field numbers. Also keeps
 * track of the keys from the other factor.
 *
 * @param joinKeys join keys to be set if a key is found
 * @param otherJoinKeys join keys for the other join factor
 * @param ref1 first input reference
 * @param ref2 second input reference
 * @param firstFieldNum first field number of the factor
 * @param lastFieldNum last field number + 1 of the factor
 * @param swap if true, check for the desired input reference in the second
 * input reference parameter if the first input reference isn't the correct
 * one
 */
private void setJoinKey(
    ImmutableBitSet.Builder joinKeys,
    ImmutableBitSet.Builder otherJoinKeys,
    int ref1,
    int ref2,
    int firstFieldNum,
    int lastFieldNum,
    boolean swap) {
  if ((ref1 >= firstFieldNum) && (ref1 < lastFieldNum)) {
    if (!((ref2 >= firstFieldNum) && (ref2 < lastFieldNum))) {
      joinKeys.set(ref1 - firstFieldNum);
      otherJoinKeys.set(ref2);
    }
    return;
  }
  if (swap) {
    setJoinKey(
        joinKeys,
        otherJoinKeys,
        ref2,
        ref1,
        firstFieldNum,
        lastFieldNum,
        false);
  }
}
 

/** Finds which columns of a correlation variable are used within a
 * relational expression. */
public static ImmutableBitSet correlationColumns(CorrelationId id, RelNode rel) {
    final CorrelationCollector collector = new CorrelationCollector();
    rel.accept(collector);
    final ImmutableBitSet.Builder builder = ImmutableBitSet.builder();
    for (int field : collector.vuv.variableFields.get(id)) {
        if (field >= 0) {
            builder.set(field);
        }
    }
    return builder.build();
}
 

public static ImmutableBitSet projectSwap(ImmutableBitSet projectedFields, int numFieldLeft, int size) {
  int[] adjustments = new int[size];
  Arrays.fill(adjustments, 0, numFieldLeft, size-numFieldLeft);
  Arrays.fill(adjustments, numFieldLeft, numFieldLeft + size-numFieldLeft, -numFieldLeft);
  ImmutableBitSet.Builder builder = ImmutableBitSet.builder();
  for (int i : projectedFields.asList()) {
    builder.set(adjustments[i]+i);
  }
  return builder.build();
}
 

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

  final int count = union.getRowType().getFieldCount();
  if (count == 1) {
    // No room for optimization since we cannot create an empty Project
    // operator. If we created a Project with one column, this rule would
    // cycle.
    return;
  }

  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 = ImmutableList.copyOf(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());
}
 

/**
 * Infers predicates for a project.
 *
 * <ol>
 * <li>create a mapping from input to projection. Map only positions that
 * directly reference an input column.
 * <li>Expressions that only contain above columns are retained in the
 * Project's pullExpressions list.
 * <li>For e.g. expression 'a + e = 9' below will not be pulled up because 'e'
 * is not in the projection list.
 *
 * <blockquote><pre>
 * inputPullUpExprs:      {a &gt; 7, b + c &lt; 10, a + e = 9}
 * projectionExprs:       {a, b, c, e / 2}
 * projectionPullupExprs: {a &gt; 7, b + c &lt; 10}
 * </pre></blockquote>
 *
 * </ol>
 */
public RelOptPredicateList getPredicates(Project project,
    RelMetadataQuery mq) {
  final RelNode input = project.getInput();
  final RexBuilder rexBuilder = project.getCluster().getRexBuilder();
  final RelOptPredicateList inputInfo = mq.getPulledUpPredicates(input);
  final List<RexNode> projectPullUpPredicates = new ArrayList<>();

  ImmutableBitSet.Builder columnsMappedBuilder = ImmutableBitSet.builder();
  Mapping m = Mappings.create(MappingType.PARTIAL_FUNCTION,
      input.getRowType().getFieldCount(),
      project.getRowType().getFieldCount());

  for (Ord<RexNode> expr : Ord.zip(project.getProjects())) {
    if (expr.e instanceof RexInputRef) {
      int sIdx = ((RexInputRef) expr.e).getIndex();
      m.set(sIdx, expr.i);
      columnsMappedBuilder.set(sIdx);
    // Project can also generate constants. We need to include them.
    } else if (RexLiteral.isNullLiteral(expr.e)) {
      projectPullUpPredicates.add(
          rexBuilder.makeCall(SqlStdOperatorTable.IS_NULL,
              rexBuilder.makeInputRef(project, expr.i)));
    } else if (RexUtil.isConstant(expr.e)) {
      final List<RexNode> args =
          ImmutableList.of(rexBuilder.makeInputRef(project, expr.i), expr.e);
      final SqlOperator op = args.get(0).getType().isNullable()
          || args.get(1).getType().isNullable()
          ? SqlStdOperatorTable.IS_NOT_DISTINCT_FROM
          : SqlStdOperatorTable.EQUALS;
      projectPullUpPredicates.add(rexBuilder.makeCall(op, args));
    }
  }

  // Go over childPullUpPredicates. If a predicate only contains columns in
  // 'columnsMapped' construct a new predicate based on mapping.
  final ImmutableBitSet columnsMapped = columnsMappedBuilder.build();
  for (RexNode r : inputInfo.pulledUpPredicates) {
    RexNode r2 = projectPredicate(rexBuilder, input, r, columnsMapped);
    if (!r2.isAlwaysTrue()) {
      r2 = r2.accept(new RexPermuteInputsShuttle(m, input));
      projectPullUpPredicates.add(r2);
    }
  }
  return RelOptPredicateList.of(rexBuilder, projectPullUpPredicates);
}
 

@Override public void onMatch(RelOptRuleCall call) {
  final Project project = call.rel(0);
  final Join join = call.rel(1);
  final boolean isLeftJoin = join.getJoinType() == JoinRelType.LEFT;
  int lower = isLeftJoin
      ? join.getLeft().getRowType().getFieldCount() - 1 : 0;
  int upper = isLeftJoin
      ? join.getRowType().getFieldCount()
      : join.getLeft().getRowType().getFieldCount();

  // Check whether the project uses columns whose index is between
  // lower(included) and upper(excluded).
  for (RexNode expr: project.getProjects()) {
    if (RelOptUtil.InputFinder.bits(expr).asList().stream().anyMatch(
        i -> i >= lower && i < upper)) {
      return;
    }
  }

  final List<Integer> leftKeys = new ArrayList<>();
  final List<Integer> rightKeys = new ArrayList<>();
  RelOptUtil.splitJoinCondition(join.getLeft(), join.getRight(),
      join.getCondition(), leftKeys, rightKeys,
      new ArrayList<>());

  final List<Integer> joinKeys = isLeftJoin ? rightKeys : leftKeys;
  final ImmutableBitSet.Builder columns = ImmutableBitSet.builder();
  joinKeys.forEach(key -> columns.set(key));

  final RelMetadataQuery mq = call.getMetadataQuery();
  if (!mq.areColumnsUnique(isLeftJoin ? join.getRight() : join.getLeft(),
      columns.build())) {
    return;
  }

  RelNode node;
  if (isLeftJoin) {
    node = project
        .copy(project.getTraitSet(), join.getLeft(), project.getProjects(),
            project.getRowType());
  } else {
    final int offset = join.getLeft().getRowType().getFieldCount();
    final List<RexNode> newExprs = project.getProjects().stream()
        .map(expr -> RexUtil.shift(expr, -offset))
        .collect(Collectors.toList());
    node = project.copy(project.getTraitSet(), join.getRight(), newExprs,
        project.getRowType());
  }
  call.transformTo(node);
}
 

private Boolean areProjectColumnsUnique(
    SingleRel rel, RelMetadataQuery mq,
    ImmutableBitSet columns, boolean ignoreNulls, List<RexNode> projExprs) {
  RelDataTypeFactory typeFactory = rel.getCluster().getTypeFactory();
  ImmutableBitSet.Builder childColumns = ImmutableBitSet.builder();
  for (int bit : columns) {
    RexNode projExpr = projExprs.get(bit);
    if (projExpr instanceof RexInputRef) {
      childColumns.set(((RexInputRef) projExpr).getIndex());
    } else if (projExpr instanceof RexCall && ignoreNulls) {
      // If the expression is a cast such that the types are the same
      // except for the nullability, then if we're ignoring nulls,
      // it doesn't matter whether the underlying column reference
      // is nullable.  Check that the types are the same by making a
      // nullable copy of both types and then comparing them.
      RexCall call = (RexCall) projExpr;
      if (call.getOperator() != SqlStdOperatorTable.CAST) {
        continue;
      }
      RexNode castOperand = call.getOperands().get(0);
      if (!(castOperand instanceof RexInputRef)) {
        continue;
      }
      RelDataType castType =
          typeFactory.createTypeWithNullability(
              projExpr.getType(), true);
      RelDataType origType = typeFactory.createTypeWithNullability(
          castOperand.getType(),
          true);
      if (castType.equals(origType)) {
        childColumns.set(((RexInputRef) castOperand).getIndex());
      }
    } else {
      // If the expression will not influence uniqueness of the
      // projection, then skip it.
      continue;
    }
  }

  // If no columns can affect uniqueness, then return unknown
  if (childColumns.cardinality() == 0) {
    return null;
  }

  return mq.areColumnsUnique(rel.getInput(), childColumns.build(),
      ignoreNulls);
}
 

public Double getPopulationSize(Aggregate rel, RelMetadataQuery mq,
    ImmutableBitSet groupKey) {
  ImmutableBitSet.Builder childKey = ImmutableBitSet.builder();
  RelMdUtil.setAggChildKeys(groupKey, rel, childKey);
  return mq.getPopulationSize(rel.getInput(), childKey.build());
}
 

/**
 * Determines which join filters can be added to the current join tree. Note
 * that the join filter still reflects the original join ordering. It will
 * only be adjusted to reflect the new join ordering if the "adjust"
 * parameter is set to true.
 *
 * @param multiJoin join factors being optimized
 * @param leftTree left subtree of the join tree
 * @param leftIdx if ≥ 0, only consider filters that reference leftIdx in
 * leftTree; otherwise, consider all filters that reference any factor in
 * leftTree
 * @param rightTree right subtree of the join tree
 * @param filtersToAdd remaining join filters that need to be added; those
 * that are added are removed from the list
 * @param adjust if true, adjust filter to reflect new join ordering
 *
 * @return AND'd expression of the join filters that can be added to the
 * current join tree
 */
private RexNode addFilters(
    LoptMultiJoin multiJoin,
    LoptJoinTree leftTree,
    int leftIdx,
    LoptJoinTree rightTree,
    List<RexNode> filtersToAdd,
    boolean adjust) {
  // loop through the remaining filters to be added and pick out the
  // ones that reference only the factors in the new join tree
  final RexBuilder rexBuilder =
      multiJoin.getMultiJoinRel().getCluster().getRexBuilder();
  final ImmutableBitSet.Builder childFactorBuilder =
      ImmutableBitSet.builder();
  childFactorBuilder.addAll(rightTree.getTreeOrder());
  if (leftIdx >= 0) {
    childFactorBuilder.set(leftIdx);
  } else {
    childFactorBuilder.addAll(leftTree.getTreeOrder());
  }
  for (int child : rightTree.getTreeOrder()) {
    childFactorBuilder.set(child);
  }

  final ImmutableBitSet childFactor = childFactorBuilder.build();
  RexNode condition = null;
  final ListIterator<RexNode> filterIter = filtersToAdd.listIterator();
  while (filterIter.hasNext()) {
    RexNode joinFilter = filterIter.next();
    ImmutableBitSet filterBitmap =
        multiJoin.getFactorsRefByJoinFilter(joinFilter);

    // if all factors in the join filter are in the join tree,
    // AND the filter to the current join condition
    if (childFactor.contains(filterBitmap)) {
      if (condition == null) {
        condition = joinFilter;
      } else {
        condition =
            rexBuilder.makeCall(
                SqlStdOperatorTable.AND,
                condition,
                joinFilter);
      }
      filterIter.remove();
    }
  }

  if (adjust && (condition != null)) {
    int [] adjustments = new int[multiJoin.getNumTotalFields()];
    if (needsAdjustment(
        multiJoin,
        adjustments,
        leftTree,
        rightTree,
        false)) {
      condition =
          condition.accept(
              new RelOptUtil.RexInputConverter(
                  rexBuilder,
                  multiJoin.getMultiJoinFields(),
                  leftTree.getJoinTree().getRowType().getFieldList(),
                  rightTree.getJoinTree().getRowType().getFieldList(),
                  adjustments));
    }
  }

  if (condition == null) {
    condition = rexBuilder.makeLiteral(true);
  }

  return condition;
}
 

/**
 * Applies a mapping to an {@code ImmutableBitSet}.
 *
 * <p>If the mapping does not affect the bit set, returns the original.
 * Never changes the original.
 *
 * @param mapping Mapping
 * @param bitSet  Bit set
 * @return Bit set with mapping applied
 */
public static ImmutableBitSet apply(Mapping mapping, ImmutableBitSet bitSet) {
  final ImmutableBitSet.Builder builder = ImmutableBitSet.builder();
  for (int source : bitSet) {
    final int target = mapping.getTarget(source);
    builder.set(target);
  }
  return builder.build(bitSet);
}