From 9738a03c833ee6e7ba421413b811c2522a01d873 Mon Sep 17 00:00:00 2001
From: Maytas Monsereenusorn <52679095+maytasm@users.noreply.github.com>
Date: Thu, 4 Jun 2020 13:28:05 -1000
Subject: [PATCH] Fix groupBy with literal in subquery grouping (#9986)

* fix groupBy with literal in subquery grouping

* fix groupBy with literal in subquery grouping

* fix groupBy with literal in subquery grouping

* address comments

* update javadocs
---
 .../apache/druid/query/QueryToolChest.java    |  9 ++-
 .../druid/sql/calcite/rel/DruidQuery.java     | 10 ++-
 .../druid/sql/calcite/CalciteQueryTest.java   | 63 +++++++++++++++++++
 3 files changed, 75 insertions(+), 7 deletions(-)

diff --git a/processing/src/main/java/org/apache/druid/query/QueryToolChest.java b/processing/src/main/java/org/apache/druid/query/QueryToolChest.java
index 509738ee680..ee400f814fb 100644
--- a/processing/src/main/java/org/apache/druid/query/QueryToolChest.java
+++ b/processing/src/main/java/org/apache/druid/query/QueryToolChest.java
@@ -179,11 +179,10 @@ public abstract class QueryToolChest<ResultType, QueryType extends Query<ResultT
   );
 
   /**
-   * Generally speaking this is the exact same thing as makePreComputeManipulatorFn.  It is leveraged in
-   * order to compute PostAggregators on results after they have been completely merged together, which
-   * should actually be done in the mergeResults() call instead of here.
-   * <p>
-   * This should never actually be overridden and it should be removed as quickly as possible.
+   * Generally speaking this is the exact same thing as makePreComputeManipulatorFn.  It is leveraged in order to
+   * compute PostAggregators on results after they have been completely merged together. To minimize walks of segments,
+   * it is recommended to use mergeResults() call instead of this method if possible. However, this may not always be
+   * possible as we don’t always want to run PostAggregators and other stuff that happens there when you mergeResults.
    *
    * @param query The Query that is currently being processed
    * @param fn    The function that should be applied to all metrics in the results
diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/rel/DruidQuery.java b/sql/src/main/java/org/apache/druid/sql/calcite/rel/DruidQuery.java
index b25bfb82330..137f9b70bd3 100644
--- a/sql/src/main/java/org/apache/druid/sql/calcite/rel/DruidQuery.java
+++ b/sql/src/main/java/org/apache/druid/sql/calcite/rel/DruidQuery.java
@@ -191,7 +191,7 @@ public class DruidQuery
           computeGrouping(
               partialQuery,
               plannerContext,
-              computeOutputRowSignature(sourceRowSignature, selectProjection, null, null),
+              computeOutputRowSignature(sourceRowSignature, null, null, null),
               virtualColumnRegistry,
               rexBuilder,
               finalizeAggregations
@@ -433,7 +433,13 @@ public class DruidQuery
     final Aggregate aggregate = partialQuery.getAggregate();
 
     // dimBitMapping maps from input field position to group set position (dimension number).
-    final int[] dimBitMapping = new int[rowSignature.size()];
+    final int[] dimBitMapping;
+    if (partialQuery.getSelectProject() != null) {
+      dimBitMapping = new int[partialQuery.getSelectProject().getRowType().getFieldCount()];
+    } else {
+      dimBitMapping = new int[rowSignature.size()];
+    }
+
     int i = 0;
     for (int dimBit : aggregate.getGroupSet()) {
       dimBitMapping[dimBit] = i++;
diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java
index 8068a0d9e38..e3579aad6a9 100644
--- a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java
+++ b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java
@@ -13419,6 +13419,69 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
     );
   }
 
+  @Test
+  public void testGroupByWithLiteralInSubqueryGrouping() throws Exception
+  {
+    testQuery(
+        "SELECT \n"
+        + "   t1, t2\n"
+        + "  FROM\n"
+        + "   ( SELECT\n"
+        + "     'dummy' as t1,\n"
+        + "     CASE\n"
+        + "       WHEN \n"
+        + "         dim4 = 'b'\n"
+        + "       THEN dim4\n"
+        + "       ELSE NULL\n"
+        + "     END AS t2\n"
+        + "     FROM\n"
+        + "       numfoo\n"
+        + "     GROUP BY\n"
+        + "       dim4\n"
+        + "   )\n"
+        + " GROUP BY\n"
+        + "   t1,t2\n",
+        ImmutableList.of(
+            GroupByQuery.builder()
+                        .setDataSource(
+                            GroupByQuery.builder()
+                                        .setDataSource(CalciteTests.DATASOURCE3)
+                                        .setInterval(querySegmentSpec(Filtration.eternity()))
+                                        .setGranularity(Granularities.ALL)
+                                        .setDimensions(new DefaultDimensionSpec("dim4", "_d0", ValueType.STRING))
+                                        .setContext(QUERY_CONTEXT_DEFAULT)
+                                        .build()
+                        )
+                        .setVirtualColumns(
+                            expressionVirtualColumn(
+                                "v0",
+                                "\'dummy\'",
+                                ValueType.STRING
+                            ),
+                            expressionVirtualColumn(
+                                "v1",
+                                "case_searched((\"_d0\" == 'b'),\"_d0\",null)",
+                                ValueType.STRING
+                            )
+                        )
+                        .setInterval(querySegmentSpec(Filtration.eternity()))
+                        .setDimensions(
+                            dimensions(
+                                new DefaultDimensionSpec("v0", "d0", ValueType.STRING),
+                                new DefaultDimensionSpec("v1", "d1", ValueType.STRING)
+                            )
+                        )
+                        .setGranularity(Granularities.ALL)
+                        .setContext(QUERY_CONTEXT_DEFAULT)
+                        .build()
+        ),
+        ImmutableList.of(
+            new Object[]{"dummy", NULL_STRING},
+            new Object[]{"dummy", "b"}
+        )
+    );
+  }
+
   @Test
   public void testMultiValueStringWorksLikeStringGroupBy() throws Exception
   {