SERVER-45447 Add $accumulator for user-defined Javascript accumulators

3 files changed, 215 insertions, 1 deletions

diff --git a/jstests/aggregation/accumulators/accumulator_js.js b/jstests/aggregation/accumulators/accumulator_js.js
new file mode 100644
index 00000000000..2942ef641dd
--- /dev/null
+++ b/jstests/aggregation/accumulators/accumulator_js.js
@@ -0,0 +1,127 @@
+// Test the behavior of user-defined (Javascript) accumulators.
+(function() {
+"use strict";
+
+load('jstests/aggregation/extras/utils.js');
+
+db.accumulator_js.drop();
+
+for (const word of ["hello", "world", "world", "hello", "hi"]) {
+    db.accumulator_js.insert({word: word, val: 1});
+}
+
+const command = {
+    aggregate: 'accumulator_js',
+    cursor: {},
+    pipeline: [{
+        $group: {
+            _id: "$word",
+            wordCount: {
+                $accumulator: {
+                    init: function() {
+                        return 0;
+                    },
+                    accumulateArgs: ["$val"],
+                    accumulate: function(state, val) {
+                        return state + val;
+                    },
+                    merge: function(state1, state2) {
+                        return state1 + state2;
+                    },
+                    finalize: function(state) {
+                        return state;
+                    }
+                }
+            }
+        }
+    }],
+};
+
+const expectedResults = [
+    {_id: "hello", wordCount: 2},
+    {_id: "world", wordCount: 2},
+    {_id: "hi", wordCount: 1},
+];
+
+let res = assert.commandWorked(db.runCommand(command));
+assert(resultsEq(res.cursor.firstBatch, expectedResults), res.cursor);
+
+// Test that the functions can be passed as strings.
+{
+    const accumulatorSpec = command.pipeline[0].$group.wordCount.$accumulator;
+    accumulatorSpec.init = accumulatorSpec.init.toString();
+    accumulatorSpec.accumulate = accumulatorSpec.accumulate.toString();
+    accumulatorSpec.merge = accumulatorSpec.merge.toString();
+    accumulatorSpec.finalize = accumulatorSpec.finalize.toString();
+}
+res = assert.commandWorked(db.runCommand(command));
+assert(resultsEq(res.cursor.firstBatch, expectedResults), res.cursor);
+
+// Test that finalize is optional.
+delete command.pipeline[0].$group.wordCount.$accumulator.finalize;
+res = assert.commandWorked(db.runCommand(command));
+assert(resultsEq(res.cursor.firstBatch, expectedResults), res.cursor);
+
+// Test a finalizer other than the identity function. Finalizers are useful when the intermediate
+// state needs to be a different format from the final result.
+res = assert.commandWorked(db.runCommand(Object.merge(command, {
+    pipeline: [{
+        $group: {
+            _id: 1,
+            avgWordLen: {
+                $accumulator: {
+                    init: function() {
+                        return {count: 0, sum: 0};
+                    },
+                    accumulateArgs: [{$strLenCP: "$word"}],
+                    accumulate: function({count, sum}, wordLen) {
+                        return {count: count + 1, sum: sum + wordLen};
+                    },
+                    merge: function(s1, s2) {
+                        return {count: s1.count + s2.count, sum: s1.sum + s2.sum};
+                    },
+                    finalize: function({count, sum}) {
+                        return sum / count;
+                    },
+                    lang: 'js',
+                }
+            },
+        }
+    }],
+})));
+assert(resultsEq(res.cursor.firstBatch, [{_id: 1, avgWordLen: 22 / 5}]), res.cursor);
+
+// Test that a null word is considered a valid value.
+assert.commandWorked(db.accumulator_js.insert({word: null, val: 1}));
+expectedResults.push({_id: null, wordCount: 1});
+res = assert.commandWorked(db.runCommand(command));
+assert(resultsEq(res.cursor.firstBatch, expectedResults), res.cursor);
+
+// Test that missing fields become JS null.
+// This is similar to how most other agg operators work.
+// TODO SERVER-45450 is this a problem for mapreduce?
+assert(db.accumulator_js.drop());
+assert.commandWorked(db.accumulator_js.insert({sentinel: 1}));
+command.pipeline = [{
+    $group: {
+        _id: 1,
+        value: {
+            $accumulator: {
+                init: function() {
+                    return [];
+                },
+                accumulateArgs: ["$no_such_field"],
+                accumulate: function(state, value) {
+                    return state.concat([value]);
+                },
+                merge: function(s1, s2) {
+                    return s1.concat(s2);
+                },
+                lang: 'js',
+            }
+        }
+    }
+}];
+res = assert.commandWorked(db.runCommand(command));
+assert(resultsEq(res.cursor.firstBatch, [{_id: 1, value: [null]}]), res.cursor);
+})();
diff --git a/jstests/aggregation/accumulators/accumulator_js_size_limits.js b/jstests/aggregation/accumulators/accumulator_js_size_limits.js
new file mode 100644
index 00000000000..80bbcc49e07
--- /dev/null
+++ b/jstests/aggregation/accumulators/accumulator_js_size_limits.js
@@ -0,0 +1,87 @@
+// Test several different kinds of size limits on user-defined (Javascript) accumulators.
+(function() {
+"use strict";
+
+const coll = db.accumulator_js_size_limits;
+
+function runExample(groupKey, accumulatorSpec) {
+    return coll.runCommand({
+        aggregate: coll.getName(),
+        cursor: {},
+        pipeline: [{
+            $group: {
+                _id: groupKey,
+                accumulatedField: {$accumulator: accumulatorSpec},
+            }
+        }]
+    });
+}
+
+// Accumulator tries to create too long a String; it can't be serialized to BSON.
+coll.drop();
+assert.commandWorked(coll.insert({}));
+let res = runExample(1, {
+    init: function() {
+        return "a".repeat(20 * 1024 * 1024);
+    },
+    accumulate: function() {
+        throw 'accumulate should not be called';
+    },
+    accumulateArgs: [],
+    merge: function() {
+        throw 'merge should not be called';
+    },
+    finalize: function() {
+        throw 'finalize should not be called';
+    },
+    lang: 'js',
+});
+assert.commandFailedWithCode(res, [10334]);
+
+// Accumulator tries to return BSON larger than 16MB from JS.
+assert(coll.drop());
+assert.commandWorked(coll.insert({}));
+res = runExample(1, {
+    init: function() {
+        const str = "a".repeat(1 * 1024 * 1024);
+        return Array.from({length: 20}, () => str);
+    },
+    accumulate: function() {
+        throw 'accumulate should not be called';
+    },
+    accumulateArgs: [],
+    merge: function() {
+        throw 'merge should not be called';
+    },
+    finalize: function() {
+        throw 'finalize should not be called';
+    },
+    lang: 'js',
+});
+assert.commandFailedWithCode(res, [17260]);
+
+// $group size limit exceeded, and cannot spill.
+assert(coll.drop());
+assert.commandWorked(coll.insert(Array.from({length: 200}, (_, i) => ({_id: i}))));
+// By grouping on _id, each group contains only 1 document. This means it creates many
+// AccumulatorState instances.
+res = runExample("$_id", {
+    init: function() {
+        // Each accumulator state is big enough to be expensive, but not big enough to hit the BSON
+        // size limit.
+        return "a".repeat(1 * 1024 * 1024);
+    },
+    accumulate: function(state) {
+        return state;
+    },
+    accumulateArgs: [1],
+    merge: function(state1, state2) {
+        return state1;
+    },
+    finalize: function(state) {
+        return state.length;
+    },
+    lang: 'js',
+});
+assert.commandFailedWithCode(res, [16945]);
+})();
diff --git a/jstests/aggregation/accumulators/internal_js_reduce.js b/jstests/aggregation/accumulators/internal_js_reduce.js
index ae470b0950f..b315adacf70 100644
--- a/jstests/aggregation/accumulators/internal_js_reduce.js
+++ b/jstests/aggregation/accumulators/internal_js_reduce.js
@@ -43,7 +43,7 @@ const expectedResults = [
 ];
 
 let res = assert.commandWorked(db.runCommand(command));
-assert(resultsEq(res.cursor.firstBatch, expectedResults, res.cursor));
+assert(resultsEq(res.cursor.firstBatch, expectedResults), res.cursor);
 
 //
 // Test that the reduce function also accepts a string argument.