Java Code Examples for org.apache.calcite.rel.core.Project#getProjects()

The following examples show how to use org.apache.calcite.rel.core.Project#getProjects() . You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example 1
Source File: RelToSqlConverter.java    From calcite with Apache License 2.0 6 votes vote down vote up
/** @see #dispatch */
public Result visit(Project e) {
  e.getVariablesSet();
  Result x = visitChild(0, e.getInput());
  parseCorrelTable(e, x);
  if (isStar(e.getProjects(), e.getInput().getRowType(), e.getRowType())) {
    return x;
  }
  final Builder builder =
      x.builder(e, Clause.SELECT);
  final List<SqlNode> selectList = new ArrayList<>();
  for (RexNode ref : e.getProjects()) {
    SqlNode sqlExpr = builder.context.toSql(null, ref);
    if (SqlUtil.isNullLiteral(sqlExpr, false)) {
      sqlExpr = castNullType(sqlExpr, e.getRowType().getFieldList().get(selectList.size()));
    }
    addSelect(selectList, sqlExpr, e.getRowType());
  }

  builder.setSelect(new SqlNodeList(selectList, POS));
  return builder.result();
}
 
Example 2
Source File: RelMdColumnUniqueness.java    From Bats with Apache License 2.0 6 votes vote down vote up
private boolean simplyProjects(RelNode rel, ImmutableBitSet columns) {
  if (!(rel instanceof Project)) {
    return false;
  }
  Project project = (Project) rel;
  final List<RexNode> projects = project.getProjects();
  for (int column : columns) {
    if (column >= projects.size()) {
      return false;
    }
    if (!(projects.get(column) instanceof RexInputRef)) {
      return false;
    }
    final RexInputRef ref = (RexInputRef) projects.get(column);
    if (ref.getIndex() != column) {
      return false;
    }
  }
  return true;
}
 
Example 3
Source File: DrillRelOptUtil.java    From Bats with Apache License 2.0 6 votes vote down vote up
/**
 * Find whether the given project rel has unknown output schema. This would happen if the
 * project has CONVERT_FROMJSON which can only derive the schema after evaluation is performed
 * @param project : The project rel
 * @return : Return true if the project output schema is unknown. Otherwise, false.
 */
public static boolean isProjectOutputSchemaUnknown(Project project) {
    try {
        RexVisitor<Void> visitor = new RexVisitorImpl<Void>(true) {
            @Override
            public Void visitCall(RexCall call) {
                if ("convert_fromjson".equals(call.getOperator().getName().toLowerCase())) {
                    throw new Util.FoundOne(call); /* throw exception to interrupt tree walk (this is similar to
                                                   other utility methods in RexUtil.java */
                }
                return super.visitCall(call);
            }
        };
        for (RexNode rex : project.getProjects()) {
            rex.accept(visitor);
        }
    } catch (Util.FoundOne e) {
        Util.swallow(e, null);
        return true;
    }
    return false;
}
 
Example 4
Source File: RexUtil.java    From calcite with Apache License 2.0 5 votes vote down vote up
/** Returns whether a {@link Project} contains a sub-query. */
public static boolean containsSubQuery(Project project) {
  for (RexNode node : project.getProjects()) {
    try {
      node.accept(INSTANCE);
    } catch (Util.FoundOne e) {
      return true;
    }
  }
  return false;
}
 
Example 5
Source File: JdbcRules.java    From calcite with Apache License 2.0 5 votes vote down vote up
private static boolean userDefinedFunctionInProject(Project project) {
  CheckingUserDefinedFunctionVisitor visitor = new CheckingUserDefinedFunctionVisitor();
  for (RexNode node : project.getProjects()) {
    node.accept(visitor);
    if (visitor.containsUserDefinedFunction()) {
      return true;
    }
  }
  return false;
}
 
Example 6
Source File: JdbcRules.java    From calcite with Apache License 2.0 5 votes vote down vote up
public RelNode convert(RelNode rel) {
  final Project project = (Project) rel;

  return new JdbcProject(
      rel.getCluster(),
      rel.getTraitSet().replace(out),
      convert(
          project.getInput(),
          project.getInput().getTraitSet().replace(out)),
      project.getProjects(),
      project.getRowType());
}
 
