SERVER-45450 $accumulator performance

3 files changed, 258 insertions, 9 deletions

diff --git a/jstests/aggregation/accumulators/accumulator_js.js b/jstests/aggregation/accumulators/accumulator_js.js
index 2942ef641dd..df9e3d85775 100644
--- a/jstests/aggregation/accumulators/accumulator_js.js
+++ b/jstests/aggregation/accumulators/accumulator_js.js
@@ -37,7 +37,7 @@ const command = {
     }],
 };
 
-const expectedResults = [
+let expectedResults = [
     {_id: "hello", wordCount: 2},
     {_id: "world", wordCount: 2},
     {_id: "hi", wordCount: 1},
@@ -99,7 +99,6 @@ 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 = [{
@@ -124,4 +123,220 @@ command.pipeline = [{
 }];
 res = assert.commandWorked(db.runCommand(command));
 assert(resultsEq(res.cursor.firstBatch, [{_id: 1, value: [null]}]), res.cursor);
+
+// Test that initArgs must evaluate to an array.
+command.pipeline = [{
+    $group: {
+        _id: 1,
+        value: {
+            $accumulator: {
+                init: function() {},
+                initArgs: {$const: 5},
+                accumulateArgs: [],
+                accumulate: function() {},
+                merge: function() {},
+                lang: 'js',
+            }
+        }
+    }
+}];
+assert.commandFailedWithCode(db.runCommand(command), 4544711);
+
+// Test that initArgs is passed to init.
+command.pipeline = [{
+    $group: {
+        _id: 1,
+        value: {
+            $accumulator: {
+                init: function(str1, str2) {
+                    return "initial_state_set_from_" + str1 + "_and_" + str2;
+                },
+                initArgs: ["ABC", "DEF"],
+                accumulateArgs: [],
+                accumulate: function(state) {
+                    return state;
+                },
+                merge: function(s1, s2) {
+                    return s1;
+                },
+                lang: 'js',
+            }
+        }
+    }
+}];
+res = assert.commandWorked(db.runCommand(command));
+assert(resultsEq(res.cursor.firstBatch, [{_id: 1, value: "initial_state_set_from_ABC_and_DEF"}]),
+       res.cursor);
+
+// Test that when initArgs errors, we fail gracefully, and don't call init.
+command.pipeline = [{
+    $group: {
+        _id: 1,
+        value: {
+            $accumulator: {
+                init: function() {
+                    throw 'init should not be called'
+                },
+                // Use $cond to thwart constant folding, to ensure we are testing evaluate rather
+                // than optimize.
+                initArgs: {$add: {$cond: ["$foo", "", ""]}},
+                accumulateArgs: [],
+                accumulate: function() {
+                    throw 'accumulate should not be called'
+                },
+                merge: function() {
+                    throw 'merge should not be called'
+                },
+                lang: 'js',
+            }
+        }
+    }
+}];
+// 16554 means "$add only supports numeric or date types"
+assert.commandFailedWithCode(db.runCommand(command), 16554);
+
+// Test that initArgs can have a different length per group.
+assert(db.accumulator_js.drop());
+assert.commandWorked(db.accumulator_js.insert([
+    {_id: 1, a: ['A', 'B', 'C']},
+    {_id: 2, a: ['A', 'B', 'C']},
+    {_id: 3, a: ['X', 'Y']},
+    {_id: 4, a: ['X', 'Y']},
+]));
+command.pipeline = [{
+    $group: {
+        _id: {a: "$a"},
+        value: {
+            $accumulator: {
+                init: function(...args) {
+                    return args.toString();
+                },
+                initArgs: "$a",
+                accumulateArgs: [],
+                accumulate: function(state) {
+                    return state;
+                },
+                merge: function(s1, s2) {
+                    return s1;
+                },
+                lang: 'js',
+            }
+        }
+    }
+}];
+res = assert.commandWorked(db.runCommand(command));
+assert(resultsEq(res.cursor.firstBatch,
+                 [{_id: ['A', 'B', 'C'], value: "A,B,C"}, {_id: ['X', 'Y'], value: "X,Y"}]),
+       res.cursor);
+
+// Test that accumulateArgs must evaluate to an array.
+command.pipeline = [{
+    $group: {
+        _id: 1,
+        value: {
+            $accumulator: {
+                init: function() {},
+                accumulateArgs: {$const: 5},
+                accumulate: function(state, value) {},
+                merge: function(s1, s2) {},
+                lang: 'js',
+            }
+        }
+    }
+}];
+assert.commandFailedWithCode(db.runCommand(command), 4544712);
+
+// Test that accumulateArgs can have more than one element.
+command.pipeline = [{
+    $group: {
+        _id: 1,
+        value: {
+            $accumulator: {
+                init: function() {},
+                accumulateArgs: ["ABC", "DEF"],
+                accumulate: function(state, str1, str2) {
+                    return str1 + str2
+                },
+                merge: function(s1, s2) {
+                    return s1 || s2;
+                },
+                lang: 'js',
+            }
+        }
+    }
+}];
+res = assert.commandWorked(db.runCommand(command));
+expectedResults = [
+    {_id: 1, value: "ABCDEF"},
+];
+assert(resultsEq(res.cursor.firstBatch, expectedResults), res.cursor);
+
+// Test that accumulateArgs can have a different length per document.
+command.pipeline = [{
+    $group: {
+        _id: 1,
+        value: {
+            $accumulator: {
+                init: function() {
+                    return [];
+                },
+                accumulateArgs: "$a",
+                accumulate: function(state, ...values) {
+                    state.push(values);
+                    state.sort();
+                    return state;
+                },
+                merge: function(s1, s2) {
+                    return s1.concat(s2);
+                },
+                lang: 'js',
+            }
+        }
+    }
+}];
+res = assert.commandWorked(db.runCommand(command));
+expectedResults = [
+    {_id: 1, value: [['A', 'B', 'C'], ['A', 'B', 'C'], ['X', 'Y'], ['X', 'Y']]},
+];
+assert(resultsEq(res.cursor.firstBatch, expectedResults), res.cursor);
+
+// Test that accumulateArgs can contain expressions that evaluate to null or missing.
+// The behavior is the same as ExpressionArray: arrays can't contain missing, so any expressions
+// that evaluate to missing get converted to null. Then, the nulls get serialized to BSON and passed
+// to JS as usual.
+assert(db.accumulator_js.drop());
+assert.commandWorked(db.accumulator_js.insert({}));
+command.pipeline = [{
+    $group: {
+        _id: 1,
+        value: {
+            $accumulator: {
+                init: function() {
+                    return null;
+                },
+                accumulateArgs: [
+                    null,
+                    "$no_such_field",
+                    {$let: {vars: {not_an_object: 5}, in : "$not_an_object.field"}}
+                ],
+                accumulate: function(state, ...values) {
+                    return {
+                        len: values.length,
+                        types: values.map(v => typeof v),
+                        values: values,
+                    };
+                },
+                merge: function(s1, s2) {
+                    return s1 || s2;
+                },
+                lang: 'js',
+            }
+        }
+    }
+}];
+res = assert.commandWorked(db.runCommand(command));
+expectedResults = [
+    {_id: 1, value: {len: 3, types: ['object', 'object', 'object'], values: [null, null, null]}},
+];
+assert(resultsEq(res.cursor.firstBatch, expectedResults), res.cursor);
 })();
diff --git a/jstests/aggregation/accumulators/accumulator_js_size_limits.js b/jstests/aggregation/accumulators/accumulator_js_size_limits.js
index c276898fe2c..cdf3680cf63 100644
--- a/jstests/aggregation/accumulators/accumulator_js_size_limits.js
+++ b/jstests/aggregation/accumulators/accumulator_js_size_limits.js
@@ -60,6 +60,30 @@ res = runExample(1, {
 });
 assert.commandFailedWithCode(res, [17260]);
 
+// Accumulator state and argument together exceed max BSON size.
+assert(coll.drop());
+const oneMBString = "a".repeat(1 * 1024 * 1024);
+const tenMBArray = Array.from({length: 10}, () => oneMBString);
+assert.commandWorked(coll.insert([{arr: tenMBArray}, {arr: tenMBArray}]));
+res = runExample(1, {
+    init: function() {
+        return [];
+    },
+    accumulate: function(state, input) {
+        state.push(input);
+        return state;
+    },
+    accumulateArgs: ["$arr"],
+    merge: function(state1, state2) {
+        return state1.concat(state2);
+    },
+    finalize: function() {
+        throw 'finalize should not be called';
+    },
+    lang: 'js',
+});
+assert.commandFailedWithCode(res, [4545000]);
+
 // $group size limit exceeded, and cannot spill.
 assert(coll.drop());
 assert.commandWorked(coll.insert(Array.from({length: 200}, (_, i) => ({_id: i}))));
diff --git a/jstests/core/system_js_access.js b/jstests/core/system_js_access.js
index a0d205cee8b..cf900bba288 100644
--- a/jstests/core/system_js_access.js
+++ b/jstests/core/system_js_access.js
@@ -55,7 +55,8 @@ assert.commandWorked(db.runCommand({
     cursor: {}
 }));
 
-// Mixed queries with both $where and $function should fail.
+// Mixed queries with both $where and $function should fail, because $where has to provide system.js
+// to user code, and $function has to not provide it.
 assert.commandFailedWithCode(db.runCommand({
     find: coll.getName(),
     filter: {
@@ -77,8 +78,8 @@ assert.commandFailedWithCode(db.runCommand({
 }),
                              4649200);
 
-// Mixed queries with both $function and $accumulator should succeed.
-// TODO SERVER-45450: Change $_internalJsReduce to $accumulator.
+// Queries with both $function and $accumulator should succeed, because both of these operators
+// provide system.js to user code.
 assert.commandWorked(db.runCommand({
     aggregate: coll.getName(),
     pipeline: [
@@ -99,9 +100,18 @@ assert.commandWorked(db.runCommand({
             $group: {
                 _id: "$name",
                 age: {
-                    $_internalJsReduce: {
-                        data: {k: "$name", v: "$age"},
-                        eval: "function concat(key, values) {return Array.sum(values);}"
+                    $accumulator: {
+                        init: function() {
+                            return 0;
+                        },
+                        accumulate: function(state, value) {
+                            return state + value;
+                        },
+                        accumulateArgs: ["$age"],
+                        merge: function(state1, state2) {
+                            return state1 + state2;
+                        },
+                        lang: 'js',
                     }
                 }
             }
@@ -111,4 +121,4 @@ assert.commandWorked(db.runCommand({
 }));
 
 assert.commandWorked(db.system.js.remove({_id: "isAdult"}));
-}());
\ No newline at end of file
+}());