org.apache.calcite.util.BitSets Java Examples

@Override
public void populatePartitionVectors(ValueVector[] vectors, List<PartitionLocation> partitions,
                                     BitSet partitionColumnBitSet, Map<Integer, String> fieldNameMap) {
  int record = 0;
  for (PartitionLocation partitionLocation: partitions) {
    for (int partitionColumnIndex : BitSets.toIter(partitionColumnBitSet)) {
      if (partitionLocation.getPartitionValue(partitionColumnIndex) == null) {
        // set null if dirX does not exist for the location.
        ((NullableVarCharVector) vectors[partitionColumnIndex]).getMutator().setNull(record);
      } else {
        byte[] bytes = (partitionLocation.getPartitionValue(partitionColumnIndex)).getBytes(Charsets.UTF_8);
        ((NullableVarCharVector) vectors[partitionColumnIndex]).getMutator().setSafe(record, bytes, 0, bytes.length);
      }
    }
    record++;
  }

  for (ValueVector v : vectors) {
    if (v == null) {
      continue;
    }
    v.getMutator().setValueCount(partitions.size());
  }
}
 

@Override
public void populatePartitionVectors(ValueVector[] vectors, List<PartitionLocation> partitions,
                                     BitSet partitionColumnBitSet, Map<Integer, String> fieldNameMap) {
  int record = 0;
  for (PartitionLocation partitionLocation: partitions) {
    for (int partitionColumnIndex : BitSets.toIter(partitionColumnBitSet)) {
      SchemaPath column = SchemaPath.getSimplePath(fieldNameMap.get(partitionColumnIndex));
      populatePruningVector(vectors[partitionColumnIndex], record, column, partitionLocation.getEntirePartitionLocation());
    }
    record++;
  }

  for (ValueVector v : vectors) {
    if (v == null) {
      continue;
    }
    v.getMutator().setValueCount(partitions.size());
  }

}
 

@Override
public LogicalOperator implement(DrillImplementor implementor) {

  GroupingAggregate.Builder builder = GroupingAggregate.builder();
  builder.setInput(implementor.visitChild(this, 0, getInput()));
  final List<String> childFields = getInput().getRowType().getFieldNames();
  final List<String> fields = getRowType().getFieldNames();

  for (int group : BitSets.toIter(groupSet)) {
    FieldReference fr = new FieldReference(childFields.get(group), ExpressionPosition.UNKNOWN);
    builder.addKey(fr, fr);
  }

  for (Ord<AggregateCall> aggCall : Ord.zip(aggCalls)) {
    FieldReference ref = new FieldReference(fields.get(groupSet.cardinality() + aggCall.i));
    LogicalExpression expr = toDrill(aggCall.e, childFields, implementor);
    builder.addExpr(ref, expr);
  }

  return builder.build();
}
 

/** Compose the array of partition values for the directories that are referenced by filter:
 *  e.g suppose the dir hierarchy is year/quarter/month and the query is:
 *     SELECT * FROM T WHERE dir0=2015 AND dir1 = 'Q1',
 * then for 2015/Q1/Feb, this will have ['2015', 'Q1', null]
 * If the query filter condition is WHERE dir1 = 'Q2'  (i.e no dir0 condition) then the array will
 * have [null, 'Q2', null]
 */
private Pair<String[], Integer> composePartition(BitSet referencedDirsBitSet,
    Map<Integer, Integer> partitionMap,
    ValueVector[] vectors,
    int recordCount) {
  String[] partition = new String[vectors.length];
  int maxIndex = -1;
  for (int referencedDirsIndex : BitSets.toIter(referencedDirsBitSet)) {
    int partitionColumnIndex = partitionMap.get(referencedDirsIndex);
    ValueVector vv = vectors[partitionColumnIndex];
    if (vv.getAccessor().getValueCount() > 0 &&
        vv.getAccessor().getObject(recordCount) != null) {
      String value = vv.getAccessor().getObject(recordCount).toString();
      partition[partitionColumnIndex] = value;
      maxIndex = Math.max(maxIndex, partitionColumnIndex);
    }
  }
  return Pair.of(partition, maxIndex);
}
 

/**
 * Determines how much each input reference needs to be adjusted as a result
 * of projection
 *
 * @return array indicating how much each input needs to be adjusted by
 */
public int[] getAdjustments() {
  int[] adjustments = new int[nChildFields];
  int newIdx = 0;
  int rightOffset = childPreserveExprs.size();
  for (int pos : BitSets.toIter(projRefs)) {
    adjustments[pos] = -(pos - newIdx);
    if (pos >= nSysFields + nFields) {
      adjustments[pos] += rightOffset;
    }
    newIdx++;
  }
  return adjustments;
}
 