Example 7
Source File: RelMdSize.java    From calcite with Apache License 2.0 5 votes vote down vote up
public List<Double> averageColumnSizes(Project rel, RelMetadataQuery mq) {
  final List<Double> inputColumnSizes =
      mq.getAverageColumnSizesNotNull(rel.getInput());
  final ImmutableNullableList.Builder<Double> sizes =
      ImmutableNullableList.builder();
  for (RexNode project : rel.getProjects()) {
    sizes.add(averageRexSize(project, inputColumnSizes));
  }
  return sizes.build();
}
 
Example 8
Source File: DrillRelOptUtil.java    From Bats with Apache License 2.0 5 votes vote down vote up
/**
 * Find whether the given project rel can produce non-scalar output (hence unknown rowcount). This
 * would happen if the project has a flatten
 * @param project : The project rel
 * @return : Return true if the rowcount is unknown. Otherwise, false.
 */
public static boolean isProjectOutputRowcountUnknown(Project project) {
    for (RexNode rex : project.getProjects()) {
        if (rex instanceof RexCall) {
            if ("flatten".equals(((RexCall) rex).getOperator().getName().toLowerCase())) {
                return true;
            }
        }
    }
    return false;
}
 
Example 9
Source File: RexUtil.java    From Bats with Apache License 2.0 5 votes vote down vote up
/** Returns whether a {@link Project} contains a sub-query. */
public static boolean containsSubQuery(Project project) {
    for (RexNode node : project.getProjects()) {
        try {
            node.accept(INSTANCE);
        } catch (Util.FoundOne e) {
            return true;
        }
    }
    return false;
}
 
Example 10
Source File: RelMdPopulationSize.java    From Bats with Apache License 2.0 5 votes vote down vote up
public Double getPopulationSize(Project rel, RelMetadataQuery mq,
    ImmutableBitSet groupKey) {
  ImmutableBitSet.Builder baseCols = ImmutableBitSet.builder();
  ImmutableBitSet.Builder projCols = ImmutableBitSet.builder();
  List<RexNode> projExprs = rel.getProjects();
  RelMdUtil.splitCols(projExprs, groupKey, baseCols, projCols);

  Double population =
      mq.getPopulationSize(rel.getInput(), baseCols.build());
  if (population == null) {
    return null;
  }

  // No further computation required if the projection expressions are
  // all column references
  if (projCols.cardinality() == 0) {
    return population;
  }

  for (int bit : projCols.build()) {
    Double subRowCount =
        RelMdUtil.cardOfProjExpr(mq, rel, projExprs.get(bit));
    if (subRowCount == null) {
      return null;
    }
    population *= subRowCount;
  }

  // REVIEW zfong 6/22/06 - Broadbase did not have the call to
  // numDistinctVals.  This is needed; otherwise, population can be
  // larger than the number of rows in the RelNode.
  return RelMdUtil.numDistinctVals(population, mq.getRowCount(rel));
}
 
Example 11
Source File: RelDecorrelator.java    From calcite with Apache License 2.0 5 votes vote down vote up
@Override public RelNode visit(RelNode other) {
  if (other instanceof Join) {
    Join join = (Join) other;
    try {
      stack.push(join);
      join.getCondition().accept(rexVisitor(join));
    } finally {
      stack.pop();
    }
    return visitJoin(join);
  } else if (other instanceof Correlate) {
    Correlate correlate = (Correlate) other;
    mapCorToCorRel.put(correlate.getCorrelationId(), correlate);
    return visitJoin(correlate);
  } else if (other instanceof Filter) {
    Filter filter = (Filter) other;
    try {
      stack.push(filter);
      filter.getCondition().accept(rexVisitor(filter));
    } finally {
      stack.pop();
    }
  } else if (other instanceof Project) {
    Project project = (Project) other;
    try {
      stack.push(project);
      for (RexNode node : project.getProjects()) {
        node.accept(rexVisitor(project));
      }
    } finally {
      stack.pop();
    }
  }
  return super.visit(other);
}
 
