From bce5175f2f031e909404c2892416b50281829c2a Mon Sep 17 00:00:00 2001
From: Iago <6860957+igonro@users.noreply.github.com>
Date: Wed, 23 Feb 2022 09:34:29 +0100
Subject: [PATCH] Fix bug in groupby checking wrong colDtype

---
 src/danfojs-base/aggregators/groupby.ts          |  3 ++-
 src/danfojs-base/core/frame.ts                   |  4 +---
 .../tests/aggregators/groupby.test.js            | 15 +++++++++++++++
 .../test/aggregators/groupby.test.ts             | 16 ++++++++++++++++
 4 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/src/danfojs-base/aggregators/groupby.ts b/src/danfojs-base/aggregators/groupby.ts
index 0b965fe7..63671914 100644
--- a/src/danfojs-base/aggregators/groupby.ts
+++ b/src/danfojs-base/aggregators/groupby.ts
@@ -262,7 +262,8 @@ export default class Groupby {
         let colName = groupColNames[colKey]
         let colIndex = this.columnName.indexOf(colName)
         let colDtype = this.colDtype[colIndex]
-        if (colDtype === "string") throw new Error(`Can't perform math operation on column ${colName}`)
+        let operationVal = (typeof operation === "string") ? operation : operation[colName]
+        if (colDtype === "string" && operationVal !== "count") throw new Error(`Can't perform math operation on column ${colName}`)
 
         if (typeof operation === "string") {
           let colName2 = `${colName}_${operation}`
diff --git a/src/danfojs-base/core/frame.ts b/src/danfojs-base/core/frame.ts
index 7b82c736..f68e00dc 100644
--- a/src/danfojs-base/core/frame.ts
+++ b/src/danfojs-base/core/frame.ts
@@ -3162,9 +3162,7 @@ export default class DataFrame extends NDframe implements DataFrameInterface {
     groupby(col: Array<string>): Groupby {
         const columns = this.columns
         const colIndex = col.map((val) => columns.indexOf(val))
-        const colDtype = this.dtypes.filter((val, index) => {
-            return colIndex.includes(index)
-        })
+        const colDtype = this.dtypes
 
         return new Groupby(
             col,
diff --git a/src/danfojs-browser/tests/aggregators/groupby.test.js b/src/danfojs-browser/tests/aggregators/groupby.test.js
index e56a2aa8..d9932697 100644
--- a/src/danfojs-browser/tests/aggregators/groupby.test.js
+++ b/src/danfojs-browser/tests/aggregators/groupby.test.js
@@ -353,4 +353,19 @@ describe("groupby", function () {
     ];
     assert.deepEqual(group_df.size().values, rslt);
   });
+  it("issue 396: fix groupby when first column is int", function () {
+    let data = {
+      'hours': [5, 6, 2, 8, 4, 3],
+      'worker': ["david", "david", "john", "alice", "john", "david"],
+      'day': ["monday", "tuesday", "wednesday", "thursday", "friday", "friday"]
+    };
+    let df = new dfd.DataFrame(data);
+    let group_df = df.groupby(["worker"]);
+    let rslt = [
+      ["david", 3, 3],
+      ["john", 2, 2],
+      ["alice", 1, 1]
+    ];
+    assert.deepEqual(group_df.count().values, rslt);
+  });
 });
diff --git a/src/danfojs-node/test/aggregators/groupby.test.ts b/src/danfojs-node/test/aggregators/groupby.test.ts
index ad7953f6..d267c5cb 100644
--- a/src/danfojs-node/test/aggregators/groupby.test.ts
+++ b/src/danfojs-node/test/aggregators/groupby.test.ts
@@ -364,4 +364,20 @@ describe("groupby", function () {
         ]
         assert.deepEqual(group_df.size().values, rslt);
     });
+
+    it("issue 396: fix groupby when first column is int", function () {
+        let data = {
+            'hours': [5, 6, 2, 8, 4, 3],
+            'worker': ["david", "david", "john", "alice", "john", "david"],
+            'day': ["monday", "tuesday", "wednesday", "thursday", "friday", "friday"]
+        };
+        let df = new DataFrame(data);
+        let group_df = df.groupby(["worker"]);
+        let rslt = [
+            ["david", 3, 3],
+            ["john", 2, 2],
+            ["alice", 1, 1]
+        ];
+        assert.deepEqual(group_df.count().values, rslt);
+    });
 })
\ No newline at end of file