private List<DistributionTrait.DistributionField> getDistributionFields(Window.Group window) {
  List<DistributionTrait.DistributionField> groupByFields = Lists.newArrayList();
  for (int group : BitSets.toIter(window.keys)) {
    DistributionTrait.DistributionField field = new DistributionTrait.DistributionField(group);
    groupByFields.add(field);
  }

  return groupByFields;
}
 

/**
 * Create a RelCollation that has partition-by as the leading keys followed by order-by keys
 * @param window The window specification
 * @return a RelCollation with {partition-by keys, order-by keys}
 */
private RelCollation getCollation(Window.Group window) {
  List<RelFieldCollation> fields = Lists.newArrayList();
  for (int group : BitSets.toIter(window.keys)) {
    fields.add(new RelFieldCollation(group));
  }

  fields.addAll(window.orderKeys.getFieldCollations());

  return RelCollations.of(fields);
}
 

public static List<NamedExpression> groupSetToExpr(RelNode input, ImmutableBitSet groupSet) {
  final List<String> childFields = input.getRowType().getFieldNames();
  final List<NamedExpression> keys = Lists.newArrayList();

  for (int group : BitSets.toIter(groupSet)) {
    FieldReference fr = FieldReference.getWithQuotedRef(childFields.get(group));
    keys.add(new NamedExpression(fr, fr));
  }
  return keys;
}
 

private List<DrillDistributionTrait.DistributionField> getDistributionFields(Window.Group window) {
    List<DrillDistributionTrait.DistributionField> groupByFields = Lists.newArrayList();
    for (int group : BitSets.toIter(window.keys)) {
        DrillDistributionTrait.DistributionField field = new DrillDistributionTrait.DistributionField(group);
        groupByFields.add(field);
    }

    return groupByFields;
}
 

/**
 * Create a RelCollation that has partition-by as the leading keys followed by order-by keys
 * @param window The window specification
 * @return a RelCollation with {partition-by keys, order-by keys}
 */
private RelCollation getCollation(Window.Group window) {
    List<RelFieldCollation> fields = Lists.newArrayList();
    for (int group : BitSets.toIter(window.keys)) {
        fields.add(new RelFieldCollation(group));
    }

    fields.addAll(window.orderKeys.getFieldCollations());

    return RelCollations.of(fields);
}
 

/**
 * Determines how much each input reference needs to be adjusted as a result
 * of projection
 *
 * @return array indicating how much each input needs to be adjusted by
 */
public int[] getAdjustments() {
  int[] adjustments = new int[nChildFields];
  int newIdx = 0;
  int rightOffset = childPreserveExprs.size();
  for (int pos : BitSets.toIter(projRefs)) {
    adjustments[pos] = -(pos - newIdx);
    if (pos >= nSysFields + nFields) {
      adjustments[pos] += rightOffset;
    }
    newIdx++;
  }
  return adjustments;
}
 

private RelCollation getCollation(DrillAggregateRel rel) {

    List<RelFieldCollation> fields = Lists.newArrayList();
    for (int group : BitSets.toIter(rel.getGroupSet())) {
      fields.add(new RelFieldCollation(group));
    }
    return RelCollations.of(fields);
  }
 

@Override
public LogicalOperator implement(DrillImplementor implementor) {
  final LogicalOperator inputOp = implementor.visitChild(this, 0, getInput());
  org.apache.drill.common.logical.data.Window.Builder builder = new org.apache.drill.common.logical.data.Window.Builder();
  final List<String> fields = getRowType().getFieldNames();
  final List<String> childFields = getInput().getRowType().getFieldNames();
  for (Group window : groups) {

    for(RelFieldCollation orderKey : window.orderKeys.getFieldCollations()) {
      builder.addOrdering(new Order.Ordering(orderKey.getDirection(), new FieldReference(fields.get(orderKey.getFieldIndex()))));
    }

    for (int group : BitSets.toIter(window.keys)) {
      FieldReference fr = new FieldReference(childFields.get(group), ExpressionPosition.UNKNOWN);
      builder.addWithin(fr, fr);
    }

    int groupCardinality = window.keys.cardinality();
    for (Ord<AggregateCall> aggCall : Ord.zip(window.getAggregateCalls(this))) {
      FieldReference ref = new FieldReference(fields.get(groupCardinality + aggCall.i));
      LogicalExpression expr = toDrill(aggCall.e, childFields);
      builder.addAggregation(ref, expr);
    }
  }
  builder.setInput(inputOp);
  org.apache.drill.common.logical.data.Window frame = builder.build();
  return frame;
}
 

