CXX-186 Add group command helper

TylerBrock · TylerBrock · commit 1e974de7cc32 · 2014-06-12T11:25:09.000-04:00
diff --git a/src/mongo/client/dbclient.cpp b/src/mongo/client/dbclient.cpp
@@ -791,6 +791,78 @@ namespace mongo {
         return info;
     }
 
+    void DBClientWithCommands::group(
+        const std::string& ns,
+        const std::string& jsreduce,
+        std::vector<BSONObj>* output,
+        const BSONObj& initial,
+        const BSONObj& cond,
+        const BSONObj& key,
+        const std::string& finalize
+    ) {
+        BSONObjBuilder groupObjBuilder;
+        _buildGroupObj(ns, jsreduce, initial, cond, finalize, &groupObjBuilder);
+
+        if (!key.isEmpty())
+            groupObjBuilder.append("key", key);
+
+        _runGroup(ns, groupObjBuilder.obj(), output);
+    }
+
+    void DBClientWithCommands::groupWithKeyFunction(
+        const std::string& ns,
+        const std::string& jsreduce,
+        std::vector<BSONObj>* output,
+        const BSONObj& initial,
+        const BSONObj& cond,
+        const std::string& jskey,
+        const std::string& finalize
+    ) {
+        BSONObjBuilder groupBuilder;
+        _buildGroupObj(ns, jsreduce, initial, cond, finalize, &groupBuilder);
+
+        if (!jskey.empty())
+            groupBuilder.append("$keyf", jskey);
+
+        _runGroup(ns, groupBuilder.obj(), output);
+    }
+
+    void DBClientWithCommands::_buildGroupObj(
+        const std::string& ns,
+        const std::string& jsreduce,
+        const BSONObj& initial,
+        const BSONObj& cond,
+        const std::string& finalize,
+        BSONObjBuilder* groupObj
+    ) {
+        groupObj->append("ns", nsGetCollection(ns));
+        groupObj->appendCode("$reduce", jsreduce);
+        groupObj->append("initial", initial);
+
+        if (!cond.isEmpty())
+            groupObj->append("cond", cond);
+        if (!finalize.empty())
+            groupObj->append("finalize", finalize);
+    }
+
+    void DBClientWithCommands::_runGroup(const std::string ns, const BSONObj& group, std::vector<BSONObj>* output) {
+        BSONObjBuilder commandBuilder;
+        commandBuilder.append("group", group);
+
+        BSONObj result;
+        bool ok = runCommand(nsGetDB(ns), commandBuilder.obj(), result);
+
+        if (!ok)
+            throw OperationException(result);
+
+        BSONElement resultArray = result.getField("retval");
+        BSONObjIterator resultIterator(resultArray.Obj());
+
+        while (resultIterator.more()) {
+            output->push_back(resultIterator.next().Obj().getOwned());
+        }
+    }
+
     bool DBClientWithCommands::eval(const string &dbname, const string &jscode, BSONObj& info, BSONElement& retValue, BSONObj *args) {
         BSONObjBuilder b;
         b.appendCode("$eval", jscode);
diff --git a/src/mongo/client/dbclientinterface.h b/src/mongo/client/dbclientinterface.h
@@ -821,6 +821,50 @@ namespace mongo {
         */
         BSONObj mapreduce(const std::string &ns, const std::string &jsmapf, const std::string &jsreducef, BSONObj query = BSONObj(), MROutput output = MRInline);
 
+        /**
+         * Groups documents in a collection by the specified keys and performs simple aggregation
+         * functions such as computing counts and sums.
+         *
+         * See: http://docs.mongodb.org/manual/reference/method/db.collection.group
+         *
+         * @param ns The namespace to group
+         * @param key The field or fields to group specified as a projection document: { field: 1 }
+         * @param jsreduce An aggregation function that operates on the documents during the group
+         * ing operation. The function should take two arguments: the current document and an
+         * aggregation result for that group.
+         * @param output The output vector.
+         * @param initial Initial aggregation result document.
+         * @param cond Optional selection criteria to determine which documents to process.
+         * @param finalize Optional function that runs for each item in the result set before
+         * returning the final values in the output vector.
+         */
+        void group(
+            const std::string& ns,
+            const std::string& jsreduce,
+            std::vector<BSONObj>* output,
+            const BSONObj& initial = BSONObj(),
+            const BSONObj& cond = BSONObj(),
+            const BSONObj& key = BSONObj(),
+            const std::string& finalize = ""
+        );
+
+        /**
+         * Does the same thing as group but accepts a key function that can be used to create the
+         * object representing the key. This allows for grouping on calculated fields rather than
+         * existing fields alone.
+         *
+         * @see DBClientWithCommands::group
+         */
+        void groupWithKeyFunction(
+            const std::string& ns,
+            const std::string& jsreduce,
+            std::vector<BSONObj>* output,
+            const BSONObj& initial = BSONObj(),
+            const BSONObj& cond = BSONObj(),
+            const std::string& jskey = "",
+            const std::string& finalize = ""
+        );
+
         /** Run javascript code on the database server.
            dbname    database SavedContext in which the code runs. The javascript variable 'db' will be assigned
                      to this database when the function is invoked.
@@ -1012,6 +1056,17 @@ namespace mongo {
     private:
         enum QueryOptions _cachedAvailableOptions;
         bool _haveCachedAvailableOptions;
+
+        void _buildGroupObj(
+            const std::string& ns,
+            const std::string& jsreduce,
+            const BSONObj& initial,
+            const BSONObj& cond,
+            const std::string& finalize,
+            BSONObjBuilder* groupObj
+        );
+
+        void _runGroup(const std::string ns, const BSONObj& group, std::vector<BSONObj>* output);
     };
 
     class DBClientWriter;
diff --git a/src/mongo/unittest/dbclient_test.cpp b/src/mongo/unittest/dbclient_test.cpp
@@ -281,6 +281,85 @@ namespace {
         ASSERT_EQUALS(c.count(TEST_NS), 2U);
     }
 
+    TEST_F(DBClientTest, SimpleGroup) {
+        c.insert(TEST_NS, BSON("a" << 1 << "color" << "green"));
+        c.insert(TEST_NS, BSON("a" << 2 << "color" << "green"));
+        c.insert(TEST_NS, BSON("a" << 3 << "color" << "green"));
+        c.insert(TEST_NS, BSON("a" << 1 << "color" << "blue"));
+
+        std::string reduce = "function(current, aggregate) {"
+            "if (current.color === 'green') { aggregate.green++; }"
+            "if (current.a) { aggregate.a += current.a; }"
+        "}";
+
+        BSONObj initial = BSON("a" << 0 << "green" << 0);
+        BSONObj cond = BSON("a" << LT << 3);
+        BSONObj key = BSON("color" << 1);
+
+        std::string finalize = "function(result) {"
+            "result.combined = result.green + result.a;"
+        "}";
+
+        vector<BSONObj> results;
+        c.group(TEST_NS, reduce, &results, initial, cond, key, finalize);
+
+        vector<BSONObj>::const_iterator it = results.begin();
+        while (it != results.end()) {
+            BSONObj current = *it;
+            if (current.getStringField("color") == std::string("green")) {
+                ASSERT_EQUALS(current.getField("a").Double(), 3.0);
+                ASSERT_EQUALS(current.getField("green").Double(), 2.0);
+                ASSERT_EQUALS(current.getField("combined").Double(), 5.0);
+            } else {
+                ASSERT_EQUALS(current.getField("a").Double(), 1.0);
+                ASSERT_EQUALS(current.getField("green").Double(), 0.0);
+                ASSERT_EQUALS(current.getField("combined").Double(), 1.0);
+            }
+            ++it;
+        }
+    }
+
+    TEST_F(DBClientTest, GroupWithKeyFunction) {
+        c.insert(TEST_NS, BSON("a" << 1 << "color" << "green"));
+        c.insert(TEST_NS, BSON("a" << 2 << "color" << "green"));
+        c.insert(TEST_NS, BSON("a" << 3 << "color" << "green"));
+        c.insert(TEST_NS, BSON("a" << 1 << "color" << "blue"));
+
+        std::string reduce = "function(current, aggregate) {"
+            "if (current.color === 'green') { aggregate.green++; }"
+            "if (current.a) { aggregate.a += current.a; }"
+        "}";
+
+        BSONObj initial = BSON("a" << 0 << "green" << 0);
+        BSONObj cond = BSON("a" << LT << 3);
+
+        std::string key = "function(doc) {"
+            "return { 'color': doc.color }"
+        "}";
+
+        std::string finalize = "function(result) {"
+            "result.combined = result.green + result.a;"
+        "}";
+
+        vector<BSONObj> results;
+        c.groupWithKeyFunction(TEST_NS, reduce, &results, initial, cond, key, finalize);
+
+        vector<BSONObj>::const_iterator it = results.begin();
+        while (it != results.end()) {
+            BSONObj current = *it;
+            if (current.getStringField("color") == std::string("green")) {
+                ASSERT_EQUALS(current.getField("a").Double(), 3.0);
+                ASSERT_EQUALS(current.getField("green").Double(), 2.0);
+                ASSERT_EQUALS(current.getField("combined").Double(), 5.0);
+            } else {
+                ASSERT_EQUALS(current.getField("a").Double(), 1.0);
+                ASSERT_EQUALS(current.getField("green").Double(), 0.0);
+                ASSERT_EQUALS(current.getField("combined").Double(), 1.0);
+            }
+            ++it;
+        }
+    }
+
     TEST_F(DBClientTest, InsertVectorContinueOnError) {
         vector<BSONObj> v;
         v.push_back(BSON("_id" << 1));
