[GitHub] [druid] ccaominh commented on a change in pull request #9122: Add SQL GROUPING SETS support.

[GitBox <gi...@apache.org>]

ccaominh commented on a change in pull request #9122: Add SQL GROUPING SETS support.
URL: https://github.com/apache/druid/pull/9122#discussion_r384192373
 
 

 ##########
 File path: sql/src/main/java/org/apache/druid/sql/calcite/rel/Grouping.java
 ##########
 @@ -141,36 +159,95 @@ public RowSignature getOutputRowSignature()
     return outputRowSignature;
   }
 
+  /**
+   * Applies a post-grouping projection.
+   *
+   * @see DruidQuery#computeGrouping which uses this
+   */
+  public Grouping applyProject(final PlannerContext plannerContext, final Project project)
+  {
+    final List<DimensionExpression> newDimensions = new ArrayList<>();
+    final List<Aggregation> newAggregations = new ArrayList<>(aggregations);
+    final Subtotals newSubtotals;
+
+    final Projection postAggregationProjection = Projection.postAggregation(
+        project,
+        plannerContext,
+        outputRowSignature,
+        "p"
+    );
+
+    postAggregationProjection.getPostAggregators().forEach(
+        postAggregator -> newAggregations.add(Aggregation.create(postAggregator))
+    );
+
+    // Remove literal dimensions that did not appear in the projection. This is useful for queries
+    // like "SELECT COUNT(*) FROM tbl GROUP BY 'dummy'" which some tools can generate, and for which we don't
+    // actually want to include a dimension 'dummy'.
+    final ImmutableBitSet aggregateProjectBits = RelOptUtil.InputFinder.bits(project.getChildExps(), null);
+    final int[] newDimIndexes = new int[dimensions.size()];
+
+    for (int i = 0; i < dimensions.size(); i++) {
+      final DimensionExpression dimension = dimensions.get(i);
+      if (Parser.parse(dimension.getDruidExpression().getExpression(), plannerContext.getExprMacroTable())
+                .isLiteral() && !aggregateProjectBits.get(i)) {
+        newDimIndexes[i] = -1;
+      } else {
+        newDimIndexes[i] = newDimensions.size();
+        newDimensions.add(dimension);
+      }
+    }
+
+    // Renumber subtotals, if needed, to account for removed dummy dimensions.
+    if (newDimensions.size() != dimensions.size()) {
+      final List<IntList> newSubtotalsList = new ArrayList<>();
+
+      for (IntList subtotal : subtotals.getSubtotals()) {
+        final IntList newSubtotal = new IntArrayList();
+        for (int dimIndex : subtotal) {
+          final int newDimIndex = newDimIndexes[dimIndex];
+          if (newDimIndex >= 0) {
+            newSubtotal.add(newDimIndex);
+          }
+        }
+
+        newSubtotalsList.add(newSubtotal);
+      }
+
+      newSubtotals = new Subtotals(newSubtotalsList);
+    } else {
+      newSubtotals = subtotals;
+    }
+
+    return Grouping.create(
+        newDimensions,
+        newSubtotals,
+        newAggregations,
+        havingFilter,
+        postAggregationProjection.getOutputRowSignature()
+    );
+  }
+
   @Override
-  public boolean equals(final Object o)
+  public boolean equals(Object o)
 
 Review comment:
   Is there a test for this?

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services

---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@druid.apache.org
For additional commands, e-mail: commits-help@druid.apache.org