SERVER-23654 Add $facet aggregation stage

3 files changed, 208 insertions, 0 deletions

diff --git a/jstests/aggregation/sources/facet/inner_graphlookup.js b/jstests/aggregation/sources/facet/inner_graphlookup.js
new file mode 100644
index 00000000000..42d10149240
--- /dev/null
+++ b/jstests/aggregation/sources/facet/inner_graphlookup.js
@@ -0,0 +1,39 @@
+/**
+ * Tests that using a $graphLookup stage inside of a $facet stage will yield the same results as
+ * using the $graphLookup stage outside of the $facet stage.
+ */
+(function() {
+    "use strict";
+
+    // We will only use one collection, the $graphLookup will look up from the same collection.
+    var graphColl = db.facetGraphLookup;
+
+    // The graph in ASCII form: 0 --- 1 --- 2    3
+    graphColl.drop();
+    assert.writeOK(graphColl.insert({_id: 0, edges: [1]}));
+    assert.writeOK(graphColl.insert({_id: 1, edges: [0, 2]}));
+    assert.writeOK(graphColl.insert({_id: 2, edges: [1]}));
+    assert.writeOK(graphColl.insert({_id: 3}));
+
+    // For each document in the collection, this will compute all the other documents that are
+    // reachable from this one.
+    const graphLookupStage = {
+        $graphLookup: {
+            from: graphColl.getName(),
+            startWith: "$_id",
+            connectFromField: "edges",
+            connectToField: "_id",
+            as: "connected"
+        }
+    };
+    const normalResults = graphColl.aggregate([graphLookupStage]).toArray();
+    const facetedResults = graphColl.aggregate([{$facet: {nested: [graphLookupStage]}}]).toArray();
+    assert.eq(facetedResults, [{nested: normalResults}]);
+
+    const normalResultsUnwound =
+        graphColl.aggregate([graphLookupStage, {$unwind: "$connected"}]).toArray();
+    const facetedResultsUnwound =
+        graphColl.aggregate([{$facet: {nested: [graphLookupStage, {$unwind: "$connected"}]}}])
+            .toArray();
+    assert.eq(facetedResultsUnwound, [{nested: normalResultsUnwound}]);
+}());
diff --git a/jstests/aggregation/sources/facet/inner_lookup.js b/jstests/aggregation/sources/facet/inner_lookup.js
new file mode 100644
index 00000000000..f4890324b65
--- /dev/null
+++ b/jstests/aggregation/sources/facet/inner_lookup.js
@@ -0,0 +1,32 @@
+/**
+ * Tests that using a $lookup stage inside of a $facet stage will yield the same results as using
+ * the $lookup stage outside of the $facet stage.
+ */
+(function() {
+    "use strict";
+
+    var local = db.facetLookupLocal;
+    var foreign = db.facetLookupForeign;
+
+    local.drop();
+    assert.writeOK(local.insert({_id: 0}));
+    assert.writeOK(local.insert({_id: 1}));
+
+    foreign.drop();
+    assert.writeOK(foreign.insert({_id: 0, foreignKey: 0}));
+    assert.writeOK(foreign.insert({_id: 1, foreignKey: 1}));
+    assert.writeOK(foreign.insert({_id: 2, foreignKey: 2}));
+
+    const lookupStage = {
+        $lookup:
+            {from: foreign.getName(), localField: "_id", foreignField: "foreignKey", as: "joined"}
+    };
+    const lookupResults = local.aggregate([lookupStage]).toArray();
+    const facetedLookupResults = local.aggregate([{$facet: {nested: [lookupStage]}}]).toArray();
+    assert.eq(facetedLookupResults, [{nested: lookupResults}]);
+
+    const lookupResultsUnwound = local.aggregate([lookupStage, {$unwind: "$joined"}]).toArray();
+    const facetedLookupResultsUnwound =
+        local.aggregate([{$facet: {nested: [lookupStage, {$unwind: "$joined"}]}}]).toArray();
+    assert.eq(facetedLookupResultsUnwound, [{nested: lookupResultsUnwound}]);
+}());
diff --git a/jstests/aggregation/sources/facet/use_cases.js b/jstests/aggregation/sources/facet/use_cases.js
new file mode 100644
index 00000000000..1bb3cefdd40
--- /dev/null
+++ b/jstests/aggregation/sources/facet/use_cases.js
@@ -0,0 +1,137 @@
+/**
+ * Tests some practical use cases of the $facet stage.
+ */
+(function() {
+    "use strict";
+    const dbName = "test";
+    const collName = jsTest.name();
+
+    Random.setRandomSeed();
+
+    /**
+     * Helper to get a random entry out of an array.
+     */
+    function randomChoice(array) {
+        return array[Random.randInt(array.length)];
+    }
+
+    /**
+     * Helper to generate a randomized document with the following schema:
+     * {
+     *   manufacturer: <string>,
+     *   price: <double>,
+     *   screenSize: <double>
+     * }
+     */
+    function generateRandomDocument() {
+        const manufacturers =
+            ["Sony", "Samsung", "LG", "Panasonic", "Mitsubishi", "Vizio", "Toshiba", "Sharp"];
+        const minPrice = 100;
+        const maxPrice = 4000;
+        const minScreenSize = 18;
+        const maxScreenSize = 40;
+
+        return {
+            manufacturer: randomChoice(manufacturers),
+            price: Random.randInt(maxPrice - minPrice + 1) + minPrice,
+            screenSize: Random.randInt(maxScreenSize - minScreenSize + 1) + minScreenSize,
+        };
+    }
+
+    function doExecutionTest(conn) {
+        var coll = conn.getDB(dbName).getCollection(collName);
+        coll.drop();
+
+        const nDocs = 1000 * 10;
+        var bulk = coll.initializeUnorderedBulkOp();
+        for (var i = 0; i < nDocs; i++) {
+            const doc = generateRandomDocument();
+            bulk.insert(doc);
+        }
+        assert.writeOK(bulk.execute());
+
+        //
+        // Compute the most common manufacturers, and the number of TVs in each price range.
+        //
+
+        // First compute each separately, to make sure we have the correct results.
+        const manufacturerPipe =
+            [{$group: {_id: "$manufacturer", count: {$sum: 1}}}, {$sort: {count: -1}}];
+        const mostCommonManufacturers = coll.aggregate(manufacturerPipe).toArray();
+
+        const pricePipe = [
+            {
+              $project: {
+                  priceBucket: {
+                      $switch: {
+                          branches: [
+                              {case: {$lt: ["$price", 500]}, then: "< 500"},
+                              {case: {$lt: ["$price", 1000]}, then: "500-1000"},
+                              {case: {$lt: ["$price", 1500]}, then: "1000-1500"},
+                              {case: {$lt: ["$price", 2000]}, then: "1500-2000"}
+                          ],
+                          default: "> 2000"
+                      }
+                  }
+              }
+            },
+            {$group: {_id: "$priceBucket", count: {$sum: 1}}},
+            {$sort: {count: -1}}
+        ];
+        const numTVsByPriceRange = coll.aggregate(pricePipe).toArray();
+
+        // Then compute the results using $facet.
+        const facetResult =
+            coll.aggregate([{$facet: {manufacturers: manufacturerPipe, prices: pricePipe}}])
+                .toArray();
+        assert.eq(facetResult.length, 1);
+        const facetManufacturers = facetResult[0].manufacturers;
+        const facetPrices = facetResult[0].prices;
+
+        // Then assert they are the same.
+        assert.eq(facetManufacturers, mostCommonManufacturers);
+        assert.eq(facetPrices, numTVsByPriceRange);
+    }
+
+    // Test against the standalone started by resmoke.py.
+    const conn = db.getMongo();
+    doExecutionTest(conn);
+
+    // Test against a sharded cluster.
+    const st = new ShardingTest({shards: 2});
+    doExecutionTest(st.s0);
+
+    // Test that $facet stage propagates information about involved collections, preventing users
+    // from doing things like $lookup from a sharded collection.
+    const shardedDBName = "sharded";
+    const shardedCollName = "collection";
+    const shardedColl = st.getDB(shardedDBName).getCollection(shardedCollName);
+    const unshardedColl = st.getDB(shardedDBName).getCollection(collName);
+
+    assert.commandWorked(st.admin.runCommand({enableSharding: shardedDBName}));
+    assert.commandWorked(
+        st.admin.runCommand({shardCollection: shardedColl.getFullName(), key: {_id: 1}}));
+    assert.commandFailed(unshardedColl.runCommand({
+        aggregate: unshardedColl,
+        pipline: [{
+            $lookup:
+                {from: shardedCollName, localField: "_id", foreignField: "_id", as: "results"}
+        }]
+    }));
+    assert.commandFailed(unshardedColl.runCommand({
+        aggregate: unshardedColl,
+        pipeline: [{
+            $facet: {
+                a: [{
+                    $lookup: {
+                        from: shardedCollName,
+                        localField: "_id",
+                        foreignField: "_id",
+                        as: "results"
+                    }
+                }]
+            }
+        }]
+    }));
+    st.stop();
+}());