SERVER-23349 add integration tests for aggregation stages respecting the collation

5 files changed, 249 insertions, 0 deletions

diff --git a/jstests/aggregation/sources/bucket/collation_bucket.js b/jstests/aggregation/sources/bucket/collation_bucket.js
new file mode 100644
index 00000000000..8fb791e0901
--- /dev/null
+++ b/jstests/aggregation/sources/bucket/collation_bucket.js
@@ -0,0 +1,52 @@
+// Test that the $bucket stage defines and sorts buckets according to the collation.
+(function() {
+    "use strict";
+
+    var results;
+    const numericOrdering = {collation: {locale: "en_US", numericOrdering: true}};
+
+    var coll = db.collation_bucket;
+    coll.drop();
+
+    function insertData() {
+        assert.writeOK(coll.insert({num: "1"}));
+        assert.writeOK(coll.insert({num: "2"}));
+        assert.writeOK(coll.insert({num: "5"}));
+        assert.writeOK(coll.insert({num: "10"}));
+        assert.writeOK(coll.insert({num: "20"}));
+        assert.writeOK(coll.insert({num: "50"}));
+        assert.writeOK(coll.insert({num: "100"}));
+        assert.writeOK(coll.insert({num: "200"}));
+        assert.writeOK(coll.insert({num: "500"}));
+    }
+
+    insertData();
+
+    // Test that $bucket respects an explicit collation.
+    results = coll.aggregate([{$bucket: {groupBy: "$num", boundaries: ["1", "10", "100", "1000"]}}],
+                             numericOrdering)
+                  .toArray();
+    assert.eq(3, results.length);
+    assert.eq({_id: "1", count: 3}, results[0]);
+    assert.eq({_id: "10", count: 3}, results[1]);
+    assert.eq({_id: "100", count: 3}, results[2]);
+
+    coll.drop();
+    assert.commandWorked(db.createCollection(coll.getName(), numericOrdering));
+    insertData();
+
+    // Test that $bucket respects the inherited collation.
+    results = coll.aggregate([{$bucket: {groupBy: "$num", boundaries: ["1", "10", "100", "1000"]}}])
+                  .toArray();
+    assert.eq(3, results.length);
+    assert.eq({_id: "1", count: 3}, results[0]);
+    assert.eq({_id: "10", count: 3}, results[1]);
+    assert.eq({_id: "100", count: 3}, results[2]);
+
+    // Test that the collection default can be overridden with the simple collation. In this case,
+    // the $bucket should fail, because under a lexicographical comparison strings like "2" or "5"
+    // won't fall into any of the buckets.
+    assert.throws(
+        () => coll.aggregate([{$bucket: {groupBy: "$num", boundaries: ["1", "10", "100", "1000"]}}],
+                             {collation: {locale: "simple"}}));
+})();
diff --git a/jstests/aggregation/sources/collation/collation_group.js b/jstests/aggregation/sources/group/collation_group.js
index 10081103bdf..4bf02ee5c97 100644
--- a/jstests/aggregation/sources/collation/collation_group.js
+++ b/jstests/aggregation/sources/group/collation_group.js
@@ -24,6 +24,25 @@
     // Ensure that equality of groups respects an explicit collation.
     assert.eq(2, coll.aggregate([{$group: {_id: "$str2"}}], diacriticInsensitive).itcount());
 
+    // Ensure that equality of groups created by $sortByCount respects the inherited collation.
+    assert.eq(2, coll.aggregate([{$sortByCount: "$str"}]).itcount());
+    assert.eq(4, coll.aggregate([{$sortByCount: "$str2"}]).itcount());
+
+    // Ensure that equality of groups created by $sortByCount respects an explicit collation.
+    assert.eq(4, coll.aggregate([{$sortByCount: "$str"}], diacriticInsensitive).itcount());
+    assert.eq(2, coll.aggregate([{$sortByCount: "$str2"}], diacriticInsensitive).itcount());
+
+    // Ensure that equality of groups inside $facet stage respects the inherited collation.
+    results =
+        coll.aggregate([{
+                $facet:
+                    {facetStr: [{$group: {_id: "$str"}}], facetStr2: [{$group: {_id: "$str2"}}]}
+            }])
+            .toArray();
+    assert.eq(1, results.length);
+    assert.eq(2, results[0].facetStr.length);
+    assert.eq(4, results[0].facetStr2.length);
+
     // Test that the $addToSet accumulator respects the inherited collation.
     results = coll.aggregate([{$group: {_id: null, set: {$addToSet: "$str"}}}]).toArray();
     assert.eq(1, results.length);
