Java Code Examples for org.apache.calcite.tools.RelBuilder#aggregate()

/**
 * Rewrites an EXISTS RexSubQuery into a {@link Join}.
 *
 * @param e            EXISTS sub-query to rewrite
 * @param variablesSet A set of variables used by a relational
 *                     expression of the specified RexSubQuery
 * @param logic        Logic for evaluating
 * @param builder      Builder
 *
 * @return Expression that may be used to replace the RexSubQuery
 */
private RexNode rewriteExists(RexSubQuery e, Set<CorrelationId> variablesSet, RelOptUtil.Logic logic,
        RelBuilder builder) {
    builder.push(e.getRel());

    builder.project(builder.alias(builder.literal(true), "i"));
    switch (logic) {
    case TRUE:
        // Handles queries with single EXISTS in filter condition:
        // select e.deptno from emp as e
        // where exists (select deptno from emp)
        builder.aggregate(builder.groupKey(0));
        builder.as("dt");
        builder.join(JoinRelType.INNER, builder.literal(true), variablesSet);
        return builder.literal(true);
    default:
        builder.distinct();
    }

    builder.as("dt");

    builder.join(JoinRelType.LEFT, builder.literal(true), variablesSet);

    return builder.isNotNull(Util.last(builder.fields()));
}
 

/**
 * Rewrites an EXISTS RexSubQuery into a {@link Join}.
 *
 * @param e            EXISTS sub-query to rewrite
 * @param variablesSet A set of variables used by a relational
 *                     expression of the specified RexSubQuery
 * @param logic        Logic for evaluating
 * @param builder      Builder
 *
 * @return Expression that may be used to replace the RexSubQuery
 */
private RexNode rewriteExists(RexSubQuery e, Set<CorrelationId> variablesSet,
    RelOptUtil.Logic logic, RelBuilder builder) {
  builder.push(e.rel);

  builder.project(builder.alias(builder.literal(true), "i"));
  switch (logic) {
  case TRUE:
    // Handles queries with single EXISTS in filter condition:
    // select e.deptno from emp as e
    // where exists (select deptno from emp)
    builder.aggregate(builder.groupKey(0));
    builder.as("dt");
    builder.join(JoinRelType.INNER, builder.literal(true), variablesSet);
    return builder.literal(true);
  default:
    builder.distinct();
  }

  builder.as("dt");

  builder.join(JoinRelType.LEFT, builder.literal(true), variablesSet);

  return builder.isNotNull(Util.last(builder.fields()));
}
 

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

public void onMatch(RelOptRuleCall call) {
  final Aggregate topAggRel = call.rel(0);
  final Union union = call.rel(1);

  // If distincts haven't been removed yet, defer invoking this rule
  if (!union.all) {
    return;
  }

  final RelBuilder relBuilder = call.builder();
  final Aggregate bottomAggRel;
  if (call.rel(3) instanceof Aggregate) {
    // Aggregate is the second input
    bottomAggRel = call.rel(3);
    relBuilder.push(call.rel(2))
        .push(call.rel(3).getInput(0));
  } else if (call.rel(2) instanceof Aggregate) {
    // Aggregate is the first input
    bottomAggRel = call.rel(2);
    relBuilder.push(call.rel(2).getInput(0))
        .push(call.rel(3));
  } else {
    return;
  }

  // Only pull up aggregates if they are there just to remove distincts
  if (!topAggRel.getAggCallList().isEmpty()
      || !bottomAggRel.getAggCallList().isEmpty()) {
    return;
  }

  relBuilder.union(true);
  relBuilder.aggregate(relBuilder.groupKey(topAggRel.getGroupSet()),
      topAggRel.getAggCallList());
  call.transformTo(relBuilder.build());
}
 

