Java Code Examples for org.apache.calcite.util.ImmutableBitSet#get()

@Override
public void onMatch(RelOptRuleCall call) {
  FilterPrel filter = call.rel(0);
  NestedLoopJoinPrel join = call.rel(1);

  if ((join.getProjectedFields() == null) || join.getProjectedFields().cardinality() == join.getInputRowType().getFieldCount()) {
    call.transformTo(NestedLoopJoinPrel.create(join.getCluster(), join.getTraitSet(), join.getLeft(), join.getRight(), join.getJoinType(), RelOptUtil.andJoinFilters(join.getCluster().getRexBuilder(), join.getCondition(), filter.getCondition()), join.getProjectedFields()));
  } else {
    // Current filter condition is written based on projected fields on join. In order to push this filter down we need to rewrite filter condition
    final ImmutableBitSet topProjectedColumns = RelOptUtil.InputFinder.bits(filter.getCondition());
    final ImmutableBitSet bottomProjectedColumns = join.getProjectedFields();

    Mapping mapping = Mappings.create(MappingType.SURJECTION, join.getRowType().getFieldCount(), join.getInputRowType().getFieldCount());
    for (Ord<Integer> ord : Ord.zip(bottomProjectedColumns)) {
      if (topProjectedColumns.get(ord.i)) {
        mapping.set(ord.i, ord.e);
      }
    }

    RexShuttle shuttle = new RexPermuteInputsShuttle(mapping);
    RexNode updatedCondition = shuttle.apply(filter.getCondition());

    call.transformTo(NestedLoopJoinPrel.create(join.getCluster(), join.getTraitSet(), join.getLeft(), join.getRight(), join.getJoinType(), RelOptUtil.andJoinFilters(join.getCluster().getRexBuilder(), join.getCondition(), updatedCondition), join.getProjectedFields()));
  }
}
 

public static RelNode addPartialProjectOnJoin(JoinPrel join, ImmutableBitSet projected) {
  if ((projected != null) && ((projected.cardinality() != 0) || (projected.cardinality() != join.getRowType().getFieldCount()))) {
    List<RelDataTypeField> fields = join.getRowType().getFieldList();
    List<RelDataType> dataTypes = new ArrayList<>();
    List<String> fieldNames = new ArrayList<>();
    List<RexNode> exprs = new ArrayList<>();
    for (int i = 0; i < fields.size(); i++) {
      if (projected.get(i)) {
        RelDataTypeField field = fields.get(i);
        dataTypes.add(field.getType());
        fieldNames.add(field.getName());
        exprs.add(join.getCluster().getRexBuilder().makeInputRef(field.getType(), i));
      }
    }
    RelDataType rowType = join.getCluster().getTypeFactory().createStructType(dataTypes, fieldNames);
    return ProjectPrel.create(join.getCluster(), join.getTraitSet(), join, exprs, rowType);
  }
  return join;
}
 

public static RelDataType rowTypeFromProjected(RelNode left, RelNode right, RelDataType curRowType, ImmutableBitSet projected, RelDataTypeFactory typeFactory) {
  int leftSize = left.getRowType().getFieldCount();
  int rightSize = right.getRowType().getFieldCount();
  List<RelDataTypeField> fields = curRowType.getFieldList();
  if (projected.asSet().size() != fields.size()) {;
    List<RelDataType> dataTypes = new ArrayList<>();
    List<String> fieldNames = new ArrayList<>();
    for (int i = 0; i < fields.size(); i++) {
      if (projected.get(i)) {
        RelDataTypeField field = fields.get(i);
        dataTypes.add(field.getType());
        fieldNames.add(field.getName());
      }
    }
    return typeFactory.createStructType(dataTypes, fieldNames);
  }
  return curRowType;
}
 

private static boolean allContain(List<ImmutableBitSet> groupSets,
    int groupKey) {
  for (ImmutableBitSet groupSet : groupSets) {
    if (!groupSet.get(groupKey)) {
      return false;
    }
  }
  return true;
}
 

private static boolean allContain(List<ImmutableBitSet> bitSets, int bit) {
  for (ImmutableBitSet bitSet : bitSets) {
    if (!bitSet.get(bit)) {
      return false;
    }
  }
  return true;
}
 

