Implement $accumulator

JAVA-3640

Author: jstewart148

driver-core/src/main/com/mongodb/client/model/Accumulators.java

@@ -16,6 +16,15 @@
 
 package com.mongodb.client.model;
 
+import com.mongodb.lang.Nullable;
+import org.bson.BsonArray;
+import org.bson.BsonDocument;
+import org.bson.BsonString;
+
+import java.util.List;
+
+import static java.util.stream.Collectors.toList;
+
 /**
  * Builders for accumulators used in the group pipeline stage of an aggregation pipeline.
  *
@@ -38,7 +47,7 @@ public final class Accumulators {
      * @mongodb.driver.manual reference/operator/aggregation/sum/ $sum
      */
     public static <TExpression> BsonField sum(final String fieldName, final TExpression expression) {
-        return accumulator("$sum", fieldName, expression);
+        return accumulatorOperator("$sum", fieldName, expression);
     }
 
     /**
@@ -52,7 +61,7 @@ public static <TExpression> BsonField sum(final String fieldName, final TExpress
      * @mongodb.driver.manual reference/operator/aggregation/avg/ $avg
      */
     public static <TExpression> BsonField avg(final String fieldName, final TExpression expression) {
-        return accumulator("$avg", fieldName, expression);
+        return accumulatorOperator("$avg", fieldName, expression);
     }
 
     /**
@@ -66,7 +75,7 @@ public static <TExpression> BsonField avg(final String fieldName, final TExpress
      * @mongodb.driver.manual reference/operator/aggregation/first/ $first
      */
     public static <TExpression> BsonField first(final String fieldName, final TExpression expression) {
-        return accumulator("$first", fieldName, expression);
+        return accumulatorOperator("$first", fieldName, expression);
     }
 
     /**
@@ -80,7 +89,7 @@ public static <TExpression> BsonField first(final String fieldName, final TExpre
      * @mongodb.driver.manual reference/operator/aggregation/last/ $last
      */
     public static <TExpression> BsonField last(final String fieldName, final TExpression expression) {
-        return accumulator("$last", fieldName, expression);
+        return accumulatorOperator("$last", fieldName, expression);
     }
 
     /**
@@ -94,7 +103,7 @@ public static <TExpression> BsonField last(final String fieldName, final TExpres
      * @mongodb.driver.manual reference/operator/aggregation/max/ $max
      */
     public static <TExpression> BsonField max(final String fieldName, final TExpression expression) {
-        return accumulator("$max", fieldName, expression);
+        return accumulatorOperator("$max", fieldName, expression);
     }
 
     /**
@@ -108,7 +117,7 @@ public static <TExpression> BsonField max(final String fieldName, final TExpress
      * @mongodb.driver.manual reference/operator/aggregation/min/ $min
      */
     public static <TExpression> BsonField min(final String fieldName, final TExpression expression) {
-        return accumulator("$min", fieldName, expression);
+        return accumulatorOperator("$min", fieldName, expression);
     }
 
     /**
@@ -122,7 +131,7 @@ public static <TExpression> BsonField min(final String fieldName, final TExpress
      * @mongodb.driver.manual reference/operator/aggregation/push/ $push
      */
     public static <TExpression> BsonField push(final String fieldName, final TExpression expression) {
-        return accumulator("$push", fieldName, expression);
+        return accumulatorOperator("$push", fieldName, expression);
     }
 
     /**
@@ -136,7 +145,7 @@ public static <TExpression> BsonField push(final String fieldName, final TExpres
      * @mongodb.driver.manual reference/operator/aggregation/addToSet/ $addToSet
      */
     public static <TExpression> BsonField addToSet(final String fieldName, final TExpression expression) {
-        return accumulator("$addToSet", fieldName, expression);
+        return accumulatorOperator("$addToSet", fieldName, expression);
     }
 
     /**
@@ -155,7 +164,7 @@ public static <TExpression> BsonField addToSet(final String fieldName, final TEx
      * @since 3.2
      */
     public static <TExpression> BsonField stdDevPop(final String fieldName, final TExpression expression) {
-        return accumulator("$stdDevPop", fieldName, expression);
+        return accumulatorOperator("$stdDevPop", fieldName, expression);
     }
 
     /**
@@ -173,10 +182,125 @@ public static <TExpression> BsonField stdDevPop(final String fieldName, final TE
      * @since 3.2
      */
     public static <TExpression> BsonField stdDevSamp(final String fieldName, final TExpression expression) {
-        return accumulator("$stdDevSamp", fieldName, expression);
+        return accumulatorOperator("$stdDevSamp", fieldName, expression);
+    }
+
+    /**
+     * Creates an $accumulator pipeline stage
+     *
+     * @param fieldName            the field name
+     * @param initFunction         a function used to initialize the state
+     * @param accumulateFunction   a function used to accumulate documents
+     * @param mergeFunction        a function used to merge two internal states, e.g. accumulated on different shards or threads. It
+     *                             returns the resulting state of the accumulator.
+     * @return the $accumulator pipeline stage
+     * @mongodb.driver.manual reference/operator/aggregation/accumulator/ $accumulator
+     * @mongodb.server.release 4.4
+     * @since 4.1
+     */
+    public static BsonField accumulator(final String fieldName, final String initFunction, final String accumulateFunction,
+                                        final String mergeFunction) {
+        return accumulator(fieldName, initFunction, null, accumulateFunction, null, mergeFunction, null, "js");
+    }
+
+    /**
+     * Creates an $accumulator pipeline stage
+     *
+     * @param fieldName            the field name
+     * @param initFunction         a function used to initialize the state
+     * @param accumulateFunction   a function used to accumulate documents
+     * @param mergeFunction        a function used to merge two internal states, e.g. accumulated on different shards or threads. It
+     *                             returns the resulting state of the accumulator.
+     * @param finalizeFunction     a function used to finalize the state and return the result (may be null)
+     * @return the $accumulator pipeline stage
+     * @mongodb.driver.manual reference/operator/aggregation/accumulator/ $accumulator
+     * @mongodb.server.release 4.4
+     * @since 4.1
+     */
+    public static BsonField accumulator(final String fieldName, final String initFunction, final String accumulateFunction,
+                                        final String mergeFunction, @Nullable final String finalizeFunction) {
+        return accumulator(fieldName, initFunction, null, accumulateFunction, null, mergeFunction, finalizeFunction, "js");
+    }
+
+    /**
+     * Creates an $accumulator pipeline stage
+     *
+     * @param fieldName            the field name
+     * @param initFunction         a function used to initialize the state
+     * @param initArgs             init function’s arguments (may be null)
+     * @param accumulateFunction   a function used to accumulate documents
+     * @param accumulateArgs       additional accumulate function’s arguments (may be null). The first argument to the function
+     *                             is ‘state’.
+     * @param mergeFunction        a function used to merge two internal states, e.g. accumulated on different shards or threads. It
+     *                             returns the resulting state of the accumulator.
+     * @param finalizeFunction     a function used to finalize the state and return the result (may be null)
+     * @return the $accumulator pipeline stage
+     * @mongodb.driver.manual reference/operator/aggregation/accumulator/ $accumulator
+     * @mongodb.server.release 4.4
+     * @since 4.1
+     */
+    public static BsonField accumulator(final String fieldName, final String initFunction, @Nullable final List<String> initArgs,
+                                        final String accumulateFunction, @Nullable final List<String> accumulateArgs,
+                                        final String mergeFunction, @Nullable final String finalizeFunction) {
+        return accumulator(fieldName, initFunction, initArgs, accumulateFunction, accumulateArgs, mergeFunction, finalizeFunction, "js");
+    }
+
+    /**
+     * Creates an $accumulator pipeline stage
+     *
+     * @param fieldName            the field name
+     * @param initFunction         a function used to initialize the state
+     * @param accumulateFunction   a function used to accumulate documents
+     * @param mergeFunction        a function used to merge two internal states, e.g. accumulated on different shards or threads. It
+     *                             returns the resulting state of the accumulator.
+     * @param finalizeFunction     a function used to finalize the state and return the result (may be null)
+     * @param lang                 a language specifier
+     * @return the $accumulator pipeline stage
+     * @mongodb.driver.manual reference/operator/aggregation/accumulator/ $accumulator
+     * @mongodb.server.release 4.4
+     * @since 4.1
+     */
+    public static BsonField accumulator(final String fieldName, final String initFunction, final String accumulateFunction,
+                                        final String mergeFunction, @Nullable final String finalizeFunction, final String lang) {
+        return accumulator(fieldName, initFunction, null, accumulateFunction, null, mergeFunction, finalizeFunction, lang);
+    }
+
+    /**
+     * Creates an $accumulator pipeline stage
+     *
+     * @param fieldName            the field name
+     * @param initFunction         a function used to initialize the state
+     * @param initArgs             init function’s arguments (may be null)
+     * @param accumulateFunction   a function used to accumulate documents
+     * @param accumulateArgs       additional accumulate function’s arguments (may be null). The first argument to the function
+     *                             is ‘state’.
+     * @param mergeFunction        a function used to merge two internal states, e.g. accumulated on different shards or threads. It
+     *                             returns the resulting state of the accumulator.
+     * @param finalizeFunction     a function used to finalize the state and return the result (may be null)
+     * @param lang                 a language specifier
+     * @return the $accumulator pipeline stage
+     * @mongodb.driver.manual reference/operator/aggregation/accumulator/ $accumulator
+     * @mongodb.server.release 4.4
+     * @since 4.1
+     */
+    public static BsonField accumulator(final String fieldName, final String initFunction, @Nullable final List<String> initArgs,
+                                        final String accumulateFunction, @Nullable final List<String> accumulateArgs,
+                                        final String mergeFunction, @Nullable final String finalizeFunction, final String lang) {
+        BsonDocument accumulatorStage = new BsonDocument("init", new BsonString(initFunction))
+                .append("initArgs", initArgs != null ? new BsonArray(initArgs.stream().map(initArg ->
+                        new BsonString(initArg)).collect(toList())) : new BsonArray())
+                .append("accumulate", new BsonString(accumulateFunction))
+                .append("accumulateArgs", accumulateArgs != null ? new BsonArray(accumulateArgs.stream().map(accumulateArg ->
+                        new BsonString(accumulateArg)).collect(toList())) : new BsonArray())
+                .append("merge", new BsonString(mergeFunction))
+                .append("lang", new BsonString(lang));
+        if (finalizeFunction != null) {
+            accumulatorStage.append("finalize", new BsonString(finalizeFunction));
+        }
+        return accumulatorOperator("$accumulator", fieldName, accumulatorStage);
     }
 
