cryptoe commented on code in PR #15470: URL: https://github.com/apache/druid/pull/15470#discussion_r1542283346
########## extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/groupby/GroupByQueryKit.java: ########## @@ -164,39 +168,96 @@ public QueryDefinition makeQueryDefinition( partitionBoost ); - queryDefBuilder.add( - StageDefinition.builder(firstStageNumber + 1) - .inputs(new StageInputSpec(firstStageNumber)) - .signature(resultSignature) - .maxWorkerCount(maxWorkerCount) - .shuffleSpec( - shuffleSpecFactoryPostAggregation != null - ? shuffleSpecFactoryPostAggregation.build(resultClusterBy, false) - : null - ) - .processorFactory(new GroupByPostShuffleFrameProcessorFactory(queryToRun)) - ); + final ShuffleSpec nextShuffleWindowSpec = getShuffleSpecForNextWindow(originalQuery, maxWorkerCount); - if (doLimitOrOffset) { - final DefaultLimitSpec limitSpec = (DefaultLimitSpec) queryToRun.getLimitSpec(); + if (nextShuffleWindowSpec == null) { queryDefBuilder.add( - StageDefinition.builder(firstStageNumber + 2) - .inputs(new StageInputSpec(firstStageNumber + 1)) + StageDefinition.builder(firstStageNumber + 1) + .inputs(new StageInputSpec(firstStageNumber)) .signature(resultSignature) - .maxWorkerCount(1) - .shuffleSpec(null) // no shuffling should be required after a limit processor. - .processorFactory( - new OffsetLimitFrameProcessorFactory( - limitSpec.getOffset(), - limitSpec.isLimited() ? (long) limitSpec.getLimit() : null - ) + .maxWorkerCount(maxWorkerCount) + .shuffleSpec( + shuffleSpecFactoryPostAggregation != null + ? shuffleSpecFactoryPostAggregation.build(resultClusterBy, false) + : null ) + .processorFactory(new GroupByPostShuffleFrameProcessorFactory(queryToRun)) + ); + + if (doLimitOrOffset) { + final DefaultLimitSpec limitSpec = (DefaultLimitSpec) queryToRun.getLimitSpec(); + queryDefBuilder.add( + StageDefinition.builder(firstStageNumber + 2) + .inputs(new StageInputSpec(firstStageNumber + 1)) + .signature(resultSignature) + .maxWorkerCount(1) + .shuffleSpec(null) // no shuffling should be required after a limit processor. + .processorFactory( + new OffsetLimitFrameProcessorFactory( + limitSpec.getOffset(), + limitSpec.isLimited() ? (long) limitSpec.getLimit() : null + ) + ) + ); + } + } else { + final RowSignature stageSignature; + // sort the signature to make sure the prefix is aligned + stageSignature = QueryKitUtils.sortableSignature( + resultSignature, + nextShuffleWindowSpec.clusterBy().getColumns() + ); + + queryDefBuilder.add( + StageDefinition.builder(firstStageNumber + 1) + .inputs(new StageInputSpec(firstStageNumber)) + .signature(stageSignature) + .maxWorkerCount(maxWorkerCount) + .shuffleSpec(nextShuffleWindowSpec) Review Comment: I think the limit and offset should be applied on the grouping key. So it should be ``` shuffleSpecFactoryPostAggregation != null ? : null ``` Also we can actually short circuit the shuffle spec of the OffsetLimitProcessor to null since limit always gets applied on 1 worker and 1 partition. So we would be okay in case a window processor is the next stage since the data would already be sorted :) -- 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