/** Returns a checker that consults a bit set to find out whether particular
 * inputs may be null. */
public static Strong of(final ImmutableBitSet nullColumns) {
  return new Strong() {
    @Override public boolean isNull(RexInputRef ref) {
      return nullColumns.get(ref.getIndex());
    }
  };
}
 

private static boolean allContain(List<ImmutableBitSet> groupSets,
    int groupKey) {
  for (ImmutableBitSet groupSet : groupSets) {
    if (!groupSet.get(groupKey)) {
      return false;
    }
  }
  return true;
}
 

private static boolean allContain(List<ImmutableBitSet> bitSets, int bit) {
  for (ImmutableBitSet bitSet : bitSets) {
    if (!bitSet.get(bit)) {
      return false;
    }
  }
  return true;
}
 

/** Returns a checker that consults a bit set to find out whether particular
 * inputs may be null. */
public static Strong of(final ImmutableBitSet nullColumns) {
  return new Strong() {
    @Override public boolean isNull(RexInputRef ref) {
      return nullColumns.get(ref.getIndex());
    }
  };
}
 

/**
 * Locates pairs of joins that are self-joins where the join can be removed
 * because the join condition between the two factors is an equality join on
 * unique keys.
 *
 * @param multiJoin join factors being optimized
 */
private void findRemovableSelfJoins(RelMetadataQuery mq, LoptMultiJoin multiJoin) {
  // Candidates for self-joins must be simple factors
  Map<Integer, RelOptTable> simpleFactors = getSimpleFactors(mq, multiJoin);

  // See if a simple factor is repeated and therefore potentially is
  // part of a self-join.  Restrict each factor to at most one
  // self-join.
  final List<RelOptTable> repeatedTables = new ArrayList<>();
  final TreeSet<Integer> sortedFactors = new TreeSet<>();
  sortedFactors.addAll(simpleFactors.keySet());
  final Map<Integer, Integer> selfJoinPairs = new HashMap<>();
  Integer [] factors =
      sortedFactors.toArray(new Integer[0]);
  for (int i = 0; i < factors.length; i++) {
    if (repeatedTables.contains(simpleFactors.get(factors[i]))) {
      continue;
    }
    for (int j = i + 1; j < factors.length; j++) {
      int leftFactor = factors[i];
      int rightFactor = factors[j];
      if (simpleFactors.get(leftFactor).getQualifiedName().equals(
          simpleFactors.get(rightFactor).getQualifiedName())) {
        selfJoinPairs.put(leftFactor, rightFactor);
        repeatedTables.add(simpleFactors.get(leftFactor));
        break;
      }
    }
  }

  // From the candidate self-join pairs, determine if there is
  // the appropriate join condition between the two factors that will
  // allow the join to be removed.
  for (Integer factor1 : selfJoinPairs.keySet()) {
    final int factor2 = selfJoinPairs.get(factor1);
    final List<RexNode> selfJoinFilters = new ArrayList<>();
    for (RexNode filter : multiJoin.getJoinFilters()) {
      ImmutableBitSet joinFactors =
          multiJoin.getFactorsRefByJoinFilter(filter);
      if ((joinFactors.cardinality() == 2)
          && joinFactors.get(factor1)
          && joinFactors.get(factor2)) {
        selfJoinFilters.add(filter);
      }
    }
    if ((selfJoinFilters.size() > 0)
        && isSelfJoinFilterUnique(
          mq,
          multiJoin,
          factor1,
          factor2,
          selfJoinFilters)) {
      multiJoin.addRemovableSelfJoinPair(factor1, factor2);
    }
  }
}
 

/**
 * Determines whether a join of the dimension table in a semijoin can be
 * removed. It can be if the dimension keys are unique and the only fields
 * referenced from the dimension table are its semijoin keys. The semijoin
 * keys can be mapped to the corresponding keys from the fact table (because
 * of the equality condition associated with the semijoin keys). Therefore,
 * that's why the dimension table can be removed even though those fields
 * are referenced elsewhere in the query tree.
 *
 * @param multiJoin join factors being optimized
 * @param semiJoin semijoin under consideration
 * @param factIdx id of the fact table in the semijoin
 * @param dimIdx id of the dimension table in the semijoin
 */