/**
 * Applies a mapping to a BitSet.
 *
 * <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 BitSet apply(Mapping mapping, BitSet bitSet) {
  final BitSet newBitSet = new BitSet();
  for (int source : BitSets.toIter(bitSet)) {
    final int target = mapping.getTarget(source);
    newBitSet.set(target);
  }
  if (newBitSet.equals(bitSet)) {
    return bitSet;
  }
  return newBitSet;
}
 

/**
 * Applies a mapping to a BitSet.
 *
 * <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 BitSet apply(Mapping mapping, BitSet bitSet) {
  final BitSet newBitSet = new BitSet();
  for (int source : BitSets.toIter(bitSet)) {
    final int target = mapping.getTarget(source);
    newBitSet.set(target);
  }
  if (newBitSet.equals(bitSet)) {
    return bitSet;
  }
  return newBitSet;
}
 

/**
 * Determines the best factor to be added next into a join tree.
 *
 * @param multiJoin join factors being optimized
 * @param factorsToAdd factors to choose from to add next
 * @param factorsAdded factors that have already been added to the join tree
 * @param semiJoinOpt optimal semijoins for each factor
 * @param joinTree join tree constructed thus far
 * @param filtersToAdd remaining filters that need to be added
 *
 * @return index of the best factor to add next
 */
private int getBestNextFactor(
    RelMetadataQuery mq,
    LoptMultiJoin multiJoin,
    BitSet factorsToAdd,
    BitSet factorsAdded,
    LoptSemiJoinOptimizer semiJoinOpt,
    LoptJoinTree joinTree,
    List<RexNode> filtersToAdd) {
  // iterate through the remaining factors and determine the
  // best one to add next
  int nextFactor = -1;
  int bestWeight = 0;
  Double bestCardinality = null;
  int [][] factorWeights = multiJoin.getFactorWeights();
  for (int factor : BitSets.toIter(factorsToAdd)) {
    // if the factor corresponds to a dimension table whose
    // join we can remove, make sure the the corresponding fact
    // table is in the current join tree
    Integer factIdx = multiJoin.getJoinRemovalFactor(factor);
    if (factIdx != null) {
      if (!factorsAdded.get(factIdx)) {
        continue;
      }
    }

    // can't add a null-generating factor if its dependent,
    // non-null generating factors haven't been added yet
    if (multiJoin.isNullGenerating(factor)
        && !BitSets.contains(factorsAdded,
            multiJoin.getOuterJoinFactors(factor))) {
      continue;
    }

    // determine the best weight between the current factor
    // under consideration and the factors that have already
    // been added to the tree
    int dimWeight = 0;
    for (int prevFactor : BitSets.toIter(factorsAdded)) {
      if (factorWeights[prevFactor][factor] > dimWeight) {
        dimWeight = factorWeights[prevFactor][factor];
      }
    }

    // only compute the join cardinality if we know that
    // this factor joins with some part of the current join
    // tree and is potentially better than other factors
    // already considered
    Double cardinality = null;
    if ((dimWeight > 0)
        && ((dimWeight > bestWeight) || (dimWeight == bestWeight))) {
      cardinality =
          computeJoinCardinality(
            mq,
              multiJoin,
              semiJoinOpt,
              joinTree,
              filtersToAdd,
              factor);
    }

    // if two factors have the same weight, pick the one
    // with the higher cardinality join key, relative to
    // the join being considered
    if ((dimWeight > bestWeight)
        || ((dimWeight == bestWeight)
        && ((bestCardinality == null)
        || ((cardinality != null)
        && (cardinality > bestCardinality))))) {
      nextFactor = factor;
      bestWeight = dimWeight;
      bestCardinality = cardinality;
    }
  }

  return nextFactor;
}
 

