From 986bc62b8898595e8fc8169a7c0c4e52f9a52cab Mon Sep 17 00:00:00 2001 From: Akshat Jain Date: Wed, 25 Sep 2024 16:31:37 +0530 Subject: [PATCH] MSQ window functions: Fix boost column not being written to the frame in window stage (#17155) --- .../WindowOperatorQueryFrameProcessor.java | 1 + .../msq/exec/MSQDrillWindowQueryTest.java | 99 +++++++++++++++++++ .../apache/druid/msq/exec/MSQWindowTest.java | 70 ++++++++++--- 3 files changed, 156 insertions(+), 14 deletions(-) diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java index aab8f1f1a6bb..a0572a91b4df 100644 --- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java +++ b/extensions-core/multi-stage-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/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQDrillWindowQueryTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQDrillWindowQueryTest.java index 3e8e2796eeac..a6e1e5ea3acf 100644 --- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQDrillWindowQueryTest.java +++ b/extensions-core/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 index 2ba9d56ac69a..486014101c5d 100644 --- 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(); }