Apache
/
druid
mirror of https://github.com/apache/druid.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

MSQ window functions: Fix boost column not being written to the frame in window stage (#17155)

This commit is contained in:
Akshat Jain 2024-09-25 16:31:37 +05:30 committed by GitHub
parent 627752922c
commit 986bc62b88
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 156 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryFrameProcessor.java
View File

@ -138,6 +138,7 @@ public class WindowOperatorQueryFrameProcessor implements FrameProcessor<Object>
if (segmentGranularityVirtualColumn != null) {
frameWriterVirtualColumns.add(segmentGranularityVirtualColumn);
}
frameWriterVirtualColumns.add(this.partitionBoostVirtualColumn);
this.frameWriterVirtualColumns = VirtualColumns.create(frameWriterVirtualColumns);
}

99
extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQDrillWindowQueryTest.java
View File

@ -167,6 +167,105 @@ public class MSQDrillWindowQueryTest extends DrillWindowQueryTest
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.

70
extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQWindowTest.java
View File

@ -1242,20 +1242,20 @@ public class MSQWindowTest extends MSQTestBase
.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 class MSQWindowTest extends MSQTestBase
.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 class MSQWindowTest extends MSQTestBase
// 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();
}