SERVER-74725 Avoid duplicate keydefs for sort

3 files changed, 46 insertions, 16 deletions

diff --git a/src/mongo/db/pipeline/document_source_densify.cpp b/src/mongo/db/pipeline/document_source_densify.cpp
index f035cff11f7..e67b05c812e 100644
--- a/src/mongo/db/pipeline/document_source_densify.cpp
+++ b/src/mongo/db/pipeline/document_source_densify.cpp
@@ -28,13 +28,13 @@
  */
 
 #include "mongo/db/pipeline/document_source_densify.h"
-#include "mongo/base/exact_cast.h"
 #include "mongo/db/pipeline/document_source_sort.h"
 #include "mongo/db/pipeline/field_path.h"
 #include "mongo/db/query/sort_pattern.h"
 #include "mongo/stdx/variant.h"
 #include "mongo/util/assert_util.h"
 #include "mongo/util/overloaded_visitor.h"
+#include "mongo/util/string_map.h"
 
 using boost::intrusive_ptr;
 using boost::optional;
@@ -190,10 +190,16 @@ SortPattern getSortPatternForDensify(RangeStatement rangeStatement,
         }
     }
 
-    // Add field path to sort spec.
-    SortPatternPart part;
-    part.fieldPath = field.fullPath();
-    sortParts.push_back(std::move(part));
+    // Add field path to sort spec if it is not yet in the sort spec.
+    const auto inserted = std::find_if(
+        sortParts.begin(), sortParts.end(), [&field](const SortPatternPart& s) -> bool {
+            return s.fieldPath->fullPath().compare(field.fullPath()) == 0;
+        });
+    if (inserted == sortParts.end()) {
+        SortPatternPart part;
+        part.fieldPath = field.fullPath();
+        sortParts.push_back(std::move(part));
+    }
     return SortPattern{sortParts};
 }
 
diff --git a/src/mongo/db/pipeline/document_source_internal_unpack_bucket.cpp b/src/mongo/db/pipeline/document_source_internal_unpack_bucket.cpp
index b29d43394f5..c7f9a6f6d64 100644
--- a/src/mongo/db/pipeline/document_source_internal_unpack_bucket.cpp
+++ b/src/mongo/db/pipeline/document_source_internal_unpack_bucket.cpp
@@ -168,18 +168,8 @@ boost::intrusive_ptr<DocumentSourceSort> createMetadataSortForReorder(
         : sort.getSortKeyPattern();
     std::vector<SortPattern::SortPatternPart> updatedPattern;
 
-    if (groupIdField) {
-        auto groupId = FieldPath(groupIdField.value());
-        SortPattern::SortPatternPart patternPart;
-        patternPart.isAscending = !flipSort;
-        patternPart.fieldPath = groupId;
-        updatedPattern.push_back(patternPart);
-    }
-
-
     for (const auto& entry : sortPattern) {
         updatedPattern.push_back(entry);
-
         if (lastpointTimeField && entry.fieldPath->fullPath() == lastpointTimeField.value()) {
             updatedPattern.back().fieldPath =
                 FieldPath((entry.isAscending ? timeseries::kControlMinFieldNamePrefix
@@ -199,6 +189,33 @@ boost::intrusive_ptr<DocumentSourceSort> createMetadataSortForReorder(
             updatedPattern.back().fieldPath = updated;
         }
     }
+    // After the modifications of the sortPattern are completed, for the lastPoint
+    // optimizations, the group field needs to be added to the beginning of the sort pattern.
+    // Do note that the modified sort pattern is for sorting within a group (within the bucket)
+    // and the plan is to do the grouping and sort within on go.
+    // If the group field is already in the sortPattern then it needs to moved to the first
+    // position. A flip in the later case is not necessary anymore as the sort order was
+    // already flipped.
+    // Example 1: $group: {a:1}, $sort{b: 1, a: -1} --> modifiedPattern: {a: -1, b: 1}
+    // Example 2: $group: {c:1}, $sort{d: -1, e: 1} --> modifiedPattern: {c: 1, d: -1, e: 1}
+    if (groupIdField) {
+        const auto groupId = FieldPath(groupIdField.value());
+        SortPattern::SortPatternPart patternPart;
+        patternPart.fieldPath = groupId;
+        const auto pattern =
+            std::find_if(updatedPattern.begin(),
+                         updatedPattern.end(),
+                         [&groupId](const SortPattern::SortPatternPart& s) -> bool {
+                             return s.fieldPath->fullPath().compare(groupId.fullPath()) == 0;
+                         });
+        if (pattern != updatedPattern.end()) {
+            patternPart.isAscending = pattern->isAscending;
+            updatedPattern.erase(pattern);
+        } else {
+            patternPart.isAscending = !flipSort;
+        }
+        updatedPattern.insert(updatedPattern.begin(), patternPart);
+    }
 
     boost::optional<uint64_t> maxMemoryUsageBytes;
     if (auto sortStatsPtr = dynamic_cast<const SortStats*>(sort.getSpecificStats())) {
diff --git a/src/mongo/db/pipeline/document_source_set_window_fields.cpp b/src/mongo/db/pipeline/document_source_set_window_fields.cpp
index c06e89155c5..cbeeecf5134 100644
--- a/src/mongo/db/pipeline/document_source_set_window_fields.cpp
+++ b/src/mongo/db/pipeline/document_source_set_window_fields.cpp
@@ -240,7 +240,14 @@ list<intrusive_ptr<DocumentSource>> document_source_set_window_fields::create(
     }
     if (sortBy) {
         for (const auto& part : *sortBy) {
-            combined.push_back(part);
+            // Check and filter out the partition by field within the sort field. The partition
+            // field is already added to the combined sort field variable. As the partition only
+            // contains documents that have the same value in the partition key an additional sort
+            // over that field does not change the order or the documents within the partition.
+            // Hence a sort by $partion does not change the sort order.
+            if (!simplePartitionBy || *simplePartitionBy != part.fieldPath) {
+                combined.push_back(part);
+            }
         }
     }