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: commits-unsubscr...@druid.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@druid.apache.org For additional commands, e-mail: commits-h...@druid.apache.org