From 20717d59379b6c5f590e0a6771a3e0818ddc3de8 Mon Sep 17 00:00:00 2001
From: Brian Hulette <brian.hulette@ccri.com>
Date: Mon, 15 Jan 2018 17:53:22 -0500
Subject: [PATCH] Fixed countBy(string)

---
 js/perf/index.js            |  2 +-
 js/src/table.ts             |  4 ++--
 js/test/unit/table-tests.ts | 16 +++++++++++++++-
 3 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/js/perf/index.js b/js/perf/index.js
index d31b6430ec871..29d5edf56de8e 100644
--- a/js/perf/index.js
+++ b/js/perf/index.js
@@ -180,7 +180,7 @@ function createDataFrameCountByTest(table, column) {
         async: true,
         name: `name: '${column}', length: ${table.length}, type: ${table.columns[colidx].type}`,
         fn() {
-            table.countBy(col(column));
+            table.countBy(column);
         }
     };
 }
diff --git a/js/src/table.ts b/js/src/table.ts
index f00b5ef9da1df..554844be2c8c4 100644
--- a/js/src/table.ts
+++ b/js/src/table.ts
@@ -115,7 +115,7 @@ export class Table implements DataFrame {
         return this.lengths.reduce((acc, val) => acc + val);
     }
     countBy(count_by: (Col|string)): CountByResult {
-        if (count_by instanceof String) {
+        if (!(count_by instanceof Col)) {
             count_by = new Col(count_by);
         }
 
@@ -216,7 +216,7 @@ class FilteredDataFrame implements DataFrame {
     }
 
     countBy(count_by: (Col|string)): CountByResult {
-        if (count_by instanceof String) {
+        if (!(count_by instanceof Col)) {
             count_by = new Col(count_by);
         }
 
diff --git a/js/test/unit/table-tests.ts b/js/test/unit/table-tests.ts
index 33fb2d178b0d2..9b19c9e56a243 100644
--- a/js/test/unit/table-tests.ts
+++ b/js/test/unit/table-tests.ts
@@ -159,14 +159,21 @@ describe(`Table`, () => {
             expect(table.filter(col('dictionary').eq('a')).count()).toEqual(3);
         });
         test(`countBy on dictionary returns the correct counts`, () => {
+            // Make sure countBy works both with and without the Col wrapper
+            // class
             expect(table.countBy(col('dictionary')).asJSON()).toEqual({
                 'a': 3,
                 'b': 2,
                 'c': 2,
             });
+            expect(table.countBy('dictionary').asJSON()).toEqual({
+                'a': 3,
+                'b': 2,
+                'c': 2,
+            });
         });
         test(`countBy on dictionary with filter returns the correct counts`, () => {
-            expect(table.filter(col('i32').eq(1)).countBy(col('dictionary')).asJSON()).toEqual({
+            expect(table.filter(col('i32').eq(1)).countBy('dictionary').asJSON()).toEqual({
                 'a': 1,
                 'b': 1,
                 'c': 1,
@@ -354,11 +361,18 @@ describe(`Table`, () => {
             expect(table.filter(col('dictionary').eq('a')).count()).toEqual(3);
         });
         test(`countBy on dictionary returns the correct counts`, () => {
+            // Make sure countBy works both with and without the Col wrapper
+            // class
             expect(table.countBy(col('dictionary')).asJSON()).toEqual({
                 'a': 3,
                 'b': 3,
                 'c': 3,
             });
+            expect(table.countBy('dictionary').asJSON()).toEqual({
+                'a': 3,
+                'b': 3,
+                'c': 3,
+            });
         });
         test(`countBy on dictionary with filter returns the correct counts`, () => {
             expect(table.filter(col('i32').eq(1)).countBy(col('dictionary')).asJSON()).toEqual({