private void removeJoin(
    LoptMultiJoin multiJoin,
    SemiJoin semiJoin,
    int factIdx,
    int dimIdx) {
  // if the dimension can be removed because of another semijoin, then
  // no need to proceed any further
  if (multiJoin.getJoinRemovalFactor(dimIdx) != null) {
    return;
  }

  // Check if the semijoin keys corresponding to the dimension table
  // are unique.  The semijoin will filter out the nulls.
  final ImmutableBitSet dimKeys = ImmutableBitSet.of(semiJoin.getRightKeys());
  final RelNode dimRel = multiJoin.getJoinFactor(dimIdx);
  if (!RelMdUtil.areColumnsDefinitelyUniqueWhenNullsFiltered(mq, dimRel,
      dimKeys)) {
    return;
  }

  // check that the only fields referenced from the dimension table
  // in either its projection or join conditions are the dimension
  // keys
  ImmutableBitSet dimProjRefs = multiJoin.getProjFields(dimIdx);
  if (dimProjRefs == null) {
    int nDimFields = multiJoin.getNumFieldsInJoinFactor(dimIdx);
    dimProjRefs = ImmutableBitSet.range(0, nDimFields);
  }
  if (!dimKeys.contains(dimProjRefs)) {
    return;
  }
  int [] dimJoinRefCounts = multiJoin.getJoinFieldRefCounts(dimIdx);
  for (int i = 0; i < dimJoinRefCounts.length; i++) {
    if (dimJoinRefCounts[i] > 0) {
      if (!dimKeys.get(i)) {
        return;
      }
    }
  }

  // criteria met; keep track of the fact table and the semijoin that
  // allow the join of this dimension table to be removed
  multiJoin.setJoinRemovalFactor(dimIdx, factIdx);
  multiJoin.setJoinRemovalSemiJoin(dimIdx, semiJoin);

  // if the dimension table doesn't reference anything in its projection
  // and the only fields referenced in its joins are the dimension keys
  // of this semijoin, then we can decrement the join reference counts
  // corresponding to the fact table's semijoin keys, since the
  // dimension table doesn't need to use those keys
  if (dimProjRefs.cardinality() != 0) {
    return;
  }
  for (int i = 0; i < dimJoinRefCounts.length; i++) {
    if (dimJoinRefCounts[i] > 1) {
      return;
    } else if (dimJoinRefCounts[i] == 1) {
      if (!dimKeys.get(i)) {
        return;
      }
    }
  }
  int [] factJoinRefCounts = multiJoin.getJoinFieldRefCounts(factIdx);
  for (Integer key : semiJoin.getLeftKeys()) {
    factJoinRefCounts[key]--;
  }
}
 

/**
 * Estimates the number of rows which would be produced by a GROUP BY on the
 * set of columns indicated by groupKey.
 * column").
 */
private Double getDistinctRowCountInternal(DrillScanRelBase scan, RelMetadataQuery mq, DrillTable table,
    ImmutableBitSet groupKey, RelDataType type, RexNode predicate) {
  double selectivity, rowCount;
  /* If predicate is present, determine its selectivity to estimate filtered rows.
   * Thereafter, compute the number of distinct rows.
   */
  selectivity = mq.getSelectivity(scan, predicate);
  rowCount = mq.getRowCount(scan);

  if (groupKey.length() == 0) {
    return selectivity * rowCount;
  }

  /* If predicate is present, determine its selectivity to estimate filtered rows. Thereafter,
   * compute the number of distinct rows
   */
  selectivity = mq.getSelectivity(scan, predicate);
  TableMetadata tableMetadata;
  try {
    tableMetadata = table.getGroupScan().getTableMetadata();
  } catch (IOException e) {
    // Statistics cannot be obtained, use default behaviour
    return scan.estimateRowCount(mq) * 0.1;
  }
  double s = 1.0;

  for (int i = 0; i < groupKey.length(); i++) {
    final String colName = type.getFieldNames().get(i);
    // Skip NDV, if not available
    if (!groupKey.get(i)) {
      continue;
    }
    ColumnStatistics columnStatistics = tableMetadata != null ? tableMetadata.getColumnStatistics(SchemaPath.getSimplePath(colName)) : null;
    Double ndv = columnStatistics != null ? (Double) columnStatistics.getStatistic(ColumnStatisticsKind.NVD) : null;
    if (ndv == null) {
      continue;
    }
    s *= 1 - ndv / rowCount;
  }
  if (s > 0 && s < 1.0) {
    return (1 - s) * selectivity * rowCount;
  } else if (s == 1.0) {
    // Could not get any NDV estimate from stats - probably stats not present for GBY cols. So Guess!
    return scan.estimateRowCount(mq) * 0.1;
  } else {
    /* rowCount maybe less than NDV(different source), sanity check OR NDV not used at all */
    return selectivity * rowCount;
  }
}
 