public void onMatch(RelOptRuleCall call) {
  Project origProj = call.rel(0);
  final Correlate corr = call.rel(1);

  // locate all fields referenced in the projection
  // determine which inputs are referenced in the projection;
  // if all fields are being referenced and there are no
  // special expressions, no point in proceeding any further
  PushProjector pushProject =
      new PushProjector(
          origProj,
          call.builder().literal(true),
          corr,
          preserveExprCondition,
          call.builder());
  if (pushProject.locateAllRefs()) {
    return;
  }

  // create left and right projections, projecting only those
  // fields referenced on each side
  RelNode leftProjRel =
      pushProject.createProjectRefsAndExprs(
          corr.getLeft(),
          true,
          false);
  RelNode rightProjRel =
      pushProject.createProjectRefsAndExprs(
          corr.getRight(),
          true,
          true);

  Map<Integer, Integer> requiredColsMap = new HashMap<>();

  // adjust requiredColumns that reference the projected columns
  int[] adjustments = pushProject.getAdjustments();
  BitSet updatedBits = new BitSet();
  for (Integer col : corr.getRequiredColumns()) {
    int newCol = col + adjustments[col];
    updatedBits.set(newCol);
    requiredColsMap.put(col, newCol);
  }

  RexBuilder rexBuilder = call.builder().getRexBuilder();

  CorrelationId correlationId = corr.getCluster().createCorrel();
  RexCorrelVariable rexCorrel =
      (RexCorrelVariable) rexBuilder.makeCorrel(
          leftProjRel.getRowType(),
          correlationId);

  // updates RexCorrelVariable and sets actual RelDataType for RexFieldAccess
  rightProjRel = rightProjRel.accept(
      new RelNodesExprsHandler(
          new RexFieldAccessReplacer(corr.getCorrelationId(),
              rexCorrel, rexBuilder, requiredColsMap)));

  // create a new correlate with the projected children
  Correlate newCorrRel =
      corr.copy(
          corr.getTraitSet(),
          leftProjRel,
          rightProjRel,
          correlationId,
          ImmutableBitSet.of(BitSets.toIter(updatedBits)),
          corr.getJoinType());

  // put the original project on top of the correlate, converting it to
  // reference the modified projection list
  RelNode topProject =
      pushProject.createNewProject(newCorrRel, adjustments);

  call.transformTo(topProject);
}
 

@Override
public PhysicalOperator getPhysicalOperator(PhysicalPlanCreator creator) throws IOException {
  Prel child = (Prel) this.getInput();

  PhysicalOperator childPOP = child.getPhysicalOperator(creator);

  final List<String> childFields = getInput().getRowType().getFieldNames();

  // We don't support distinct partitions
  checkState(groups.size() == 1, "Only one window is expected in WindowPrel");

  Group window = groups.get(0);
  List<NamedExpression> withins = Lists.newArrayList();
  List<NamedExpression> aggs = Lists.newArrayList();
  List<Order.Ordering> orderings = Lists.newArrayList();

  for (int group : BitSets.toIter(window.keys)) {
    FieldReference fr = new FieldReference(childFields.get(group), ExpressionPosition.UNKNOWN);
    withins.add(new NamedExpression(fr, fr));
  }

  for (AggregateCall aggCall : window.getAggregateCalls(this)) {
    FieldReference ref = new FieldReference(aggCall.getName());
    LogicalExpression expr = toDrill(aggCall, childFields);
    aggs.add(new NamedExpression(expr, ref));
  }

  for (RelFieldCollation fieldCollation : window.orderKeys.getFieldCollations()) {
    orderings.add(new Order.Ordering(fieldCollation.getDirection(), new FieldReference(childFields.get(fieldCollation.getFieldIndex())), fieldCollation.nullDirection));
  }

  WindowPOP windowPOP = new WindowPOP(
      childPOP,
      withins,
      aggs,
      orderings,
      window.isRows,
      WindowPOP.newBound(window.lowerBound),
      WindowPOP.newBound(window.upperBound));

  creator.addMetadata(this, windowPOP);
  return windowPOP;
}
 

@Override
public void onMatch(RelOptRuleCall call) {
    Project origProj = call.rel(0);
    final Correlate corr = call.rel(1);

    // locate all fields referenced in the projection
    // determine which inputs are referenced in the projection;
    // if all fields are being referenced and there are no
    // special expressions, no point in proceeding any further
    PushProjector pushProject = new PushProjector(origProj, call.builder().literal(true), corr,
            preserveExprCondition, call.builder());
    if (pushProject.locateAllRefs()) {
        return;
    }

    // create left and right projections, projecting only those
    // fields referenced on each side
    RelNode leftProjRel = pushProject.createProjectRefsAndExprs(corr.getLeft(), true, false);
    RelNode rightProjRel = pushProject.createProjectRefsAndExprs(corr.getRight(), true, true);

    Map<Integer, Integer> requiredColsMap = new HashMap<>();

    // adjust requiredColumns that reference the projected columns
    int[] adjustments = pushProject.getAdjustments();
    BitSet updatedBits = new BitSet();
    for (Integer col : corr.getRequiredColumns()) {
        int newCol = col + adjustments[col];
        updatedBits.set(newCol);
        requiredColsMap.put(col, newCol);
    }

    RexBuilder rexBuilder = call.builder().getRexBuilder();

    CorrelationId correlationId = corr.getCluster().createCorrel();
    RexCorrelVariable rexCorrel = (RexCorrelVariable) rexBuilder.makeCorrel(leftProjRel.getRowType(),
            correlationId);

    // updates RexCorrelVariable and sets actual RelDataType for RexFieldAccess
    rightProjRel = rightProjRel.accept(new RelNodesExprsHandler(
            new RexFieldAccessReplacer(corr.getCorrelationId(), rexCorrel, rexBuilder, requiredColsMap)));

    // create a new correlate with the projected children
    Correlate newCorrRel = corr.copy(corr.getTraitSet(), leftProjRel, rightProjRel, correlationId,
            ImmutableBitSet.of(BitSets.toIter(updatedBits)), corr.getJoinType());

    // put the original project on top of the correlate, converting it to
    // reference the modified projection list
    RelNode topProject = pushProject.createNewProject(newCorrRel, adjustments);

    call.transformTo(topProject);
}
 

