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 (#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.
This commit is contained in:
Akshat Jain 2024-09-25 09:50:42 +05:30 committed by GitHub
parent 464b153850
commit 446ffc325c
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

@ -173,6 +173,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();
}