SERVER-23612 Make the ENSURE_SORTED PlanStage respect the query's collation

3 files changed, 17 insertions, 4 deletions

diff --git a/src/mongo/db/exec/ensure_sorted.cpp b/src/mongo/db/exec/ensure_sorted.cpp
index e57cc61727c..1fef4b191ff 100644
--- a/src/mongo/db/exec/ensure_sorted.cpp
+++ b/src/mongo/db/exec/ensure_sorted.cpp
@@ -32,6 +32,7 @@
 
 #include "mongo/db/exec/scoped_timer.h"
 #include "mongo/db/exec/working_set_computed_data.h"
+#include "mongo/db/query/collation/collator_interface.h"
 #include "mongo/db/query/find_common.h"
 #include "mongo/stdx/memory.h"
 
@@ -44,9 +45,10 @@ const char* EnsureSortedStage::kStageType = "ENSURE_SORTED";
 
 EnsureSortedStage::EnsureSortedStage(OperationContext* opCtx,
                                      BSONObj pattern,
+                                     CollatorInterface* collator,
                                      WorkingSet* ws,
                                      PlanStage* child)
-    : PlanStage(kStageType, opCtx), _ws(ws) {
+    : PlanStage(kStageType, opCtx), _ws(ws), _collator(collator) {
     _children.emplace_back(child);
     _pattern = FindCommon::transformSortSpec(pattern);
 }
@@ -95,7 +97,7 @@ const SpecificStats* EnsureSortedStage::getSpecificStats() const {
 }
 
 bool EnsureSortedStage::isInOrder(const BSONObj& lhsSortKey, const BSONObj& rhsSortKey) const {
-    return lhsSortKey.woCompare(rhsSortKey, _pattern, /*considerFieldName*/ false) <= 0;
+    return lhsSortKey.woCompare(rhsSortKey, _pattern, /*considerFieldName*/ false, _collator) <= 0;
 }
 
 }  // namespace mongo
diff --git a/src/mongo/db/exec/ensure_sorted.h b/src/mongo/db/exec/ensure_sorted.h
index 3a58f4767de..a7f9bd45107 100644
--- a/src/mongo/db/exec/ensure_sorted.h
+++ b/src/mongo/db/exec/ensure_sorted.h
@@ -32,6 +32,8 @@
 
 namespace mongo {
 
+class CollatorInterface;
+
 /**
  * Takes the output of its child and drops any results that are not in the sort order specified by
  * 'pattern'. Thus, if the pattern is {a: -1} and the following documents are inputted:
@@ -41,7 +43,11 @@ namespace mongo {
  */
 class EnsureSortedStage final : public PlanStage {
 public:
-    EnsureSortedStage(OperationContext* opCtx, BSONObj pattern, WorkingSet* ws, PlanStage* child);
+    EnsureSortedStage(OperationContext* opCtx,
+                      BSONObj pattern,
+                      CollatorInterface* collator,
+                      WorkingSet* ws,
+                      PlanStage* child);
 
     bool isEOF() final;
     StageState doWork(WorkingSetID* out) final;
@@ -68,6 +74,10 @@ private:
     // The pattern that we're sorting by.
     BSONObj _pattern;
 
+    // Null if this ensure sorted stage orders strings according to simple binary compare. If
+    // non-null, represents the collator used to compare strings.
+    CollatorInterface* _collator;
+
     // The sort key of the previous result.
     BSONObj _prevSortKey;
 
diff --git a/src/mongo/db/query/stage_builder.cpp b/src/mongo/db/query/stage_builder.cpp
index a3ed4a01361..629b214475a 100644
--- a/src/mongo/db/query/stage_builder.cpp
+++ b/src/mongo/db/query/stage_builder.cpp
@@ -331,7 +331,8 @@ PlanStage* buildStages(OperationContext* txn,
         if (NULL == childStage) {
             return NULL;
         }
-        return new EnsureSortedStage(txn, esn->pattern, ws, childStage);
+        // TODO SERVER-23613: pass the appropriate CollatorInterface* instead of nullptr.
+        return new EnsureSortedStage(txn, esn->pattern, nullptr, ws, childStage);
     } else {
         mongoutils::str::stream ss;
         root->appendToString(&ss, 0);