Java Code Examples for org.apache.calcite.rel.metadata.RelMetadataQuery#areColumnsUnique()

/**
 * Rewrites a scalar sub-query into an
 * {@link org.apache.calcite.rel.core.Aggregate}.
 *
 * @param e            IN sub-query to rewrite
 * @param variablesSet A set of variables used by a relational
 *                     expression of the specified RexSubQuery
 * @param builder      Builder
 * @param offset       Offset to shift {@link RexInputRef}
 *
 * @return Expression that may be used to replace the RexSubQuery
 */
private RexNode rewriteScalarQuery(RexSubQuery e, Set<CorrelationId> variablesSet, RelBuilder builder,
        int inputCount, int offset) {
    builder.push(e.getRel());
    final RelMetadataQuery mq = e.getRel().getCluster().getMetadataQuery();
    final Boolean unique = mq.areColumnsUnique(builder.peek(), ImmutableBitSet.of());
    if (unique == null || !unique) {
        builder.aggregate(builder.groupKey(),
                builder.aggregateCall(SqlStdOperatorTable.SINGLE_VALUE, builder.field(0)));
    }
    builder.join(JoinRelType.LEFT, builder.literal(true), variablesSet);
    return field(builder, inputCount, offset);
}
 

/**
 * Rewrites a scalar sub-query into an
 * {@link org.apache.calcite.rel.core.Aggregate}.
 *
 * @param e            IN sub-query to rewrite
 * @param variablesSet A set of variables used by a relational
 *                     expression of the specified RexSubQuery
 * @param builder      Builder
 * @param offset       Offset to shift {@link RexInputRef}
 *
 * @return Expression that may be used to replace the RexSubQuery
 */
private RexNode rewriteScalarQuery(RexSubQuery e, Set<CorrelationId> variablesSet,
    RelBuilder builder, int inputCount, int offset) {
  builder.push(e.rel);
  final RelMetadataQuery mq = e.rel.getCluster().getMetadataQuery();
  final Boolean unique = mq.areColumnsUnique(builder.peek(),
      ImmutableBitSet.of());
  if (unique == null || !unique) {
    builder.aggregate(builder.groupKey(),
        builder.aggregateCall(SqlStdOperatorTable.SINGLE_VALUE,
            builder.field(0)));
  }
  builder.join(JoinRelType.LEFT, builder.literal(true), variablesSet);
  return field(builder, inputCount, offset);
}
 

/** Asserts that {@link RelMetadataQuery#getUniqueKeys(RelNode)}
 * and {@link RelMetadataQuery#areColumnsUnique(RelNode, ImmutableBitSet)}
 * return consistent results. */
private void assertUniqueConsistent(RelNode rel) {
  final RelMetadataQuery mq = rel.getCluster().getMetadataQuery();
  final Set<ImmutableBitSet> uniqueKeys = mq.getUniqueKeys(rel);
  final ImmutableBitSet allCols =
      ImmutableBitSet.range(0, rel.getRowType().getFieldCount());
  for (ImmutableBitSet key : allCols.powerSet()) {
    Boolean result2 = mq.areColumnsUnique(rel, key);
    assertEquals(isUnique(uniqueKeys, key), SqlFunctions.isTrue(result2),
        () -> "areColumnsUnique. key: " + key + ", uniqueKeys: " + uniqueKeys
            + ", rel: " + RelOptUtil.toString(rel));
  }
}
 

@Test void testFullOuterJoinUniqueness1() {
  final String sql = "select e.empno, d.deptno\n"
      + "from (select cast(null as int) empno from sales.emp "
      + " where empno = 10 group by cast(null as int)) as e\n"
      + "full outer join (select cast (null as int) deptno from sales.dept "
      + "group by cast(null as int)) as d on e.empno = d.deptno\n"
      + "group by e.empno, d.deptno";
  RelNode rel = convertSql(sql);
  final RelMetadataQuery mq = rel.getCluster().getMetadataQuery();
  final ImmutableBitSet allCols =
      ImmutableBitSet.range(0, rel.getRowType().getFieldCount());
  Boolean areGroupByKeysUnique = mq.areColumnsUnique(rel.getInput(0), allCols);
  assertThat(areGroupByKeysUnique, is(false));
}
 

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

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