refactor: make sure value_counts always has a projection

ibis-project · Dec 6, 2022 · a70a302 · a70a302
1 parent 0234e5c
commit a70a302
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/ibis/expr/types/generic.py b/ibis/expr/types/generic.py
@@ -729,10 +729,13 @@ def value_counts(self, metric_name: str = "count") -> ir.Table:
         """
         from ibis.expr.analysis import find_first_base_table
 
-        base = find_first_base_table(self.op()).to_expr()
-        metric = base.count().name(metric_name)
-
-        return base.group_by(self).aggregate(metric)
+        return (
+            find_first_base_table(self.op())
+            .to_expr()
+            .select(self)
+            .group_by(self.get_name())
+            .agg(**{metric_name: lambda t: t.count()})
+        )
 
     def first(self) -> Column:
         return ops.FirstValue(self).to_expr()

diff --git a/ibis/tests/expr/test_table.py b/ibis/tests/expr/test_table.py
@@ -738,7 +738,7 @@ def test_group_by_column_select_api(table):
 def test_value_counts_convenience(table):
     # #152
     result = table.g.value_counts()
-    expected = table.group_by('g').aggregate(table.count().name('count'))
+    expected = table.select('g').group_by('g').aggregate(count=lambda t: t.count())
 
     assert_equal(result, expected)