summaryrefslogtreecommitdiff
path: root/src
diff options
context:
space:
mode:
Diffstat (limited to 'src')
-rw-r--r--src/mongo/db/pipeline/document_source_sort.cpp17
-rw-r--r--src/mongo/db/pipeline/pipeline_test.cpp232
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}]", "[]");
}