public RelNode align(Aggregate rel, List<RelFieldCollation> collations) {
    // 1) We extract the group by positions that are part of the collations and
    // sort them so they respect it
    LinkedHashSet<Integer> aggregateColumnsOrder = new LinkedHashSet<>();
    ImmutableList.Builder<RelFieldCollation> propagateCollations = ImmutableList.builder();
    if (rel.getGroupType() == Group.SIMPLE && !collations.isEmpty()) {
      for (RelFieldCollation c : collations) {
        if (c.getFieldIndex() < rel.getGroupCount()) {
          // Group column found
          if (aggregateColumnsOrder.add(c.getFieldIndex())) {
            propagateCollations.add(c.copy(rel.getGroupSet().nth(c.getFieldIndex())));
          }
        }
      }
    }
    for (int i = 0; i < rel.getGroupCount(); i++) {
      if (!aggregateColumnsOrder.contains(i)) {
        // Not included in the input collations, but can be propagated as this Aggregate
        // will enforce it
        propagateCollations.add(new RelFieldCollation(rel.getGroupSet().nth(i)));
      }
    }

    // 2) We propagate
    final RelNode child = dispatchAlign(rel.getInput(), propagateCollations.build());

    // 3) We annotate the Aggregate operator with this info
    final HiveAggregate newAggregate = (HiveAggregate) rel.copy(rel.getTraitSet(),
            ImmutableList.of(child));
    newAggregate.setAggregateColumnsOrder(aggregateColumnsOrder);
    return newAggregate;
  }
Execute the logic in this class. In particular, make a top-down traversal of the tree
and annotate and recreate appropriate operators.
ql/src/java/org/apache/hadoop/hive/ql/optimizer/calcite/rules/HiveRelColumnsAlignment.java