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

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

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++;

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
   {