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

… in window stage (apache#17141) In WindowOperatorQueryFrameProcessor, currently we were not adding the partition boost column to the ColumnSelectorFactory used to create the FrameWriter. Because of this, the boost column was not being written to the frame.
kfaraz · Sep 25, 2024 · 446ffc3 · 446ffc3
1 parent 464b153
commit 446ffc3
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
@@ -173,6 +173,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();
   }