@Override
public void onMatch(RelOptRuleCall call) {
  ProjectPrel project = call.rel(0);
  NestedLoopJoinPrel nlj = call.rel(1);

  ImmutableBitSet topProjectedColumns = InputFinder.bits(project.getProjects(), null);

  ImmutableBitSet bottomProjectedColumns = null;
  if (nlj.getProjectedFields() == null) {
    bottomProjectedColumns = ImmutableBitSet.range(nlj.getRowType().getFieldCount());
  } else {
    bottomProjectedColumns = nlj.getProjectedFields();
  }

  ImmutableBitSet.Builder builder = ImmutableBitSet.builder();
  int field = 0;
  Mapping mapping = Mappings.create(MappingType.SURJECTION, bottomProjectedColumns.cardinality(), topProjectedColumns.cardinality());
  for (Ord<Integer> ord : Ord.zip(bottomProjectedColumns)) {
    if (topProjectedColumns.get(ord.i)) {
      builder.set(ord.e);
      mapping.set(ord.i, field);
      field++;
    }
  }

  if (builder.cardinality() == 0) {
    //project at least one column
    builder.set(0);
  }

  ImmutableBitSet newJoinProjectedFields = builder.build();

  if (newJoinProjectedFields.equals(nlj.getProjectedFields())) {
    return;
  }

  RexShuttle shuttle = new RexPermuteInputsShuttle(mapping);
  List<RexNode> newProjects = shuttle.apply(project.getProjects());

  NestedLoopJoinPrel newJoin = (NestedLoopJoinPrel) nlj.copy(newJoinProjectedFields);
  ProjectPrel newProject = ProjectPrel.create(nlj.getCluster(), project.getTraitSet(), newJoin, newProjects, project.getRowType());
  call.transformTo(newProject);
}
 

/**
 * Locates pairs of joins that are self-joins where the join can be removed
 * because the join condition between the two factors is an equality join on
 * unique keys.
 *
 * @param multiJoin join factors being optimized
 */
private void findRemovableSelfJoins(RelMetadataQuery mq, LoptMultiJoin multiJoin) {
  // Candidates for self-joins must be simple factors
  Map<Integer, RelOptTable> simpleFactors = getSimpleFactors(mq, multiJoin);

  // See if a simple factor is repeated and therefore potentially is
  // part of a self-join.  Restrict each factor to at most one
  // self-join.
  final List<RelOptTable> repeatedTables = new ArrayList<>();
  final TreeSet<Integer> sortedFactors = new TreeSet<>();
  sortedFactors.addAll(simpleFactors.keySet());
  final Map<Integer, Integer> selfJoinPairs = new HashMap<>();
  Integer [] factors =
      sortedFactors.toArray(new Integer[0]);
  for (int i = 0; i < factors.length; i++) {
    if (repeatedTables.contains(simpleFactors.get(factors[i]))) {
      continue;
    }
    for (int j = i + 1; j < factors.length; j++) {
      int leftFactor = factors[i];
      int rightFactor = factors[j];
      if (simpleFactors.get(leftFactor).getQualifiedName().equals(
          simpleFactors.get(rightFactor).getQualifiedName())) {
        selfJoinPairs.put(leftFactor, rightFactor);
        repeatedTables.add(simpleFactors.get(leftFactor));
        break;
      }
    }
  }

  // From the candidate self-join pairs, determine if there is
  // the appropriate join condition between the two factors that will
  // allow the join to be removed.
  for (Integer factor1 : selfJoinPairs.keySet()) {
    final int factor2 = selfJoinPairs.get(factor1);
    final List<RexNode> selfJoinFilters = new ArrayList<>();
    for (RexNode filter : multiJoin.getJoinFilters()) {
      ImmutableBitSet joinFactors =
          multiJoin.getFactorsRefByJoinFilter(filter);
      if ((joinFactors.cardinality() == 2)
          && joinFactors.get(factor1)
          && joinFactors.get(factor2)) {
        selfJoinFilters.add(filter);
      }
    }
    if ((selfJoinFilters.size() > 0)
        && isSelfJoinFilterUnique(
          mq,
          multiJoin,
          factor1,
          factor2,
          selfJoinFilters)) {
      multiJoin.addRemovableSelfJoinPair(factor1, factor2);
    }
  }
}
 

