Akshat-Jain commented on code in PR #16804:
URL: https://github.com/apache/druid/pull/16804#discussion_r1698260880
##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/WindowOperatorQueryKit.java:
##########
@@ -328,4 +321,29 @@ private ShuffleSpec
findShuffleSpecForNextWindow(List<OperatorFactory> operatorF
return new HashShuffleSpec(new ClusterBy(keyColsOfWindow, 0),
maxWorkerCount);
}
+
+ /**
+ * Override the shuffle spec of the last stage based on the shuffling
required by the first window stage.
+ * @param queryId
+ * @param dataSourcePlan
+ * @param shuffleSpec
+ * @return
+ */
+ private QueryDefinitionBuilder makeQueryDefinitionBuilder(String queryId,
DataSourcePlan dataSourcePlan, ShuffleSpec shuffleSpec)
+ {
+ final QueryDefinitionBuilder queryDefBuilder =
QueryDefinition.builder(queryId);
+ int previousStageNumber =
dataSourcePlan.getSubQueryDefBuilder().get().build().getFinalStageDefinition().getStageNumber();
+ for (final StageDefinition stageDef :
dataSourcePlan.getSubQueryDefBuilder().get().build().getStageDefinitions()) {
+ if (stageDef.getStageNumber() == previousStageNumber) {
+ RowSignature rowSignature = QueryKitUtils.sortableSignature(
+ stageDef.getSignature(),
+ shuffleSpec.clusterBy().getColumns()
+ );
+
queryDefBuilder.add(StageDefinition.builder(stageDef).shuffleSpec(shuffleSpec).signature(rowSignature));
Review Comment:
If we throw error when the shuffleSpec is non-null for the stage that's
getting overridden, a lot of tests fail in MSQDrillWindowQueryTest: 754 tests
failed, 114 tests passed.
We can't throw an error when it's non-null 😅
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