diff --git a/jstests/aggregation/sources/match/collation_match.js b/jstests/aggregation/sources/match/collation_match.js
new file mode 100644
index 00000000000..7136e03f9da
--- /dev/null
+++ b/jstests/aggregation/sources/match/collation_match.js
@@ -0,0 +1,46 @@
+// Test that the $match stage respects the collation.
+(function() {
+    "use strict";
+
+    var caseInsensitive = {collation: {locale: "en_US", strength: 2}};
+
+    var coll = db.collation_match;
+    coll.drop();
+    assert.writeOK(coll.insert({a: "a"}));
+
+    // Test that the $match respects an explicit collation when it can be pushed down into the query
+    // layer.
+    assert.eq(1, coll.aggregate([{$match: {a: "A"}}], caseInsensitive).itcount());
+
+    // Test that the $match respects an explicit collation when it cannot be pushed down into the
+    // query layer.
+    assert.eq(
+        1, coll.aggregate([{$project: {b: "B"}}, {$match: {b: "b"}}], caseInsensitive).itcount());
+
+    // Test that $match inside a $facet respects the collation.
+    assert.eq(1,
+              coll.aggregate([{$facet: {fct: [{$match: {a: "A"}}]}}], caseInsensitive)
+                  .toArray()[0]
+                  .fct.length);
+
+    // Test that when a $match can be split to be part before the $unwind and part after, both
+    // pieces of the split respect the collation.
+    coll.drop();
+    assert.writeOK(coll.insert({a: "foo", b: ["bar"]}));
+    assert.eq(1,
+              coll.aggregate([{$limit: 1}, {$unwind: "$b"}, {$match: {a: "FOO", b: "BAR"}}],
+                             caseInsensitive)
+                  .itcount());
+
+    coll.drop();
+    assert.commandWorked(db.createCollection(coll.getName(), caseInsensitive));
+    assert.writeOK(coll.insert({a: "a"}));
+
+    // Test that the $match respects the inherited collation when it can be pushed down into the
+    // query layer.
+    assert.eq(1, coll.aggregate([{$match: {a: "A"}}]).itcount());
+
+    // Test that the $match respects the inherited collation when it cannot be pushed down into the
+    // query layer.
+    assert.eq(1, coll.aggregate([{$project: {b: "B"}}, {$match: {b: "b"}}]).itcount());
+})();
diff --git a/jstests/aggregation/sources/redact/collation_redact.js b/jstests/aggregation/sources/redact/collation_redact.js
new file mode 100644
index 00000000000..d990c78959a
--- /dev/null
+++ b/jstests/aggregation/sources/redact/collation_redact.js
@@ -0,0 +1,37 @@
+// Test that the $redact stage respects the collation.
+(function() {
+    "use strict";
+
+    var caseInsensitive = {collation: {locale: "en_US", strength: 2}};
+
+    var coll = db.collation_redact;
+    coll.drop();
+    assert.writeOK(coll.insert({a: "a"}));
+
+    // Test that $redact respects an explicit collation. Since the top-level of the document gets
+    // pruned, we end up redacting the entire document and returning no results.
+    assert.eq(0,
+              coll.aggregate([{$redact: {$cond: [{$eq: ["A", "a"]}, "$$PRUNE", "$$KEEP"]}}],
+                             caseInsensitive)
+                  .itcount());
+
+    coll.drop();
+    assert.commandWorked(db.createCollection(coll.getName(), caseInsensitive));
+    assert.writeOK(coll.insert({a: "a"}));
+
+    // Test that $redact respects the inherited collation. Since the top-level of the document gets
+    // pruned, we end up redacting the entire document and returning no results.
+    assert.eq(
+        0,
+        coll.aggregate([{$redact: {$cond: [{$eq: ["A", "a"]}, "$$PRUNE", "$$KEEP"]}}]).itcount());
+
+    // Test that a $match which can be optimized to be pushed before the $redact respects the
+    // collation.
+    assert.eq(1, coll.aggregate([{$redact: "$$KEEP"}, {$match: {a: "A"}}]).itcount());
+
+    // Comparison to the internal constants bound to the $$KEEP, $$PRUNE, and $$DESCEND variable
+    // should not respect the collation.
+    assert.throws(() => coll.aggregate([{$redact: "KEEP"}], caseInsensitive));
+    assert.throws(() => coll.aggregate([{$redact: "PRUNE"}], caseInsensitive));
+    assert.throws(() => coll.aggregate([{$redact: "REDACT"}], caseInsensitive));
+})();
diff --git a/jstests/aggregation/sources/sort/collation_sort.js b/jstests/aggregation/sources/sort/collation_sort.js
new file mode 100644
index 00000000000..8febbafb857
--- /dev/null
+++ b/jstests/aggregation/sources/sort/collation_sort.js
@@ -0,0 +1,95 @@
+// Test that the $sort stage respects the collation.
+(function() {
+    "use strict";
+
+    // In French, words are sometimes ordered on the secondary level (a.k.a. at the level of
+    // diacritical marks) by the *last* accent difference rather than the first. This is specified
+    // by the {backwards: true} option.
+    //
+    // For example, côte < coté, since the last accent difference is "e" < "é". Without the reverse
+    // accent weighting turned on, these two words would sort in the opposite order, since "ô" >
+    // "o".
+    var frenchAccentOrdering = {collation: {locale: "fr", backwards: true}};
+
+    var coll = db.collation_sort;
+    coll.drop();
+    assert.writeOK(coll.insert({_id: 1, word1: "pêche", word2: "côté"}));
+    assert.writeOK(coll.insert({_id: 2, word1: "pêche", word2: "coté"}));
+    assert.writeOK(coll.insert({_id: 3, word1: "pêche", word2: "côte"}));
+    assert.writeOK(coll.insert({_id: 4, word1: "pèché", word2: "côté"}));
+    assert.writeOK(coll.insert({_id: 5, word1: "pèché", word2: "coté"}));
+    assert.writeOK(coll.insert({_id: 6, word1: "pèché", word2: "côte"}));
+    assert.writeOK(coll.insert({_id: 7, word1: "pêché", word2: "côté"}));
+    assert.writeOK(coll.insert({_id: 8, word1: "pêché", word2: "coté"}));
+    assert.writeOK(coll.insert({_id: 9, word1: "pêché", word2: "côte"}));
+
+    // Test that ascending sort respects the collation.
+    assert.eq([{_id: "pèché"}, {_id: "pêche"}, {_id: "pêché"}],
+              coll.aggregate([{$group: {_id: "$word1"}}, {$sort: {_id: 1}}]).toArray());
+    assert.eq([{_id: "pêche"}, {_id: "pèché"}, {_id: "pêché"}],
+              coll.aggregate([{$group: {_id: "$word1"}}, {$sort: {_id: 1}}], frenchAccentOrdering)
+                  .toArray());
+
+    // Test that descending sort respects the collation.
+    assert.eq([{_id: "pêché"}, {_id: "pêche"}, {_id: "pèché"}],
+              coll.aggregate([{$group: {_id: "$word1"}}, {$sort: {_id: -1}}]).toArray());
+    assert.eq([{_id: "pêché"}, {_id: "pèché"}, {_id: "pêche"}],
+              coll.aggregate([{$group: {_id: "$word1"}}, {$sort: {_id: -1}}], frenchAccentOrdering)
+                  .toArray());
+
+    // Test that compound, mixed ascending/descending sort respects the collation.
+    assert.eq([4, 6, 5, 1, 3, 2, 7, 9, 8],
+              coll.aggregate([
+                      {$sort: {word1: 1, word2: -1}},
+                      {$project: {_id: 1}},
+                      {$group: {_id: null, out: {$push: "$_id"}}}
+                  ])
+                  .toArray()[0]
+                  .out);
+    assert.eq([1, 2, 3, 4, 5, 6, 7, 8, 9],
+              coll.aggregate(
+                      [
+                        {$sort: {word1: 1, word2: -1}},
+                        {$project: {_id: 1}},
+                        {$group: {_id: null, out: {$push: "$_id"}}}
+                      ],
+                      frenchAccentOrdering)
+                  .toArray()[0]
+                  .out);
+
+    // Test that compound, mixed descending/ascending sort respects the collation.
+    assert.eq([8, 9, 7, 2, 3, 1, 5, 6, 4],
+              coll.aggregate([
+                      {$sort: {word1: -1, word2: 1}},
+                      {$project: {_id: 1}},
+                      {$group: {_id: null, out: {$push: "$_id"}}}
+                  ])
+                  .toArray()[0]
+                  .out);
+    assert.eq([9, 8, 7, 6, 5, 4, 3, 2, 1],
+              coll.aggregate(
+                      [
+                        {$sort: {word1: -1, word2: 1}},
+                        {$project: {_id: 1}},
+                        {$group: {_id: null, out: {$push: "$_id"}}}
+                      ],
+                      frenchAccentOrdering)
+                  .toArray()[0]
+                  .out);
+
+    // Test that sort inside a $facet respects the collation.
+    const results = coll.aggregate([{
+                                      $facet: {
+                                          fct: [
+                                              {$sort: {word1: -1, word2: 1}},
+                                              {$project: {_id: 1}},
+                                              {$group: {_id: null, out: {$push: "$_id"}}}
+                                          ]
+                                      }
+                                   }],
+                                   frenchAccentOrdering)
+                        .toArray();
+    assert.eq(1, results.length);
+    assert.eq(1, results[0].fct.length);
+    assert.eq([9, 8, 7, 6, 5, 4, 3, 2, 1], results[0].fct[0].out);
+})();