From fa4cab405fae1cc986306f865508e9e04c577d8c Mon Sep 17 00:00:00 2001
From: Clint Wylie <cwylie@apache.org>
Date: Mon, 13 Feb 2023 18:27:23 -0800
Subject: [PATCH] fix bug with sql planner when virtual column capabilities are
 null (#13797)

---
 .../groupby/NestedDataGroupByQueryTest.java   | 26 +++++++-------
 .../calcite/rel/VirtualColumnRegistry.java    | 15 +++++---
 .../calcite/CalciteNestedDataQueryTest.java   | 35 +++++++++++++++++++
 3 files changed, 58 insertions(+), 18 deletions(-)

diff --git a/processing/src/test/java/org/apache/druid/query/groupby/NestedDataGroupByQueryTest.java b/processing/src/test/java/org/apache/druid/query/groupby/NestedDataGroupByQueryTest.java
index 079307eadf7..3038ccaedf5 100644
--- a/processing/src/test/java/org/apache/druid/query/groupby/NestedDataGroupByQueryTest.java
+++ b/processing/src/test/java/org/apache/druid/query/groupby/NestedDataGroupByQueryTest.java
@@ -579,6 +579,7 @@ public class NestedDataGroupByQueryTest extends InitializedNullHandlingTest
             "java.lang.UnsupportedOperationException: GroupBy v1 does not support dimension selectors with unknown cardinality.",
             t.getMessage()
         );
+        return;
       } else if (vectorize == QueryContexts.Vectorize.FORCE) {
         Throwable t = Assert.assertThrows(
             RuntimeException.class,
@@ -591,20 +592,19 @@ public class NestedDataGroupByQueryTest extends InitializedNullHandlingTest
         );
         return;
       }
-    } else {
-
-      Sequence<ResultRow> seq = helper.runQueryOnSegmentsObjs(
-          segmentsGenerator.apply(helper, tempFolder, closer),
-          groupQuery
-      );
-
-      List<ResultRow> results = seq.toList();
-      verifyResults(
-          groupQuery.getResultRowSignature(),
-          results,
-          expectedResults
-      );
     }
+
+    Sequence<ResultRow> seq = helper.runQueryOnSegmentsObjs(
+        segmentsGenerator.apply(helper, tempFolder, closer),
+        groupQuery
+    );
+
+    List<ResultRow> results = seq.toList();
+    verifyResults(
+        groupQuery.getResultRowSignature(),
+        results,
+        expectedResults
+    );
   }
 
   private static void verifyResults(RowSignature rowSignature, List<ResultRow> results, List<Object[]> expected)
diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/rel/VirtualColumnRegistry.java b/sql/src/main/java/org/apache/druid/sql/calcite/rel/VirtualColumnRegistry.java
index b3464b9bd63..0e8e8f7ae18 100644
--- a/sql/src/main/java/org/apache/druid/sql/calcite/rel/VirtualColumnRegistry.java
+++ b/sql/src/main/java/org/apache/druid/sql/calcite/rel/VirtualColumnRegistry.java
@@ -23,6 +23,7 @@ import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.druid.math.expr.ExprMacroTable;
 import org.apache.druid.segment.VirtualColumn;
+import org.apache.druid.segment.column.ColumnCapabilities;
 import org.apache.druid.segment.column.ColumnType;
 import org.apache.druid.segment.column.RowSignature;
 import org.apache.druid.sql.calcite.expression.DruidExpression;
@@ -180,15 +181,19 @@ public class VirtualColumnRegistry
     for (Map.Entry<String, ExpressionAndTypeHint> virtualColumn : virtualColumnsByName.entrySet()) {
       final String columnName = virtualColumn.getKey();
 
+      final ColumnType typeHint = virtualColumn.getValue().getTypeHint();
       // this is expensive, maybe someday it could use the typeHint, or the inferred type, but for now use native
       // expression type inference
+      final ColumnCapabilities virtualCapabilities = virtualColumn.getValue().getExpression().toVirtualColumn(
+          columnName,
+          typeHint,
+          macroTable
+      ).capabilities(baseSignature, columnName);
+
+      // fall back to type hint
       builder.add(
           columnName,
-          virtualColumn.getValue().getExpression().toVirtualColumn(
-              columnName,
-              virtualColumn.getValue().getTypeHint(),
-              macroTable
-          ).capabilities(baseSignature, columnName).toColumnType()
+          virtualCapabilities != null ? virtualCapabilities.toColumnType() : typeHint
       );
     }
 
diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteNestedDataQueryTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteNestedDataQueryTest.java
index 36464032e36..692f3830e52 100644
--- a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteNestedDataQueryTest.java
+++ b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteNestedDataQueryTest.java
@@ -780,6 +780,41 @@ public class CalciteNestedDataQueryTest extends BaseCalciteQueryTest
     );
   }
 
+  @Test
+  public void testGroupByRootSingleTypeStringMixed2SparseJsonValueNonExistentPath()
+  {
+    testQuery(
+        "SELECT "
+        + "JSON_VALUE(string_sparse, '$[1]'), "
+        + "SUM(cnt) "
+        + "FROM druid.nested_mix_2 GROUP BY 1",
+        ImmutableList.of(
+            GroupByQuery.builder()
+                        .setDataSource(DATA_SOURCE_MIXED_2)
+                        .setInterval(querySegmentSpec(Filtration.eternity()))
+                        .setGranularity(Granularities.ALL)
+                        .setDimensions(
+                            dimensions(
+                                new DefaultDimensionSpec("v0", "d0")
+                            )
+                        )
+                        .setVirtualColumns(
+                            new NestedFieldVirtualColumn("string_sparse", "$[1]", "v0", ColumnType.STRING)
+                        )
+                        .setAggregatorSpecs(aggregators(new LongSumAggregatorFactory("a0", "cnt")))
+                        .setContext(QUERY_CONTEXT_DEFAULT)
+                        .build()
+        ),
+        ImmutableList.of(
+            new Object[]{NullHandling.defaultStringValue(), 14L}
+        ),
+        RowSignature.builder()
+                    .add("EXPR$0", ColumnType.STRING)
+                    .add("EXPR$1", ColumnType.LONG)
+                    .build()
+    );
+  }
+
   @Test
   public void testGroupByJsonValues()
   {