Example 12
Source File: ElasticProjectRule.java    From dremio-oss with Apache License 2.0 4 votes vote down vote up
ProjectConverterVisitor(Project project, RelNode input) {
  super(input);
  projectExprs = project.getProjects();
  projectDataType = project.getRowType();
}
 
Example 13
Source File: RelDecorrelator.java    From calcite with Apache License 2.0 4 votes vote down vote up
public Frame decorrelateRel(Project rel) {
  //
  // Rewrite logic:
  //
  // 1. Pass along any correlated variables coming from the input.
  //

  final RelNode oldInput = rel.getInput();
  Frame frame = getInvoke(oldInput, rel);
  if (frame == null) {
    // If input has not been rewritten, do not rewrite this rel.
    return null;
  }
  final List<RexNode> oldProjects = rel.getProjects();
  final List<RelDataTypeField> relOutput = rel.getRowType().getFieldList();

  // Project projects the original expressions,
  // plus any correlated variables the input wants to pass along.
  final List<Pair<RexNode, String>> projects = new ArrayList<>();

  // If this Project has correlated reference, create value generator
  // and produce the correlated variables in the new output.
  if (cm.mapRefRelToCorRef.containsKey(rel)) {
    frame = decorrelateInputWithValueGenerator(rel, frame);
  }

  // Project projects the original expressions
  final Map<Integer, Integer> mapOldToNewOutputs = new HashMap<>();
  int newPos;
  for (newPos = 0; newPos < oldProjects.size(); newPos++) {
    projects.add(
        newPos,
        Pair.of(
            decorrelateExpr(currentRel, map, cm, oldProjects.get(newPos)),
            relOutput.get(newPos).getName()));
    mapOldToNewOutputs.put(newPos, newPos);
  }

  // Project any correlated variables the input wants to pass along.
  final SortedMap<CorDef, Integer> corDefOutputs = new TreeMap<>();
  for (Map.Entry<CorDef, Integer> entry : frame.corDefOutputs.entrySet()) {
    projects.add(
        RexInputRef.of2(entry.getValue(),
            frame.r.getRowType().getFieldList()));
    corDefOutputs.put(entry.getKey(), newPos);
    newPos++;
  }

  RelNode newProject = relBuilder.push(frame.r)
      .projectNamed(Pair.left(projects), Pair.right(projects), true)
      .build();

  newProject = RelOptUtil.copyRelHints(rel, newProject);

  return register(rel, newProject, mapOldToNewOutputs, corDefOutputs);
}
 
Example 14
Source File: RelMdDistinctRowCount.java    From calcite with Apache License 2.0 4 votes vote down vote up
public Double getDistinctRowCount(Project rel, RelMetadataQuery mq,
    ImmutableBitSet groupKey, RexNode predicate) {
  if (predicate == null || predicate.isAlwaysTrue()) {
    if (groupKey.isEmpty()) {
      return 1D;
    }
  }
  ImmutableBitSet.Builder baseCols = ImmutableBitSet.builder();
  ImmutableBitSet.Builder projCols = ImmutableBitSet.builder();
  List<RexNode> projExprs = rel.getProjects();
  RelMdUtil.splitCols(projExprs, groupKey, baseCols, projCols);

  final List<RexNode> notPushable = new ArrayList<>();
  final List<RexNode> pushable = new ArrayList<>();
  RelOptUtil.splitFilters(
      ImmutableBitSet.range(rel.getRowType().getFieldCount()),
      predicate,
      pushable,
      notPushable);
  final RexBuilder rexBuilder = rel.getCluster().getRexBuilder();

  // get the distinct row count of the child input, passing in the
  // columns and filters that only reference the child; convert the
  // filter to reference the children projection expressions
  RexNode childPred =
      RexUtil.composeConjunction(rexBuilder, pushable, true);
  RexNode modifiedPred;
  if (childPred == null) {
    modifiedPred = null;
  } else {
    modifiedPred = RelOptUtil.pushPastProject(childPred, rel);
  }
  Double distinctRowCount =
      mq.getDistinctRowCount(rel.getInput(), baseCols.build(),
          modifiedPred);

  if (distinctRowCount == null) {
    return null;
  } else if (!notPushable.isEmpty()) {
    RexNode preds =
        RexUtil.composeConjunction(rexBuilder, notPushable, true);
    distinctRowCount *= RelMdUtil.guessSelectivity(preds);
  }

  // No further computation required if the projection expressions
  // are all column references
  if (projCols.cardinality() == 0) {
    return distinctRowCount;
  }

  // multiply by the cardinality of the non-child projection expressions
  for (int bit : projCols.build()) {
    Double subRowCount =
        RelMdUtil.cardOfProjExpr(mq, rel, projExprs.get(bit));
    if (subRowCount == null) {
      return null;
    }
    distinctRowCount *= subRowCount;
  }

  return RelMdUtil.numDistinctVals(distinctRowCount, mq.getRowCount(rel));
}
 