-    private static <TExpression> BsonField accumulator(final String name, final String fieldName, final TExpression expression) {
+    private static <TExpression> BsonField accumulatorOperator(final String name, final String fieldName, final TExpression expression) {
         return new BsonField(fieldName, new SimpleExpression<TExpression>(name, expression));
     }
 

driver-core/src/test/functional/com/mongodb/client/model/AggregatesFunctionalSpecification.groovy

@@ -26,6 +26,7 @@ import org.bson.conversions.Bson
 import spock.lang.IgnoreIf
 
 import static com.mongodb.ClusterFixture.serverVersionAtLeast
+import static com.mongodb.client.model.Accumulators.accumulator
 import static com.mongodb.client.model.Accumulators.addToSet
 import static com.mongodb.client.model.Accumulators.avg
 import static com.mongodb.client.model.Accumulators.first
@@ -787,6 +788,50 @@ class AggregatesFunctionalSpecification extends OperationFunctionalSpecification
         helper?.drop()
     }
 
+    @IgnoreIf({ !serverVersionAtLeast(4, 3) })
+    def '$accumulator'() {
+        given:
+        def helper = getCollectionHelper()
+
+        when:
+        helper.drop()
+        helper.insertDocuments(Document.parse('{_id: 1, x: "string"}'))
+        def init = 'function() { return { x: "test string" } }'
+        def accumulate = 'function(state) { return state }'
+        def merge = 'function(state1, state2) { return state1 }'
+        def accumulatorExpr = accumulator('testString', init, accumulate, merge);
+        def results1 = helper.aggregate([group('$x', asList(accumulatorExpr))])
+
+        then:
+        results1.size() == 1
+        results1.contains(Document.parse('{ _id: "string", testString: { x: "test string" } }'))
+
+        when:
+        helper.drop()
+        helper.insertDocuments(Document.parse('{_id: 8751, title: "The Banquet", author: "Dante", copies: 2}'),
+                Document.parse('{_id: 8752, title: "Divine Comedy", author: "Dante", copies: 1}'),
+                Document.parse('{_id: 8645, title: "Eclogues", author: "Dante", copies: 2}'),
+                Document.parse('{_id: 7000, title: "The Odyssey", author: "Homer", copies: 10}'),
+                Document.parse('{_id: 7020, title: "Iliad", author: "Homer", copies: 10}'))
+        def initFunction = 'function(initCount, initSum) { return { count: parseInt(initCount), sum: parseInt(initSum) } }';
+        def accumulateFunction = 'function(state, numCopies) { return { count : state.count + 1, sum : state.sum + numCopies } }';
+        def mergeFunction = 'function(state1, state2) { return { count : state1.count + state2.count, sum : state1.sum + state2.sum } }';
+        def finalizeFunction = 'function(state) { return (state.sum / state.count) }';
+        def accumulatorExpression = accumulator('avgCopies', initFunction, [ '0', '0' ], accumulateFunction,
+                [ '$copies' ], mergeFunction, finalizeFunction)
+        def results2 = helper.aggregate([group('$author', asList(
+                new BsonField('minCopies', new Document('$min', '$copies')), accumulatorExpression,
+                new BsonField('maxCopies', new Document('$max', '$copies'))))])
+
+        then:
+        results2.size() == 2
+        results2.contains(Document.parse('{_id: "Dante", minCopies: 1, avgCopies: 1.6666666666666667, maxCopies : 2}'))
+        results2.contains(Document.parse('{_id: "Homer", minCopies: 10, avgCopies: 10.0, maxCopies : 10}'))
+
+        cleanup:
+        helper?.drop()
+    }
+
     @IgnoreIf({ !serverVersionAtLeast(3, 4) })
     def '$addFields'() {
         given: