Java Code Examples for org.apache.calcite.rel.core.Aggregate#getGroupType()

private static boolean isAggregateSupported(Aggregate aggregate, boolean allowFunctions) {
  if (!allowFunctions && !aggregate.getAggCallList().isEmpty()) {
    return false;
  }
  if (aggregate.getGroupType() != Aggregate.Group.SIMPLE) {
    return false;
  }
  // If any aggregate functions do not support splitting, bail out
  // If any aggregate call has a filter or is distinct, bail out
  for (AggregateCall aggregateCall : aggregate.getAggCallList()) {
    if (aggregateCall.getAggregation().unwrap(SqlSplittableAggFunction.class)
        == null) {
      return false;
    }
    if (aggregateCall.filterArg >= 0 || aggregateCall.isDistinct()) {
      return false;
    }
  }
  return true;
}
 

private boolean canPush(Aggregate aggregate, ImmutableBitSet rCols) {
  // If the filter references columns not in the group key, we cannot push
  final ImmutableBitSet groupKeys =
      ImmutableBitSet.range(0, aggregate.getGroupSet().cardinality());
  if (!groupKeys.contains(rCols)) {
    return false;
  }

  if (aggregate.getGroupType() != Group.SIMPLE) {
    // If grouping sets are used, the filter can be pushed if
    // the columns referenced in the predicate are present in
    // all the grouping sets.
    for (ImmutableBitSet groupingSet : aggregate.getGroupSets()) {
      if (!groupingSet.contains(rCols)) {
        return false;
      }
    }
  }
  return true;
}
 

@Override
public boolean matches(RelOptRuleCall call) {
	final Aggregate aggregate = call.rel(0);
	final RelNode input = call.rel(1);
	if (aggregate.getGroupCount() == 0 || aggregate.indicator ||
			aggregate.getGroupType() != Aggregate.Group.SIMPLE) {
		return false;
	}
	for (AggregateCall aggCall : aggregate.getAggCallList()) {
		SqlKind aggCallKind = aggCall.getAggregation().getKind();
		// TODO supports more AggregateCalls
		boolean isAllowAggCall = aggCallKind == SqlKind.SUM ||
				aggCallKind == SqlKind.MIN ||
				aggCallKind == SqlKind.MAX ||
				aggCall.getAggregation() instanceof SqlAuxiliaryGroupAggFunction;
		if (!isAllowAggCall || aggCall.filterArg >= 0 || aggCall.getArgList().size() != 1) {
			return false;
		}
	}

	final RelMetadataQuery mq = call.getMetadataQuery();
	return SqlFunctions.isTrue(mq.areColumnsUnique(input, aggregate.getGroupSet()));
}
 

@Override
public boolean matches(RelOptRuleCall call) {
	final Aggregate aggregate = call.rel(0);
	final RelNode input = call.rel(1);
	if (aggregate.getGroupCount() == 0 || aggregate.indicator ||
			aggregate.getGroupType() != Aggregate.Group.SIMPLE) {
		return false;
	}
	for (AggregateCall aggCall : aggregate.getAggCallList()) {
		SqlKind aggCallKind = aggCall.getAggregation().getKind();
		// TODO supports more AggregateCalls
		boolean isAllowAggCall = aggCallKind == SqlKind.SUM ||
				aggCallKind == SqlKind.MIN ||
				aggCallKind == SqlKind.MAX ||
				aggCall.getAggregation() instanceof SqlAuxiliaryGroupAggFunction;
		if (!isAllowAggCall || aggCall.filterArg >= 0 || aggCall.getArgList().size() != 1) {
			return false;
		}
	}

	final RelMetadataQuery mq = call.getMetadataQuery();
	return SqlFunctions.isTrue(mq.areColumnsUnique(input, aggregate.getGroupSet()));
}
 

public Double getMaxRowCount(Aggregate rel, RelMetadataQuery mq) {
  if (rel.getGroupSet().isEmpty()) {
    // Aggregate with no GROUP BY always returns 1 row (even on empty table).
    return 1D;
  }

  // Aggregate with constant GROUP BY always returns 1 row
  if (rel.getGroupType() == Aggregate.Group.SIMPLE) {
    final RelOptPredicateList predicateList =
        mq.getPulledUpPredicates(rel.getInput());
    if (predicateList != null
        && allGroupKeysAreConstant(rel, predicateList)) {
      return 1D;
    }
  }
  final Double rowCount = mq.getMaxRowCount(rel.getInput());
  if (rowCount == null) {
    return null;
  }
  return rowCount * rel.getGroupSets().size();
}
 

private static boolean isAggregateSupported(Aggregate aggregate, boolean allowFunctions) {
  if (!allowFunctions && !aggregate.getAggCallList().isEmpty()) {
    return false;
  }
  if (aggregate.getGroupType() != Aggregate.Group.SIMPLE) {
    return false;
  }
  // If any aggregate functions do not support splitting, bail out
  // If any aggregate call has a filter or is distinct, bail out
  for (AggregateCall aggregateCall : aggregate.getAggCallList()) {
    if (aggregateCall.getAggregation().unwrap(SqlSplittableAggFunction.class)
        == null) {
      return false;
    }
    if (aggregateCall.filterArg >= 0 || aggregateCall.isDistinct()) {
      return false;
    }
  }
  return true;
}
 

private boolean canPush(Aggregate aggregate, ImmutableBitSet rCols) {
  // If the filter references columns not in the group key, we cannot push
  final ImmutableBitSet groupKeys =
      ImmutableBitSet.range(0, aggregate.getGroupSet().cardinality());
  if (!groupKeys.contains(rCols)) {
    return false;
  }

  if (aggregate.getGroupType() != Group.SIMPLE) {
    // If grouping sets are used, the filter can be pushed if
    // the columns referenced in the predicate are present in
    // all the grouping sets.
    for (ImmutableBitSet groupingSet : aggregate.getGroupSets()) {
      if (!groupingSet.contains(rCols)) {
        return false;
      }
    }
  }
  return true;
}
 

@Override
protected boolean isValidPlan(Project topProject, RelNode node, RelMetadataQuery mq) {
    if (!(node instanceof Aggregate)) {
        return false;
    }
    Aggregate aggregate = (Aggregate) node;
    if (aggregate.getGroupType() != Aggregate.Group.SIMPLE) {
        // TODO: Rewriting with grouping sets not supported yet
        return false;
    }
    return isValidRelNodePlan(aggregate.getInput(), mq);
}
 

private Aggregate mergeAggregate(Aggregate aggregate1, Aggregate aggregate2) {
  //Support only simple groups
  if (aggregate1.getGroupType() != Aggregate.Group.SIMPLE
      || aggregate2.getGroupType() != Aggregate.Group.SIMPLE) {
    return null;
  }

  final int callLen1 = aggregate1.getAggCallList().size();
  final int callLen2 = aggregate2.getAggCallList().size();
  final List<AggregateCall> newAggCalls = Lists.newArrayList();
  if (callLen1 <= callLen2) {
    //Create new Call list
    for (AggregateCall call : aggregate1.getAggCallList()) {
      AggregateCall newAggCall = getMergedAggCall(aggregate2, call);
      if (newAggCall == null) {
        return null;
      } else {
        newAggCalls.add(newAggCall);
      }
    }

    //Create new groupSets
    ImmutableBitSet.Builder groupSetsBuilder = ImmutableBitSet.builder();
    for (int key : aggregate1.getGroupSet()) {
      try {
        groupSetsBuilder.set(aggregate2.getGroupSet().nth(key));
      } catch (IndexOutOfBoundsException e) {
        return null;
      }
    }
    final ImmutableBitSet newGroupSets = groupSetsBuilder.build();
    return aggregate1.copy(aggregate1.getTraitSet(), aggregate2.getInput(),
        aggregate1.indicator, newGroupSets, ImmutableList.of(newGroupSets), newAggCalls);
  } else {
    return null;
  }
}
 

@Override protected boolean isValidPlan(Project topProject, RelNode node,
    RelMetadataQuery mq) {
  if (!(node instanceof Aggregate)) {
    return false;
  }
  Aggregate aggregate = (Aggregate) node;
  if (aggregate.getGroupType() != Aggregate.Group.SIMPLE) {
    // TODO: Rewriting with grouping sets not supported yet
    return false;
  }
  return isValidRelNodePlan(aggregate.getInput(), mq);
}
 

private static boolean isAggregateSupported(Aggregate aggregate) {
  if (aggregate.getGroupType() != Aggregate.Group.SIMPLE
      || aggregate.getGroupCount() == 0) {
    return false;
  }
  // If any aggregate functions do not support splitting, bail out.
  for (AggregateCall aggregateCall : aggregate.getAggCallList()) {
    if (aggregateCall.filterArg >= 0
        || aggregateCall.getAggregation()
            .unwrap(SqlSplittableAggFunction.class) == null) {
      return false;
    }
  }
  return true;
}
 

public boolean hasTrickyRollup(Sort e, Aggregate aggregate) {
  return !dialect.supportsAggregateFunction(SqlKind.ROLLUP)
      && dialect.supportsGroupByWithRollup()
      && (aggregate.getGroupType() == Aggregate.Group.ROLLUP
          || aggregate.getGroupType() == Aggregate.Group.CUBE
              && aggregate.getGroupSet().cardinality() == 1)
      && e.collation.getFieldCollations().stream().allMatch(fc ->
          fc.getFieldIndex() < aggregate.getGroupSet().cardinality());
}
 

@Override
public boolean apply(@Nullable Aggregate input) {
    return input.getGroupType() != Aggregate.Group.SIMPLE;
}
 

@Override
public boolean apply(@Nullable Aggregate input) {
    return input.getGroupType() != Aggregate.Group.SIMPLE;
}
 

public static RelNode apply(RelOptRuleCall call, Aggregate aggregate,
    Project project) {
  // Find all fields which we need to be straightforward field projections.
  final Set<Integer> interestingFields = RelOptUtil.getAllFields(aggregate);

  // Build the map from old to new; abort if any entry is not a
  // straightforward field projection.
  final Map<Integer, Integer> map = new HashMap<>();
  for (int source : interestingFields) {
    final RexNode rex = project.getProjects().get(source);
    if (!(rex instanceof RexInputRef)) {
      return null;
    }
    map.put(source, ((RexInputRef) rex).getIndex());
  }

  final ImmutableBitSet newGroupSet = aggregate.getGroupSet().permute(map);
  ImmutableList<ImmutableBitSet> newGroupingSets = null;
  if (aggregate.getGroupType() != Group.SIMPLE) {
    newGroupingSets =
        ImmutableBitSet.ORDERING.immutableSortedCopy(
            ImmutableBitSet.permute(aggregate.getGroupSets(), map));
  }

  final ImmutableList.Builder<AggregateCall> aggCalls =
      ImmutableList.builder();
  final int sourceCount = aggregate.getInput().getRowType().getFieldCount();
  final int targetCount = project.getInput().getRowType().getFieldCount();
  final Mappings.TargetMapping targetMapping =
      Mappings.target(map, sourceCount, targetCount);
  for (AggregateCall aggregateCall : aggregate.getAggCallList()) {
    aggCalls.add(aggregateCall.transform(targetMapping));
  }

  final Aggregate newAggregate =
      aggregate.copy(aggregate.getTraitSet(), project.getInput(),
          newGroupSet, newGroupingSets, aggCalls.build());

  // Add a project if the group set is not in the same order or
  // contains duplicates.
  final RelBuilder relBuilder = call.builder();
  relBuilder.push(newAggregate);
  final List<Integer> newKeys =
      Lists.transform(aggregate.getGroupSet().asList(), map::get);
  if (!newKeys.equals(newGroupSet.asList())) {
    final List<Integer> posList = new ArrayList<>();
    for (int newKey : newKeys) {
      posList.add(newGroupSet.indexOf(newKey));
    }
    for (int i = newAggregate.getGroupCount();
         i < newAggregate.getRowType().getFieldCount(); i++) {
      posList.add(i);
    }
    relBuilder.project(relBuilder.fields(posList));
  }

  return relBuilder.build();
}
 

public void onMatch(RelOptRuleCall call) {
  final Aggregate topAgg = call.rel(0);
  final Aggregate bottomAgg = call.rel(1);
  if (topAgg.getGroupCount() > bottomAgg.getGroupCount()) {
    return;
  }

  final ImmutableBitSet bottomGroupSet = bottomAgg.getGroupSet();
  final Map<Integer, Integer> map = new HashMap<>();
  bottomGroupSet.forEach(v -> map.put(map.size(), v));
  for (int k : topAgg.getGroupSet()) {
    if (!map.containsKey(k)) {
      return;
    }
  }

  // top aggregate keys must be subset of lower aggregate keys
  final ImmutableBitSet topGroupSet = topAgg.getGroupSet().permute(map);
  if (!bottomGroupSet.contains(topGroupSet)) {
    return;
  }

  boolean hasEmptyGroup = topAgg.getGroupSets()
      .stream().anyMatch(n -> n.isEmpty());

  final List<AggregateCall> finalCalls = new ArrayList<>();
  for (AggregateCall topCall : topAgg.getAggCallList()) {
    if (!isAggregateSupported(topCall)
        || topCall.getArgList().size() == 0) {
      return;
    }
    // Make sure top aggregate argument refers to one of the aggregate
    int bottomIndex = topCall.getArgList().get(0) - bottomGroupSet.cardinality();
    if (bottomIndex >= bottomAgg.getAggCallList().size()
        || bottomIndex < 0) {
      return;
    }
    AggregateCall bottomCall = bottomAgg.getAggCallList().get(bottomIndex);
    // Should not merge if top agg with empty group keys and the lower agg
    // function is COUNT, because in case of empty input for lower agg,
    // the result is empty, if we merge them, we end up with 1 result with
    // 0, which is wrong.
    if (!isAggregateSupported(bottomCall)
        || (bottomCall.getAggregation() == SqlStdOperatorTable.COUNT
             && hasEmptyGroup)) {
      return;
    }
    SqlSplittableAggFunction splitter = Objects.requireNonNull(
        bottomCall.getAggregation().unwrap(SqlSplittableAggFunction.class));
    AggregateCall finalCall = splitter.merge(topCall, bottomCall);
    // fail to merge the aggregate call, bail out
    if (finalCall == null) {
      return;
    }
    finalCalls.add(finalCall);
  }

  // re-map grouping sets
  ImmutableList<ImmutableBitSet> newGroupingSets = null;
  if (topAgg.getGroupType() != Group.SIMPLE) {
    newGroupingSets =
        ImmutableBitSet.ORDERING.immutableSortedCopy(
            ImmutableBitSet.permute(topAgg.getGroupSets(), map));
  }

  final Aggregate finalAgg =
      topAgg.copy(topAgg.getTraitSet(), bottomAgg.getInput(), topGroupSet,
          newGroupingSets, finalCalls);
  call.transformTo(finalAgg);
}