private RelNode getNewAggregate(Aggregate oldAggregate, RelBuilder relBuilder, Mapping mapping) {

		final ImmutableBitSet newGroupSet =
			Mappings.apply(mapping, oldAggregate.getGroupSet());

		final Iterable<ImmutableBitSet> newGroupSets =
			oldAggregate.getGroupSets().stream()
				.map(bitSet -> Mappings.apply(mapping, bitSet))
				.collect(Collectors.toList());

		final List<RelBuilder.AggCall> newAggCallList =
			getNewAggCallList(oldAggregate, relBuilder, mapping);

		final RelBuilder.GroupKey groupKey =
			relBuilder.groupKey(newGroupSet, newGroupSets);

		if (oldAggregate instanceof LogicalWindowAggregate) {
			if (newGroupSet.size() == 0 && newAggCallList.size() == 0) {
				// Return the old LogicalWindowAggregate directly, as we can't get an empty Aggregate
				// from the relBuilder.
				return oldAggregate;
			} else {
				relBuilder.aggregate(groupKey, newAggCallList);
				Aggregate newAggregate = (Aggregate) relBuilder.build();
				LogicalWindowAggregate oldLogicalWindowAggregate = (LogicalWindowAggregate) oldAggregate;

				return LogicalWindowAggregate.create(
					oldLogicalWindowAggregate.getWindow(),
					oldLogicalWindowAggregate.getNamedProperties(),
					newAggregate);
			}
		} else {
			relBuilder.aggregate(groupKey, newAggCallList);
			return relBuilder.build();
		}
	}
 

private RelNode getNewAggregate(Aggregate oldAggregate, RelBuilder relBuilder, Mapping mapping) {

		final ImmutableBitSet newGroupSet =
			Mappings.apply(mapping, oldAggregate.getGroupSet());

		final Iterable<ImmutableBitSet> newGroupSets =
			oldAggregate.getGroupSets().stream()
				.map(bitSet -> Mappings.apply(mapping, bitSet))
				.collect(Collectors.toList());

		final List<RelBuilder.AggCall> newAggCallList =
			getNewAggCallList(oldAggregate, relBuilder, mapping);

		final RelBuilder.GroupKey groupKey =
			relBuilder.groupKey(newGroupSet, newGroupSets);

		if (oldAggregate instanceof LogicalWindowAggregate) {
			if (newGroupSet.size() == 0 && newAggCallList.size() == 0) {
				// Return the old LogicalWindowAggregate directly, as we can't get an empty Aggregate
				// from the relBuilder.
				return oldAggregate;
			} else {
				relBuilder.aggregate(groupKey, newAggCallList);
				Aggregate newAggregate = (Aggregate) relBuilder.build();
				LogicalWindowAggregate oldLogicalWindowAggregate = (LogicalWindowAggregate) oldAggregate;

				return LogicalWindowAggregate.create(
					oldLogicalWindowAggregate.getWindow(),
					oldLogicalWindowAggregate.getNamedProperties(),
					newAggregate);
			}
		} else {
			relBuilder.aggregate(groupKey, newAggCallList);
			return relBuilder.build();
		}
	}
 

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

public void onMatch(RelOptRuleCall call) {
  final Aggregate topAggRel = call.rel(0);
  final Union union = call.rel(1);

  // If distincts haven't been removed yet, defer invoking this rule
  if (!union.all) {
    return;
  }

  final RelBuilder relBuilder = call.builder();
  final Aggregate bottomAggRel;
  if (call.rel(3) instanceof Aggregate) {
    // Aggregate is the second input
    bottomAggRel = call.rel(3);
    relBuilder.push(call.rel(2))
        .push(getInputWithSameRowType(bottomAggRel));
  } else if (call.rel(2) instanceof Aggregate) {
    // Aggregate is the first input
    bottomAggRel = call.rel(2);
    relBuilder.push(getInputWithSameRowType(bottomAggRel))
        .push(call.rel(3));
  } else {
    return;
  }

  // Only pull up aggregates if they are there just to remove distincts
  if (!topAggRel.getAggCallList().isEmpty()
      || !bottomAggRel.getAggCallList().isEmpty()) {
    return;
  }

  relBuilder.union(true);
  relBuilder.rename(union.getRowType().getFieldNames());
  relBuilder.aggregate(relBuilder.groupKey(topAggRel.getGroupSet()),
      topAggRel.getAggCallList());
  call.transformTo(relBuilder.build());
}
 

