Fixed countBy(string)

apache · Jan 15, 2018 · 20717d5 · 20717d5
1 parent 7244887
commit 20717d5
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 4 deletions.
diff --git 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
@@ -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
@@ -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({