diff options
Diffstat (limited to 'src')
-rw-r--r-- | src/mongo/db/pipeline/document_source_sort.cpp | 17 | ||||
-rw-r--r-- | src/mongo/db/pipeline/pipeline_test.cpp | 232 |
2 files changed, 249 insertions, 0 deletions
diff --git a/src/mongo/db/pipeline/document_source_sort.cpp b/src/mongo/db/pipeline/document_source_sort.cpp index dcf40ba110b..02290179439 100644 --- a/src/mongo/db/pipeline/document_source_sort.cpp +++ b/src/mongo/db/pipeline/document_source_sort.cpp @@ -169,6 +169,23 @@ Pipeline::SourceContainer::iterator DocumentSourceSort::doOptimizeAt( _sortExecutor->setLimit(*limit); } + if (std::next(itr) == container->end()) { + return container->end(); + } + + if (auto nextSort = dynamic_cast<DocumentSourceSort*>((*std::next(itr)).get())) { + // If subsequent $sort stage exists, optimize by erasing the initial one. + // Since $sort is not guaranteed to be stable, we can blindly remove the first $sort. + auto thisLim = _sortExecutor->getLimit(); + auto otherLim = nextSort->_sortExecutor->getLimit(); + // When coalescing subsequent $sort stages, retain the existing/lower limit. + if (thisLim && (!otherLim || otherLim > thisLim)) { + nextSort->_sortExecutor->setLimit(thisLim); + } + Pipeline::SourceContainer::iterator ret = std::next(itr); + container->erase(itr); + return ret; + } return std::next(itr); } diff --git a/src/mongo/db/pipeline/pipeline_test.cpp b/src/mongo/db/pipeline/pipeline_test.cpp index 43915889af4..239924ce42e 100644 --- a/src/mongo/db/pipeline/pipeline_test.cpp +++ b/src/mongo/db/pipeline/pipeline_test.cpp @@ -293,6 +293,238 @@ TEST(PipelineOptimizationTest, SortMatchProjSkipLimBecomesMatchTopKSortSkipProj) assertPipelineOptimizesAndSerializesTo(inputPipe, outputPipe, serializedPipe); } +TEST(PipelineOptimizationTest, IdenticalSortSortBecomesSort) { + std::string inputPipe = + "[{$sort: {a: 1}}" + ",{$sort: {a: 1}}" + "]"; + + std::string outputPipe = + "[{$sort: {sortKey: {a: 1}}}" + "]"; + + std::string serializedPipe = + "[{$sort: {a: 1}}" + "]"; + + assertPipelineOptimizesAndSerializesTo(inputPipe, outputPipe, serializedPipe); +} + +TEST(PipelineOptimizationTest, IdenticalSortSortSortBecomesSort) { + std::string inputPipe = + "[{$sort: {a: 1}}" + ",{$sort: {a: 1}}" + ",{$sort: {a: 1}}" + "]"; + + std::string outputPipe = + "[{$sort: {sortKey: {a: 1}}}" + "]"; + + std::string serializedPipe = + "[{$sort: {a: 1}}" + "]"; + + assertPipelineOptimizesAndSerializesTo(inputPipe, outputPipe, serializedPipe); +} + +TEST(PipelineOptimizationTest, NonIdenticalSortsOnlySortOnFinalKey) { + std::string inputPipe = + "[{$sort: {a: -1}}" + ",{$sort: {a: 1}}" + ",{$sort: {a: -1}}" + "]"; + + std::string outputPipe = + "[{$sort: {sortKey: {a: -1}}}" + "]"; + + std::string serializedPipe = + "[{$sort: {a: -1}}" + "]"; + + assertPipelineOptimizesAndSerializesTo(inputPipe, outputPipe, serializedPipe); +} + +TEST(PipelineOptimizationTest, SortSortLimitBecomesFinalKeyTopKSort) { + std::string inputPipe = + "[{$sort: {a: -1}}" + ",{$sort: {a: 1}}" + ",{$limit: 5}" + "]"; + + std::string outputPipe = + "[{$sort: {sortKey: {a: 1}, limit: 5}}" + "]"; + + std::string serializedPipe = + "[{$sort: {a: 1}}" + ",{$limit: 5}" + "]"; + + assertPipelineOptimizesAndSerializesTo(inputPipe, outputPipe, serializedPipe); +} + +TEST(PipelineOptimizationTest, SortSortSkipLimitBecomesTopKSortSkip) { + std::string inputPipe = + "[{$sort: {b: 1}}" + ",{$sort: {a: 1}}" + ",{$skip : 3}" + ",{$limit: 5}" + "]"; + + std::string outputPipe = + "[{$sort: {sortKey: {a: 1}, limit: 8}}" + ",{$skip: 3}" + "]"; + + std::string serializedPipe = + "[{$sort: {a: 1}}" + ",{$limit: 8}" + ",{$skip : 3}" + "]"; + + assertPipelineOptimizesAndSerializesTo(inputPipe, outputPipe, serializedPipe); +} + +TEST(PipelineOptimizationTest, SortLimitSortLimitBecomesTopKSort) { + std::string inputPipe = + "[{$sort: {a: 1}}" + ",{$limit: 12}" + ",{$sort: {a: 1}}" + ",{$limit: 20}" + "]"; + + std::string outputPipe = + "[{$sort: {sortKey: {a: 1}, limit: 12}}" + "]"; + + std::string serializedPipe = + "[{$sort: {a: 1}}" + ",{$limit: 12}" + "]"; + + assertPipelineOptimizesAndSerializesTo(inputPipe, outputPipe, serializedPipe); +} + +TEST(PipelineOptimizationTest, SortLimitSortRetainsLimit) { + std::string inputPipe = + "[{$sort: {a: 1}}" + ",{$limit: 12}" + ",{$sort: {a: 1}}" + "]"; + + std::string outputPipe = + "[{$sort: {sortKey: {a: 1}, limit: 12}}" + "]"; + + std::string serializedPipe = + "[{$sort: {a: 1}}" + ",{$limit: 12}" + "]"; + + assertPipelineOptimizesAndSerializesTo(inputPipe, outputPipe, serializedPipe); +} + +TEST(PipelineOptimizationTest, SortSortLimitRetainsLimit) { + std::string inputPipe = + "[{$sort: {a: 1}}" + ",{$sort: {a: 1}}" + ",{$limit: 20}" + "]"; + + std::string outputPipe = + "[{$sort: {sortKey: {a: 1}, limit: 20}}" + "]"; + + std::string serializedPipe = + "[{$sort: {a: 1}}" + ",{$limit: 20}" + "]"; + + assertPipelineOptimizesAndSerializesTo(inputPipe, outputPipe, serializedPipe); +} + +TEST(PipelineOptimizationTest, SortSortSortMatchProjSkipLimBecomesMatchTopKSortSkipProj) { + std::string inputPipe = + "[{$sort: {a: 1}}" + ",{$sort: {a: 1}}" + ",{$sort: {a: 1}}" + ",{$match: {a: 1}}" + ",{$project : {a: 1}}" + ",{$skip : 3}" + ",{$limit: 5}" + "]"; + + std::string outputPipe = + "[{$match: {a: {$eq: 1}}}" + ",{$sort: {sortKey: {a: 1}, limit: 8}}" + ",{$skip: 3}" + ",{$project: {_id: true, a: true}}" + "]"; + + std::string serializedPipe = + "[{$match: {a: 1}}" + ",{$sort: {a: 1}}" + ",{$limit: 8}" + ",{$skip : 3}" + ",{$project : {_id: true, a: true}}" + "]"; + + assertPipelineOptimizesAndSerializesTo(inputPipe, outputPipe, serializedPipe); +} + +TEST(PipelineOptimizationTest, NonIdenticalSortsBecomeFinalKeyTopKSort) { + std::string inputPipe = + "[{$sort: {a: -1}}" + ",{$sort: {b: -1}}" + ",{$sort: {b: 1}}" + ",{$sort: {a: 1}}" + ",{$limit: 7}" + ",{$project : {a: 1}}" + ",{$limit: 5}" + "]"; + + std::string outputPipe = + "[{$sort: {sortKey: {a: 1}, limit: 5}}" + ",{$project: {_id: true, a: true}}" + "]"; + + std::string serializedPipe = + "[{$sort: {a: 1}}" + ",{$limit: 5}" + ",{$project : {_id: true, a: true}}" + "]"; + + assertPipelineOptimizesAndSerializesTo(inputPipe, outputPipe, serializedPipe); +} + +TEST(PipelineOptimizationTest, SubsequentSortsMergeAndBecomeTopKSortWithFinalKeyAndLowestLimit) { + std::string inputPipe = + "[{$sort: {a: 1}}" + ",{$sort: {a: -1}}" + ",{$limit: 8}" + ",{$limit: 7}" + ",{$project : {a: 1}}" + ",{$unwind: {path: '$a'}}" + "]"; + + std::string outputPipe = + "[{$sort: {sortKey: {a: -1}, limit: 7}}" + ",{$project: {_id: true, a: true}}" + ",{$unwind: {path: '$a'}}" + "]"; + + std::string serializedPipe = + "[{$sort: {a: -1}}" + ",{$limit: 7}" + ",{$project : {_id: true, a: true}}" + ",{$unwind: {path: '$a'}}" + "]"; + + assertPipelineOptimizesAndSerializesTo(inputPipe, outputPipe, serializedPipe); +} + TEST(PipelineOptimizationTest, RemoveSkipZero) { assertPipelineOptimizesTo("[{$skip: 0}]", "[]"); } |