public void onMatch(RelOptRuleCall call) {
  Aggregate aggRel = call.rel(0);
  Union union = call.rel(1);

  if (!union.all) {
    // This transformation is only valid for UNION ALL.
    // Consider t1(i) with rows (5), (5) and t2(i) with
    // rows (5), (10), and the query
    // select sum(i) from (select i from t1) union (select i from t2).
    // The correct answer is 15.  If we apply the transformation,
    // we get
    // select sum(i) from
    // (select sum(i) as i from t1) union (select sum(i) as i from t2)
    // which yields 25 (incorrect).
    return;
  }

  int groupCount = aggRel.getGroupSet().cardinality();

  List<AggregateCall> transformedAggCalls =
      transformAggCalls(
          aggRel.copy(aggRel.getTraitSet(), aggRel.getInput(), false,
              aggRel.getGroupSet(), null, aggRel.getAggCallList()),
          groupCount, aggRel.getAggCallList());
  if (transformedAggCalls == null) {
    // we've detected the presence of something like AVG,
    // which we can't handle
    return;
  }

  // create corresponding aggregates on top of each union child
  final RelBuilder relBuilder = call.builder();
  int transformCount = 0;
  final RelMetadataQuery mq = call.getMetadataQuery();
  for (RelNode input : union.getInputs()) {
    boolean alreadyUnique =
        RelMdUtil.areColumnsDefinitelyUnique(mq, input,
            aggRel.getGroupSet());

    relBuilder.push(input);
    if (!alreadyUnique) {
      ++transformCount;
      relBuilder.aggregate(relBuilder.groupKey(aggRel.getGroupSet()),
          aggRel.getAggCallList());
    }
  }

  if (transformCount == 0) {
    // none of the children could benefit from the push-down,
    // so bail out (preventing the infinite loop to which most
    // planners would succumb)
    return;
  }

  // create a new union whose children are the aggregates created above
  relBuilder.union(true, union.getInputs().size());
  relBuilder.aggregate(
      relBuilder.groupKey(aggRel.getGroupSet(), aggRel.getGroupSets()),
      transformedAggCalls);
  call.transformTo(relBuilder.build());
}
 

@Override
public RelNode visit(LogicalAggregate aggregate) {
  RelNode input = aggregate.getInput().accept(this);

  // Create a new project with null UPDATE_COLUMN below aggregate
  final RelBuilder relBuilder = newCalciteRelBuilderWithoutContext(aggregate.getCluster());
  relBuilder.push(input);
  List<RexNode> nodes = input.getRowType().getFieldList().stream().map(q -> {
    if (UPDATE_COLUMN.equals(q.getName())) {
      return relBuilder.getRexBuilder().makeNullLiteral(q.getType());
    } else{
      return relBuilder.getRexBuilder().makeInputRef(q.getType(), q.getIndex());
    }
  }).collect(Collectors.toList());
  relBuilder.project(nodes, input.getRowType().getFieldNames());

  // create a new aggregate with null UPDATE_COLUMN in groupSet
  RelDataType incomingRowType = relBuilder.peek().getRowType();
  RelDataTypeField modField = incomingRowType.getField(UPDATE_COLUMN, false, false);
  ImmutableBitSet newGroupSet = aggregate.getGroupSet().rebuild().set(modField.getIndex()).build();
  GroupKey groupKey = relBuilder.groupKey(newGroupSet, aggregate.indicator, null);

  final int groupCount = aggregate.getGroupCount();
  final Pointer<Integer> ind = new Pointer<>(groupCount-1);
  final List<String> fieldNames = aggregate.getRowType().getFieldNames();
  final List<AggregateCall> aggCalls = aggregate.getAggCallList().stream().map(q -> {
    ind.value++;
    if (q.getName() == null) {
      return q.rename(fieldNames.get(ind.value));
    }
    return q;
  }).collect(Collectors.toList());

  relBuilder.aggregate(groupKey, aggCalls);

  // create a new project on top to preserve rowType
  Iterable<RexInputRef> projects = FluentIterable.from(aggregate.getRowType().getFieldNames())
    .transform(new Function<String, RexInputRef>() {
      @Override
      public RexInputRef apply(String fieldName) {
        return relBuilder.field(fieldName);
      }
    })
    .append(relBuilder.field(UPDATE_COLUMN));

  relBuilder.project(projects);

  return relBuilder.build();
}
 