/**
 * Determines the best factor to be added next into a join tree.
 *
 * @param multiJoin join factors being optimized
 * @param factorsToAdd factors to choose from to add next
 * @param factorsAdded factors that have already been added to the join tree
 * @param semiJoinOpt optimal semijoins for each factor
 * @param joinTree join tree constructed thus far
 * @param filtersToAdd remaining filters that need to be added
 *
 * @return index of the best factor to add next
 */
private int getBestNextFactor(
    RelMetadataQuery mq,
    LoptMultiJoin multiJoin,
    BitSet factorsToAdd,
    BitSet factorsAdded,
    LoptSemiJoinOptimizer semiJoinOpt,
    LoptJoinTree joinTree,
    List<RexNode> filtersToAdd) {
  // iterate through the remaining factors and determine the
  // best one to add next
  int nextFactor = -1;
  int bestWeight = 0;
  Double bestCardinality = null;
  int [][] factorWeights = multiJoin.getFactorWeights();
  for (int factor : BitSets.toIter(factorsToAdd)) {
    // if the factor corresponds to a dimension table whose
    // join we can remove, make sure the the corresponding fact
    // table is in the current join tree
    Integer factIdx = multiJoin.getJoinRemovalFactor(factor);
    if (factIdx != null) {
      if (!factorsAdded.get(factIdx)) {
        continue;
      }
    }

    // can't add a null-generating factor if its dependent,
    // non-null generating factors haven't been added yet
    if (multiJoin.isNullGenerating(factor)
        && !BitSets.contains(factorsAdded,
            multiJoin.getOuterJoinFactors(factor))) {
      continue;
    }

    // determine the best weight between the current factor
    // under consideration and the factors that have already
    // been added to the tree
    int dimWeight = 0;
    for (int prevFactor : BitSets.toIter(factorsAdded)) {
      if (factorWeights[prevFactor][factor] > dimWeight) {
        dimWeight = factorWeights[prevFactor][factor];
      }
    }

    // only compute the join cardinality if we know that
    // this factor joins with some part of the current join
    // tree and is potentially better than other factors
    // already considered
    Double cardinality = null;
    if ((dimWeight > 0)
        && ((dimWeight > bestWeight) || (dimWeight == bestWeight))) {
      cardinality =
          computeJoinCardinality(
            mq,
              multiJoin,
              semiJoinOpt,
              joinTree,
              filtersToAdd,
              factor);
    }

    // if two factors have the same weight, pick the one
    // with the higher cardinality join key, relative to
    // the join being considered
    if ((dimWeight > bestWeight)
        || ((dimWeight == bestWeight)
        && ((bestCardinality == null)
        || ((cardinality != null)
        && (cardinality > bestCardinality))))) {
      nextFactor = factor;
      bestWeight = dimWeight;
      bestCardinality = cardinality;
    }
  }

  return nextFactor;
}
 

/**
 * Returns true if a join tree contains all factors required
 *
 * @param joinTree join tree to be examined
 * @param factorsNeeded bitmap of factors required
 *
 * @return true if join tree contains all required factors
 */
public boolean hasAllFactors(
    LoptJoinTree joinTree,
    BitSet factorsNeeded) {
  return BitSets.contains(BitSets.of(joinTree.getTreeOrder()), factorsNeeded);
}
 

/**
 * Returns true if a join tree contains all factors required
 *
 * @param joinTree join tree to be examined
 * @param factorsNeeded bitmap of factors required
 *
 * @return true if join tree contains all required factors
 */
public boolean hasAllFactors(
    LoptJoinTree joinTree,
    BitSet factorsNeeded) {
  return BitSets.contains(BitSets.of(joinTree.getTreeOrder()), factorsNeeded);
}