From 813c73bdd73f9b37e1a9c35e35d435f1c6b942f6 Mon Sep 17 00:00:00 2001
From: Brendan Ashworth <brendan.ashworth@me.com>
Date: Tue, 29 Jun 2021 12:45:16 -0700
Subject: [PATCH] aggregate multiple columns

includes tests

Fixes: https://github.com/Gmousse/dataframe-js/issues/107
---
 src/group.js            | 24 ++++++++++++++++++------
 tests/dataframe.test.js | 28 ++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+), 6 deletions(-)

diff --git a/src/group.js b/src/group.js
index 5756903..b7cd004 100755
--- a/src/group.js
+++ b/src/group.js
@@ -136,19 +136,31 @@ class GroupedDataFrame {
 
     /**
      * Create an aggregation from a function.
-     * @param {Function} func The aggregation function.
+     * @param {Function | Object} func The aggregation function or an Object with keys representing
+     *     the output column, and values as aggregation functions.
      * @param {String} [columnName='aggregation'] The column name created by the aggregation.
      * @returns {DataFrame} A new DataFrame with a column 'aggregation' containing the result.
      * @example
      * groupedDF.aggregate(group => group.stat.sum('column1'));
+     * groupedDF.aggregate({ 'column1': group => group.stat.sum('column1') })
      */
     aggregate(func, columnName = "aggregation") {
+        // Convert into an object, where keys are the target column name, and
+        // values are the aggregation function.
+        const mappings = (typeof func == 'function') ? { [columnName]: func } : func;
+
         return this.df.__newInstance__(
-            [...this].map(({ group, groupKey }) => ({
-                ...groupKey,
-                [columnName]: func(group, groupKey)
-            })),
-            [...this.on, columnName]
+            [...this].map(({ group, groupKey }) => {
+                var transformed = { ...groupKey };
+
+                for (const column in mappings) {
+                    const aggregationFunc = mappings[column];
+                    transformed[column] = aggregationFunc(group, groupKey);
+                }
+
+                return transformed;
+            }),
+            [...this.on].concat(Object.keys(mappings))
         );
     }
 
diff --git a/tests/dataframe.test.js b/tests/dataframe.test.js
index 499cd79..13cd300 100755
--- a/tests/dataframe.test.js
+++ b/tests/dataframe.test.js
@@ -488,6 +488,18 @@ test("DataFrame rows can be ", (assert) => {
         "groupBy and compute (by aggregation) the count by group."
     );
 
+    assert.deepEqual(
+        df3
+            .groupBy("id")
+            .aggregate({ output: (group) => group.count() })
+            .toDict(),
+        {
+            id: [3, 6, 8, 1],
+            output: [2, 1, 2, 2]
+        },
+        "groupBy and compute (by aggregation) the count by group via object params."
+    );
+
     assert.deepEqual(
         df4
             .groupBy("id", "id2")
@@ -501,6 +513,22 @@ test("DataFrame rows can be ", (assert) => {
         "groupBy on multiple columns and compute the count by group."
     );
 
+    assert.deepEqual(
+        df4b
+            .groupBy('test1')
+            .aggregate({
+                test2: group => group.stat.mean('test2'),
+                test3: group => group.stat.max('test3')
+            })
+            .toDict(),
+        {
+            test1: [95, 94],
+            test2: [91, 96],
+            test3: [76, 99]
+        },
+        "groupBy on multiple columns via Object and compute independent aggregations."
+    );
+
     assert.deepEqual(
         new DataFrame(
             {