Example 15
Source File: RelMdColumnUniqueness.java    From Bats with Apache License 2.0 4 votes vote down vote up
public Boolean areColumnsUnique(Project rel, RelMetadataQuery mq,
    ImmutableBitSet columns, boolean ignoreNulls) {
  // LogicalProject maps a set of rows to a different set;
  // Without knowledge of the mapping function(whether it
  // preserves uniqueness), it is only safe to derive uniqueness
  // info from the child of a project when the mapping is f(a) => a.
  //
  // Also need to map the input column set to the corresponding child
  // references

  List<RexNode> projExprs = rel.getProjects();
  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;
      }
      RelDataTypeFactory typeFactory =
          rel.getCluster().getTypeFactory();
      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);
}
 
Example 16
Source File: RelDecorrelator.java    From calcite with Apache License 2.0 4 votes vote down vote up
private void onMatch2(
    RelOptRuleCall call,
    Correlate correlate,
    RelNode leftInput,
    Project aggOutputProject,
    Aggregate aggregate) {
  if (generatedCorRels.contains(correlate)) {
    // This Correlate was generated by a previous invocation of
    // this rule. No further work to do.
    return;
  }

  setCurrent(call.getPlanner().getRoot(), correlate);

  // check for this pattern
  // The pattern matching could be simplified if rules can be applied
  // during decorrelation,
  //
  // CorrelateRel(left correlation, condition = true)
  //   leftInput
  //   Project-A (a RexNode)
  //     Aggregate (groupby (0), agg0(), agg1()...)

  // check aggOutputProj projects only one expression
  List<RexNode> aggOutputProjExprs = aggOutputProject.getProjects();
  if (aggOutputProjExprs.size() != 1) {
    return;
  }

  JoinRelType joinType = correlate.getJoinType();
  // corRel.getCondition was here, however Correlate was updated so it
  // never includes a join condition. The code was not modified for brevity.
  RexNode joinCond = relBuilder.literal(true);
  if ((joinType != JoinRelType.LEFT)
      || (joinCond != relBuilder.literal(true))) {
    return;
  }

  // check that the agg is on the entire input
  if (!aggregate.getGroupSet().isEmpty()) {
    return;
  }

  List<AggregateCall> aggCalls = aggregate.getAggCallList();
  Set<Integer> isCount = new HashSet<>();

  // remember the count() positions
  int i = -1;
  for (AggregateCall aggCall : aggCalls) {
    ++i;
    if (aggCall.getAggregation() instanceof SqlCountAggFunction) {
      isCount.add(i);
    }
  }

  // now rewrite the plan to
  //
  // Project-A' (all LHS plus transformed original projections,
  //             replacing references to count() with case statement)
  //   Correlate(left correlation, condition = true)
  //     leftInput
  //     Aggregate(groupby (0), agg0(), agg1()...)
  //
  List<RexNode> requiredNodes =
      correlate.getRequiredColumns().asList().stream()
          .map(ord -> relBuilder.getRexBuilder().makeInputRef(correlate, ord))
          .collect(Collectors.toList());
  Correlate newCorrelate = (Correlate) relBuilder.push(leftInput)
      .push(aggregate).correlate(correlate.getJoinType(),
          correlate.getCorrelationId(),
          requiredNodes).build();


  // remember this rel so we don't fire rule on it again
  // REVIEW jhyde 29-Oct-2007: rules should not save state; rule
  // should recognize patterns where it does or does not need to do
  // work
  generatedCorRels.add(newCorrelate);

  // need to update the mapCorToCorRel Update the output position
  // for the corVars: only pass on the corVars that are not used in
  // the join key.
  if (cm.mapCorToCorRel.get(correlate.getCorrelationId()) == correlate) {
    cm.mapCorToCorRel.put(correlate.getCorrelationId(), newCorrelate);
  }

  RelNode newOutput =
      aggregateCorrelatorOutput(newCorrelate, aggOutputProject, isCount);

  call.transformTo(newOutput);
}
 