@Override
public void onMatch(RelOptRuleCall call) {
  final LogicalAggregate agg = call.rel(0);
  final RelNode input = agg.getInput();

  boolean distinctReplaced = false;
  List<AggregateCall> calls = new ArrayList<>();

  RelMetadataQuery query = null;

  final Boolean[] memo = new Boolean[agg.getInput().getRowType().getFieldCount()];
  for (AggregateCall c : agg.getAggCallList()) {
    final boolean candidate = c.isDistinct() && c.getArgList().size() == 1 && "COUNT".equals(c.getAggregation().getName());

    if(!candidate) {
      calls.add(c);
      continue;
    }

    final int inputOrdinal = c.getArgList().get(0);
    boolean allowed = false;
    if(memo[inputOrdinal] != null) {
      allowed = memo[inputOrdinal];
    } else {
      if(query == null) {
        query = agg.getCluster().getMetadataQuery();
      }

      Set<RelColumnOrigin> origins = query.getColumnOrigins(input, inputOrdinal);

      // see if any column origin allowed a transformation.
      for(RelColumnOrigin o : origins) {
        RelOptTable table = o.getOriginTable();
        NamespaceTable namespaceTable = table.unwrap(NamespaceTable.class);
        if(namespaceTable == null) {
          // unable to decide, no way to transform.
          return;
        }

        if(namespaceTable.isApproximateStatsAllowed()) {
          allowed = true;
        }
      }

      memo[inputOrdinal] = allowed;

    }


    if(allowed) {
      calls.add(AggregateCall.create(HyperLogLog.NDV, false, c.getArgList(), -1, c.getType(), c.getName()));
      distinctReplaced = true;
    } else {
      calls.add(c);
    }

  }

  if(!distinctReplaced) {
    return;
  }

  final RelBuilder builder = relBuilderFactory.create(agg.getCluster(), null);
  builder.push(agg.getInput());
  builder.aggregate(builder.groupKey(agg.getGroupSet().toArray()), calls);
  call.transformTo(builder.build());
}
 

@Override
public void onMatch(RelOptRuleCall call) {
  final LogicalAggregate agg = call.rel(0);
  final RelDataTypeFactory typeFactory = agg.getCluster().getTypeFactory();

  List<AggregateCall> calls = new ArrayList<>();
  Set<Integer> hllApplications = new HashSet<>();

  int i = agg.getGroupCount();
  for (AggregateCall c : agg.getAggCallList()) {

    final int location = i;
    i++;

    if(!"NDV".equals(c.getAggregation().getName())) {
      calls.add(c);
      continue;
    }

    hllApplications.add(location);
    calls.add(AggregateCall.create(HyperLogLog.HLL, false, c.getArgList(), -1, typeFactory.createSqlType(SqlTypeName.VARBINARY, HyperLogLog.HLL_VARBINARY_SIZE), c.getName()));
  }

  if(hllApplications.isEmpty()) {
    return;
  }

  final RelBuilder builder = relBuilderFactory.create(agg.getCluster(), null);
  builder.push(agg.getInput());
  builder.aggregate(builder.groupKey(agg.getGroupSet().toArray()), calls);

  // add the hll application project.
  final List<RexNode> nodes = new ArrayList<>();
  for(int field = 0; field < agg.getRowType().getFieldCount(); field++) {
    if(!hllApplications.contains(field)) {
      nodes.add(builder.field(field));
      continue;
    }

    nodes.add(builder.call(HyperLogLog.HLL_DECODE, builder.field(field)));
  }
  builder.project(nodes);
  call.transformTo(builder.build());
}
 

