enable CalciteArraysQueryTest for MSQ

LakshSingla · Oct 3, 2023 · ad5c590 · ad5c590
1 parent ffafadf
commit ad5c590
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 20 deletions.
diff --git a/...test/CalciteArraysSelectQueryMSQTest.java → ...d/msq/test/CalciteArraysQueryMSQTest.java b/...test/CalciteArraysSelectQueryMSQTest.java → ...d/msq/test/CalciteArraysQueryMSQTest.java
@@ -36,7 +36,7 @@
 /**
  * Runs {@link CalciteArraysQueryTest} but with MSQ engine
  */
-public class CalciteArraysSelectQueryMSQTest extends CalciteArraysQueryTest
+public class CalciteArraysQueryMSQTest extends CalciteArraysQueryTest
 {
   private TestGroupByBuffers groupByBuffers;
 

diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/CalciteArraysQueryTest.java b/sql/src/test/java/org/apache/druid/sql/calcite/CalciteArraysQueryTest.java
@@ -85,7 +85,6 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
   @Test
   public void testSelectConstantArrayExpressionFromTable()
   {
-    notMsqCompatible();
     testQuery(
         "SELECT ARRAY[1,2] as arr, dim1 FROM foo LIMIT 1",
         ImmutableList.of(
@@ -167,7 +166,8 @@ public void testSelectNonConstantArrayExpressionFromTable()
   @Test
   public void testSelectNonConstantArrayExpressionFromTableForMultival()
   {
-    notMsqCompatible();
+    // Produces nested string array, that MSQ can't infer from the selector
+    msqIncompatible();
     final String sql = "SELECT ARRAY[CONCAT(dim3, 'word'),'up'] as arr, dim1 FROM foo LIMIT 5";
     final Query<?> scanQuery = newScanQueryBuilder()
         .dataSource(CalciteTests.DATASOURCE1)
@@ -207,7 +207,6 @@ public void testSomeArrayFunctionsWithScanQuery()
     // Yes these outputs are strange sometimes, arrays are in a partial state of existence so end up a bit
     // stringy for now this is because virtual column selectors are coercing values back to stringish so that
     // multi-valued string dimensions can be grouped on.
-    notMsqCompatible();
     List<Object[]> expectedResults;
     if (useDefault) {
       expectedResults = ImmutableList.of(
@@ -387,7 +386,6 @@ public void testSomeArrayFunctionsWithScanQueryNoStringify()
     // which will still always be stringified to ultimately adhere to the varchar type
     // as array support increases in the engine this will likely change since using explict array functions should
     // probably kick it into an array
-    notMsqCompatible();
     List<Object[]> expectedResults;
     if (useDefault) {
       expectedResults = ImmutableList.of(
@@ -1019,7 +1017,6 @@ public void testArrayOffset()
   @Test
   public void testArrayGroupAsLongArray()
   {
-    notMsqCompatible();
     // Cannot vectorize as we donot have support in native query subsytem for grouping on arrays
     cannotVectorize();
     testQuery(
@@ -1071,7 +1068,6 @@ public void testArrayGroupAsDoubleArray()
   {
     // Cannot vectorize as we donot have support in native query subsytem for grouping on arrays as keys
     cannotVectorize();
-    notMsqCompatible();
     testQuery(
         "SELECT ARRAY[d1], SUM(cnt) FROM druid.numfoo GROUP BY 1 ORDER BY 2 DESC",
         QUERY_CONTEXT_NO_STRINGIFY_ARRAY,
@@ -1119,7 +1115,6 @@ public void testArrayGroupAsDoubleArray()
   @Test
   public void testArrayGroupAsFloatArray()
   {
-    notMsqCompatible();
     // Cannot vectorize as we donot have support in native query subsytem for grouping on arrays as keys
     cannotVectorize();
     testQuery(
@@ -1610,7 +1605,6 @@ public void testArrayAggMultiValue()
   @Test
   public void testArrayAggNumeric()
   {
-    notMsqCompatible();
     cannotVectorize();
     testQuery(
         "SELECT ARRAY_AGG(l1), ARRAY_AGG(DISTINCT l1), ARRAY_AGG(d1), ARRAY_AGG(DISTINCT d1), ARRAY_AGG(f1), ARRAY_AGG(DISTINCT f1) FROM numfoo",
@@ -1747,7 +1741,6 @@ public void testArrayAggNumeric()
   @Test
   public void testArrayAggQuantile()
   {
-    notMsqCompatible();
     cannotVectorize();
     testQuery(
         "SELECT ARRAY_QUANTILE(ARRAY_AGG(l1), 0.9) FROM numfoo",
@@ -1791,7 +1784,9 @@ public void testArrayAggQuantile()
   @Test
   public void testArrayAggArrays()
   {
-    notMsqCompatible();
+    // Produces nested array - ARRAY<ARRAY<LONG>>, which frame writers don't support. A way to get this query
+    // to run would be to use nested columns.
+    msqIncompatible();
     cannotVectorize();
     testQuery(
         "SELECT ARRAY_AGG(ARRAY[l1, l2]), ARRAY_AGG(DISTINCT ARRAY[l1, l2]) FROM numfoo",
@@ -1888,7 +1883,6 @@ public void testArrayAggArrays()
   @Test
   public void testArrayConcatAggArrays()
   {
-    notMsqCompatible();
     cannotVectorize();
     testQuery(
         "SELECT ARRAY_CONCAT_AGG(ARRAY[l1, l2]), ARRAY_CONCAT_AGG(DISTINCT ARRAY[l1, l2]) FROM numfoo",
@@ -2037,7 +2031,6 @@ public void testArrayAggExpression()
   public void testArrayAggMaxBytes()
   {
     cannotVectorize();
-    notMsqCompatible();
     testQuery(
         "SELECT ARRAY_AGG(l1, 128), ARRAY_AGG(DISTINCT l1, 128) FROM numfoo",
         ImmutableList.of(
@@ -2237,7 +2230,6 @@ public void testArrayAggGroupByArrayAggFromSubquery()
   @Test
   public void testArrayAggGroupByArrayAggOfLongsFromSubquery()
   {
-    notMsqCompatible();
     requireMergeBuffers(3);
     cannotVectorize();
     testQuery(
@@ -2377,7 +2369,6 @@ public void testArrayAggGroupByArrayAggOfStringsFromSubquery()
   @Test
   public void testArrayAggGroupByArrayAggOfDoubleFromSubquery()
   {
-    notMsqCompatible();
     requireMergeBuffers(3);
     cannotVectorize();
     testQuery(
@@ -2895,7 +2886,6 @@ public void testUnnestTwiceWithFiltersAndExpressions()
   @Test
   public void testUnnestThriceWithFiltersOnDimAndUnnestCol()
   {
-    notMsqCompatible();
     cannotVectorize();
     String sql = "    SELECT dimZipf, dim3_unnest1, dim3_unnest2, dim3_unnest3 FROM \n"
                  + "      ( SELECT * FROM \n"
@@ -2994,7 +2984,6 @@ public void testUnnestThriceWithFiltersOnDimAndUnnestCol()
   @Test
   public void testUnnestThriceWithFiltersOnDimAndAllUnnestColumns()
   {
-    notMsqCompatible();
     cannotVectorize();
     String sql = "    SELECT dimZipf, dim3_unnest1, dim3_unnest2, dim3_unnest3 FROM \n"
                  + "      ( SELECT * FROM \n"

diff --git a/sql/src/test/java/org/apache/druid/sql/calcite/QueryTestRunner.java b/sql/src/test/java/org/apache/druid/sql/calcite/QueryTestRunner.java
@@ -321,10 +321,15 @@ public static Pair<RowSignature, List<Object[]>> getResults(
   public static class VerifyResults implements QueryVerifyStep
   {
     protected final BaseExecuteQuery execStep;
+    protected final boolean verifyRowSignature;
 
-    public VerifyResults(BaseExecuteQuery execStep)
+    public VerifyResults(
+        BaseExecuteQuery execStep,
+        boolean verifyRowSignature
+    )
     {
       this.execStep = execStep;
+      this.verifyRowSignature = verifyRowSignature;
     }
 
     @Override
@@ -346,7 +351,9 @@ private void verifyResults(QueryResults queryResults)
       }
 
       QueryTestBuilder builder = execStep.builder();
-      builder.expectedResultsVerifier.verifyRowSignature(queryResults.signature);
+      if (verifyRowSignature) {
+        builder.expectedResultsVerifier.verifyRowSignature(queryResults.signature);
+      }
       builder.expectedResultsVerifier.verify(builder.sql, results);
     }
   }
@@ -667,7 +674,9 @@ public QueryTestRunner(QueryTestBuilder builder)
         verifySteps.add(new VerifyNativeQueries(finalExecStep));
       }
       if (builder.expectedResultsVerifier != null) {
-        verifySteps.add(new VerifyResults(finalExecStep));
+        // Don't verify the row signature when MSQ is running, since the broker receives the task id, and the signature
+        // would be {TASK:STRING} instead of the expected results signature
+        verifySteps.add(new VerifyResults(finalExecStep, !config.isRunningMSQ()));
       }
 
       if (!builder.customVerifications.isEmpty()) {