Example 17
Source File: AggregateCaseToFilterRule.java    From calcite with Apache License 2.0 4 votes vote down vote up
@Override public void onMatch(RelOptRuleCall call) {
  final Aggregate aggregate = call.rel(0);
  final Project project = call.rel(1);
  final RexBuilder rexBuilder = aggregate.getCluster().getRexBuilder();
  final List<AggregateCall> newCalls =
      new ArrayList<>(aggregate.getAggCallList().size());
  final List<RexNode> newProjects = new ArrayList<>(project.getProjects());
  final List<RexNode> newCasts = new ArrayList<>();

  for (int fieldNumber : aggregate.getGroupSet()) {
    newCasts.add(
        rexBuilder.makeInputRef(
            project.getProjects().get(fieldNumber).getType(), fieldNumber));
  }

  for (AggregateCall aggregateCall : aggregate.getAggCallList()) {
    AggregateCall newCall =
        transform(aggregateCall, project, newProjects);

    // Possibly CAST the new aggregator to an appropriate type.
    final int i = newCasts.size();
    final RelDataType oldType =
        aggregate.getRowType().getFieldList().get(i).getType();
    if (newCall == null) {
      newCalls.add(aggregateCall);
      newCasts.add(rexBuilder.makeInputRef(oldType, i));
    } else {
      newCalls.add(newCall);
      newCasts.add(
          rexBuilder.makeCast(oldType,
              rexBuilder.makeInputRef(newCall.getType(), i)));
    }
  }

  if (newCalls.equals(aggregate.getAggCallList())) {
    return;
  }

  final RelBuilder relBuilder = call.builder()
      .push(project.getInput())
      .project(newProjects);

  final RelBuilder.GroupKey groupKey =
      relBuilder.groupKey(aggregate.getGroupSet(),
          (Iterable<ImmutableBitSet>) aggregate.getGroupSets());

  relBuilder.aggregate(groupKey, newCalls)
      .convert(aggregate.getRowType(), false);

  call.transformTo(relBuilder.build());
  call.getPlanner().prune(aggregate);
}
 
Example 18
Source File: ProjectJoinRemoveRule.java    From calcite with Apache License 2.0 4 votes vote down vote up
@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);
}
 