public void onMatch(RelOptRuleCall call) {
  Aggregate aggRel = call.rel(0);
  Union union = call.rel(1);

  if (!union.all) {
    // This transformation is only valid for UNION ALL.
    // Consider t1(i) with rows (5), (5) and t2(i) with
    // rows (5), (10), and the query
    // select sum(i) from (select i from t1) union (select i from t2).
    // The correct answer is 15.  If we apply the transformation,
    // we get
    // select sum(i) from
    // (select sum(i) as i from t1) union (select sum(i) as i from t2)
    // which yields 25 (incorrect).
    return;
  }

  int groupCount = aggRel.getGroupSet().cardinality();

  List<AggregateCall> transformedAggCalls =
      transformAggCalls(
          aggRel.copy(aggRel.getTraitSet(), aggRel.getInput(),
              aggRel.getGroupSet(), null, aggRel.getAggCallList()),
          groupCount, aggRel.getAggCallList());
  if (transformedAggCalls == null) {
    // we've detected the presence of something like AVG,
    // which we can't handle
    return;
  }

  // create corresponding aggregates on top of each union child
  final RelBuilder relBuilder = call.builder();
  int transformCount = 0;
  final RelMetadataQuery mq = call.getMetadataQuery();
  for (RelNode input : union.getInputs()) {
    boolean alreadyUnique =
        RelMdUtil.areColumnsDefinitelyUnique(mq, input,
            aggRel.getGroupSet());

    relBuilder.push(input);
    if (!alreadyUnique) {
      ++transformCount;
      relBuilder.aggregate(relBuilder.groupKey(aggRel.getGroupSet()),
          aggRel.getAggCallList());
    }
  }

  if (transformCount == 0) {
    // none of the children could benefit from the push-down,
    // so bail out (preventing the infinite loop to which most
    // planners would succumb)
    return;
  }

  // create a new union whose children are the aggregates created above
  relBuilder.union(true, union.getInputs().size());
  relBuilder.aggregate(
      relBuilder.groupKey(aggRel.getGroupSet(),
          (Iterable<ImmutableBitSet>) aggRel.getGroupSets()),
      transformedAggCalls);
  call.transformTo(relBuilder.build());
}
 

/**
 * Do a shallow clone of oldAggRel and update aggCalls. Could be refactored
 * into Aggregate and subclasses - but it's only needed for some
 * subclasses.
 *
 * @param relBuilder Builder of relational expressions; at the top of its
 *                   stack is its input
 * @param oldAggregate LogicalAggregate to clone.
 * @param newCalls  New list of AggregateCalls
 */
protected void newAggregateRel(RelBuilder relBuilder,
    Aggregate oldAggregate,
    List<AggregateCall> newCalls) {
  relBuilder.aggregate(
      relBuilder.groupKey(oldAggregate.getGroupSet(),
          oldAggregate.getGroupSets()),
      newCalls);
}
 

/**
 * Do a shallow clone of oldAggRel and update aggCalls. Could be refactored
 * into Aggregate and subclasses - but it's only needed for some
 * subclasses.
 *
 * @param relBuilder Builder of relational expressions; at the top of its
 *                   stack is its input
 * @param oldAggregate LogicalAggregate to clone.
 * @param newCalls  New list of AggregateCalls
 */
protected void newAggregateRel(RelBuilder relBuilder,
    Aggregate oldAggregate,
    List<AggregateCall> newCalls) {
  relBuilder.aggregate(
      relBuilder.groupKey(oldAggregate.getGroupSet(),
          (Iterable<ImmutableBitSet>) oldAggregate.getGroupSets()),
      newCalls);
}