/**
 * Determines whether a join of the dimension table in a semijoin can be
 * removed. It can be if the dimension keys are unique and the only fields
 * referenced from the dimension table are its semijoin keys. The semijoin
 * keys can be mapped to the corresponding keys from the fact table (because
 * of the equality condition associated with the semijoin keys). Therefore,
 * that's why the dimension table can be removed even though those fields
 * are referenced elsewhere in the query tree.
 *
 * @param multiJoin join factors being optimized
 * @param semiJoin semijoin under consideration
 * @param factIdx id of the fact table in the semijoin
 * @param dimIdx id of the dimension table in the semijoin
 */
private void removeJoin(
    LoptMultiJoin multiJoin,
    LogicalJoin semiJoin,
    int factIdx,
    int dimIdx) {
  // if the dimension can be removed because of another semijoin, then
  // no need to proceed any further
  if (multiJoin.getJoinRemovalFactor(dimIdx) != null) {
    return;
  }

  // Check if the semijoin keys corresponding to the dimension table
  // are unique.  The semijoin will filter out the nulls.
  final ImmutableBitSet dimKeys = ImmutableBitSet.of(semiJoin.analyzeCondition().rightKeys);
  final RelNode dimRel = multiJoin.getJoinFactor(dimIdx);
  if (!RelMdUtil.areColumnsDefinitelyUniqueWhenNullsFiltered(mq, dimRel,
      dimKeys)) {
    return;
  }

  // check that the only fields referenced from the dimension table
  // in either its projection or join conditions are the dimension
  // keys
  ImmutableBitSet dimProjRefs = multiJoin.getProjFields(dimIdx);
  if (dimProjRefs == null) {
    int nDimFields = multiJoin.getNumFieldsInJoinFactor(dimIdx);
    dimProjRefs = ImmutableBitSet.range(0, nDimFields);
  }
  if (!dimKeys.contains(dimProjRefs)) {
    return;
  }
  int [] dimJoinRefCounts = multiJoin.getJoinFieldRefCounts(dimIdx);
  for (int i = 0; i < dimJoinRefCounts.length; i++) {
    if (dimJoinRefCounts[i] > 0) {
      if (!dimKeys.get(i)) {
        return;
      }
    }
  }

  // criteria met; keep track of the fact table and the semijoin that
  // allow the join of this dimension table to be removed
  multiJoin.setJoinRemovalFactor(dimIdx, factIdx);
  multiJoin.setJoinRemovalSemiJoin(dimIdx, semiJoin);

  // if the dimension table doesn't reference anything in its projection
  // and the only fields referenced in its joins are the dimension keys
  // of this semijoin, then we can decrement the join reference counts
  // corresponding to the fact table's semijoin keys, since the
  // dimension table doesn't need to use those keys
  if (dimProjRefs.cardinality() != 0) {
    return;
  }
  for (int i = 0; i < dimJoinRefCounts.length; i++) {
    if (dimJoinRefCounts[i] > 1) {
      return;
    } else if (dimJoinRefCounts[i] == 1) {
      if (!dimKeys.get(i)) {
        return;
      }
    }
  }
  int [] factJoinRefCounts = multiJoin.getJoinFieldRefCounts(factIdx);
  for (Integer key : semiJoin.analyzeCondition().leftKeys) {
    factJoinRefCounts[key]--;
  }
}
 

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