Example 19
Source File: ProjectJoinJoinRemoveRule.java    From calcite with Apache License 2.0 4 votes vote down vote up
@Override public void onMatch(RelOptRuleCall call) {
  final Project project = call.rel(0);
  final Join topJoin = call.rel(1);
  final Join bottomJoin = call.rel(2);
  int leftBottomChildSize = bottomJoin.getLeft().getRowType().getFieldCount();

  // Check whether the project uses columns in the right input of bottom join.
  for (RexNode expr: project.getProjects()) {
    if (RelOptUtil.InputFinder.bits(expr).asList().stream().anyMatch(
        i -> i >= leftBottomChildSize
            && i < bottomJoin.getRowType().getFieldCount())) {
      return;
    }
  }

  // Check whether the top join uses columns in the right input of bottom join.
  final List<Integer> leftKeys = new ArrayList<>();
  RelOptUtil.splitJoinCondition(topJoin.getLeft(), topJoin.getRight(),
      topJoin.getCondition(), leftKeys, new ArrayList<>(),
      new ArrayList<>());
  if (leftKeys.stream().anyMatch(s -> s >= leftBottomChildSize)) {
    return;
  }

  // Check whether left join keys in top join and bottom join are equal.
  final List<Integer> leftChildKeys = new ArrayList<>();
  final List<Integer> rightChildKeys = new ArrayList<>();
  RelOptUtil.splitJoinCondition(bottomJoin.getLeft(), bottomJoin.getRight(),
      bottomJoin.getCondition(), leftChildKeys, rightChildKeys,
      new ArrayList<>());
  if (!leftKeys.equals(leftChildKeys)) {
    return;
  }

  // Make sure that right keys of bottom join are unique.
  final ImmutableBitSet.Builder columns = ImmutableBitSet.builder();
  rightChildKeys.forEach(key -> columns.set(key));
  final RelMetadataQuery mq = call.getMetadataQuery();
  if (!mq.areColumnsUnique(bottomJoin.getRight(), columns.build())) {
    return;
  }

  int offset = bottomJoin.getRight().getRowType().getFieldCount();
  final RelBuilder relBuilder = call.builder();

  final RexNode condition = RexUtil.shift(topJoin.getCondition(),
      leftBottomChildSize, -offset);
  final RelNode join = relBuilder.push(bottomJoin.getLeft())
      .push(topJoin.getRight())
      .join(topJoin.getJoinType(), condition)
      .build();

  final List<RexNode> newExprs = project.getProjects().stream()
      .map(expr -> RexUtil.shift(expr, leftBottomChildSize, -offset))
      .collect(Collectors.toList());
  relBuilder.push(join).project(newExprs);
  call.transformTo(relBuilder.build());
}
 
Example 20
Source File: RelMdUniqueKeys.java    From Bats with Apache License 2.0 4 votes vote down vote up
public Set<ImmutableBitSet> getUniqueKeys(Project rel, RelMetadataQuery mq,
    boolean ignoreNulls) {
  // LogicalProject maps a set of rows to a different set;
  // Without knowledge of the mapping function(whether it
  // preserves uniqueness), it is only safe to derive uniqueness
  // info from the child of a project when the mapping is f(a) => a.
  //
  // Further more, the unique bitset coming from the child needs
  // to be mapped to match the output of the project.
  final Map<Integer, Integer> mapInToOutPos = new HashMap<>();
  final List<RexNode> projExprs = rel.getProjects();
  final Set<ImmutableBitSet> projUniqueKeySet = new HashSet<>();

  // Build an input to output position map.
  for (int i = 0; i < projExprs.size(); i++) {
    RexNode projExpr = projExprs.get(i);
    if (projExpr instanceof RexInputRef) {
      mapInToOutPos.put(((RexInputRef) projExpr).getIndex(), i);
    }
  }

  if (mapInToOutPos.isEmpty()) {
    // if there's no RexInputRef in the projected expressions
    // return empty set.
    return projUniqueKeySet;
  }

  Set<ImmutableBitSet> childUniqueKeySet =
      mq.getUniqueKeys(rel.getInput(), ignoreNulls);

  if (childUniqueKeySet != null) {
    // Now add to the projUniqueKeySet the child keys that are fully
    // projected.
    for (ImmutableBitSet colMask : childUniqueKeySet) {
      ImmutableBitSet.Builder tmpMask = ImmutableBitSet.builder();
      boolean completeKeyProjected = true;
      for (int bit : colMask) {
        if (mapInToOutPos.containsKey(bit)) {
          tmpMask.set(mapInToOutPos.get(bit));
        } else {
          // Skip the child unique key if part of it is not
          // projected.
          completeKeyProjected = false;
          break;
        }
      }
      if (completeKeyProjected) {
        projUniqueKeySet.add(tmpMask.build());
      }
    }
  }

  return projUniqueKeySet;
}