MSQ window functions: Fix boost column not being written to the frame…

… in window stage (apache#17155)
clintropolis · Sep 25, 2024 · 986bc62 · 986bc62
1 parent 6277529
commit 986bc62
Show file tree

Hide file tree

Showing 3 changed files with 156 additions and 14 deletions.
diff --git a/...-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java b/...-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java
@@ -138,6 +138,7 @@ public WindowOperatorQueryFrameProcessor(
     if (segmentGranularityVirtualColumn != null) {
       frameWriterVirtualColumns.add(segmentGranularityVirtualColumn);
     }
+    frameWriterVirtualColumns.add(this.partitionBoostVirtualColumn);
     this.frameWriterVirtualColumns = VirtualColumns.create(frameWriterVirtualColumns);
   }
 

diff --git a/...re/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQDrillWindowQueryTest.java b/...re/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQDrillWindowQueryTest.java
@@ -167,6 +167,105 @@ public void test_over_clause_with_only_partitioning_multiple_over_different_part
     windowQueryTest();
   }
 
+  @Override
+  @DrillTest("ntile_func/ntileFn_47")
+  @Test
+  public void test_ntile_func_ntileFn_47()
+  {
+    useSingleWorker();
+    windowQueryTest();
+  }
+
+  @Override
+  @DrillTest("ntile_func/ntileFn_49")
+  @Test
+  public void test_ntile_func_ntileFn_49()
+  {
+    useSingleWorker();
+    windowQueryTest();
+  }
+
+  @Override
+  @DrillTest("ntile_func/ntileFn_50")
+  @Test
+  public void test_ntile_func_ntileFn_50()
+  {
+    useSingleWorker();
+    windowQueryTest();
+  }
+
+  @Override
+  @DrillTest("ntile_func/ntileFn_51")
+  @Test
+  public void test_ntile_func_ntileFn_51()
+  {
+    useSingleWorker();
+    windowQueryTest();
+  }
+
+  @Override
+  @DrillTest("ntile_func/ntileFn_52")
+  @Test
+  public void test_ntile_func_ntileFn_52()
+  {
+    useSingleWorker();
+    windowQueryTest();
+  }
+
+  @Override
+  @DrillTest("ntile_func/ntileFn_53")
+  @Test
+  public void test_ntile_func_ntileFn_53()
+  {
+    useSingleWorker();
+    windowQueryTest();
+  }
+
+  @Override
+  @DrillTest("ntile_func/ntileFn_54")
+  @Test
+  public void test_ntile_func_ntileFn_54()
+  {
+    useSingleWorker();
+    windowQueryTest();
+  }
+
+  @Override
+  @DrillTest("ntile_func/ntileFn_55")
+  @Test
+  public void test_ntile_func_ntileFn_55()
+  {
+    useSingleWorker();
+    windowQueryTest();
+  }
+
+  @Override
+  @DrillTest("ntile_func/ntileFn_56")
+  @Test
+  public void test_ntile_func_ntileFn_56()
+  {
+    useSingleWorker();
+    windowQueryTest();
+  }
+
+  @Override
+  @DrillTest("ntile_func/ntileFn_57")
+  @Test
+  public void test_ntile_func_ntileFn_57()
+  {
+    useSingleWorker();
+    windowQueryTest();
+  }
+
+  @Override
+  @DrillTest("ntile_func/ntileFn_58")
+  @Test
+  public void test_ntile_func_ntileFn_58()
+  {
+    useSingleWorker();
+    windowQueryTest();
+  }
+
   /*
   Queries having window functions can give multiple correct results because of using MixShuffleSpec in the previous stage.
   So we want to use a single worker to get the same result everytime for such test cases.

diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQWindowTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQWindowTest.java
@@ -1242,20 +1242,20 @@ public void testWindowOnFooWithDim2()
         .setExpectedResultRows(
             NullHandling.replaceWithDefault() ?
             ImmutableList.of(
-                new Object[]{"a", 5.0},
+                new Object[]{"", 11.0},
                 new Object[]{"", 11.0},
                 new Object[]{"", 11.0},
                 new Object[]{"a", 5.0},
-                new Object[]{"abc", 5.0},
-                new Object[]{"", 11.0}
+                new Object[]{"a", 5.0},
+                new Object[]{"abc", 5.0}
             ) :
             ImmutableList.of(
-                new Object[]{"a", 5.0},
+                new Object[]{null, 8.0},
                 new Object[]{null, 8.0},
                 new Object[]{"", 3.0},
                 new Object[]{"a", 5.0},
-                new Object[]{"abc", 5.0},
-                new Object[]{null, 8.0}
+                new Object[]{"a", 5.0},
+                new Object[]{"abc", 5.0}
             ))
         .setQueryContext(DEFAULT_MSQ_CONTEXT)
         .verifyResults();
@@ -1891,11 +1891,11 @@ public void testSelectWithWikipediaWithPartitionKeyNotInSelect()
                                    .build())
         .setExpectedRowSignature(rowSignature)
         .setExpectedResultRows(ImmutableList.of(
-            new Object[]{"Auburn", 0L, 1698L},
-            new Object[]{"Mexico City", 0L, 6136L},
-            new Object[]{"Seoul", 663L, 5582L},
-            new Object[]{"Tokyo", 0L, 12615L},
-            new Object[]{"Santiago", 161L, 401L}
+            new Object[]{"Al Ain", 8L, 6334L},
+            new Object[]{"Dubai", 3L, 6334L},
+            new Object[]{"Dubai", 6323L, 6334L},
+            new Object[]{"Tirana", 26L, 26L},
+            new Object[]{"Benguela", 0L, 0L}
         ))
         .setQueryContext(DEFAULT_MSQ_CONTEXT)
         .verifyResults();
@@ -2225,17 +2225,59 @@ public void testWindowOnMixOfEmptyAndNonEmptyOverWithMultipleWorkers()
 
         // Stage 3, Worker 0
         .setExpectedCountersForStageWorkerChannel(
-            CounterSnapshotMatcher.with().rows(13).bytes(1327).frames(1),
+            CounterSnapshotMatcher.with().rows(3).bytes(318).frames(1),
             3, 0, "input0"
         )
         .setExpectedCountersForStageWorkerChannel(
-            CounterSnapshotMatcher.with().rows(13).bytes(1379).frames(1),
+            CounterSnapshotMatcher.with().rows(3).bytes(330).frames(1),
             3, 0, "output"
         )
         .setExpectedCountersForStageWorkerChannel(
-            CounterSnapshotMatcher.with().rows(13).bytes(1327).frames(1),
+            CounterSnapshotMatcher.with().rows(3).bytes(318).frames(1),
             3, 0, "shuffle"
         )
+
+        // Stage 3, Worker 1
+        .setExpectedCountersForStageWorkerChannel(
+            CounterSnapshotMatcher.with().rows(0, 3).bytes(0, 333).frames(0, 1),
+            3, 1, "input0"
+        )
+        .setExpectedCountersForStageWorkerChannel(
+            CounterSnapshotMatcher.with().rows(3).bytes(345).frames(1),
+            3, 1, "output"
+        )
+        .setExpectedCountersForStageWorkerChannel(
+            CounterSnapshotMatcher.with().rows(3).bytes(333).frames(1),
+            3, 1, "shuffle"
+        )
+
+        // Stage 3, Worker 2
+        .setExpectedCountersForStageWorkerChannel(
+            CounterSnapshotMatcher.with().rows(0, 0, 3).bytes(0, 0, 352).frames(0, 0, 1),
+            3, 2, "input0"
+        )
+        .setExpectedCountersForStageWorkerChannel(
+            CounterSnapshotMatcher.with().rows(3).bytes(364).frames(1),
+            3, 2, "output"
+        )
+        .setExpectedCountersForStageWorkerChannel(
+            CounterSnapshotMatcher.with().rows(3).bytes(352).frames(1),
+            3, 2, "shuffle"
+        )
+
+        // Stage 3, Worker 3
+        .setExpectedCountersForStageWorkerChannel(
+            CounterSnapshotMatcher.with().rows(0, 0, 0, 4).bytes(0, 0, 0, 426).frames(0, 0, 0, 1),
+            3, 3, "input0"
+        )
+        .setExpectedCountersForStageWorkerChannel(
+            CounterSnapshotMatcher.with().rows(4).bytes(442).frames(1),
+            3, 3, "output"
+        )
+        .setExpectedCountersForStageWorkerChannel(
+            CounterSnapshotMatcher.with().rows(4).bytes(426).frames(1),
+            3, 3, "shuffle"
+        )
         .verifyResults();
   }