Implement value_counts with duplicates and add test

Author: johnzangwill

pandas/core/groupby/generic.py

@@ -26,7 +26,10 @@
 
 import numpy as np
 
-from pandas._libs import reduction as libreduction
+from pandas._libs import (
+    lib,
+    reduction as libreduction,
+)
 from pandas._typing import (
     ArrayLike,
     Manager,
@@ -1730,7 +1733,7 @@ def value_counts(
                 observed=self.observed,
                 dropna=self.dropna,
             )
-            result = cast(Series, gb.size())
+            result = gb.size()
 
             if normalize:
                 # Normalize the results by dividing by the original group sizes.
@@ -1749,13 +1752,32 @@ def value_counts(
             if sort:
                 # Sort the values and then resort by the main grouping
                 index_level = range(len(self.grouper.groupings))
-                result = result.sort_values(ascending=ascending).sort_index(
-                    level=index_level, sort_remaining=False
+                result = (
+                    cast(Series, result)
+                    .sort_values(ascending=ascending)
+                    .sort_index(level=index_level, sort_remaining=False)
                 )
 
             if not self.as_index:
                 # Convert to frame
-                result = result.reset_index(name="proportion" if normalize else "count")
+                name = "proportion" if normalize else "count"
+                columns = result.index.names
+                if name in columns:
+                    raise ValueError(
+                        f"Column label '{name}' is duplicate of result column"
+                    )
+                columns = com.fill_missing_names(columns)
+                values = result.values
+                result_frame = DataFrame()
+                for i, column in enumerate(columns):
+                    level_values = result.index.get_level_values(i)._values
+                    if level_values.dtype == np.object_:
+                        level_values = lib.maybe_convert_objects(
+                            cast(np.ndarray, level_values)
+                        )
+                    result_frame.insert(i, column, level_values, allow_duplicates=True)
+                result = result_frame.assign(**{name: values})
+
             return result.__finalize__(self.obj, method="value_counts")
 
 

pandas/tests/groupby/test_frame_value_counts.py

@@ -413,22 +413,44 @@ def test_mixed_groupings(normalize, expected_label, expected_values):
     ],
 )
 @pytest.mark.parametrize("as_index", [False, True])
-def test_column_name_clashes(test, columns, expected_names, as_index):
+def test_column_label_duplicates(test, columns, expected_names, as_index):
+    # Test for duplicate input column labels and generated duplicate labels
     df = DataFrame([[1, 3, 5, 7, 9], [2, 4, 6, 8, 10]], columns=columns)
-
+    expected_data = [(1, 0, 7, 3, 5, 9), (2, 1, 8, 4, 6, 10)]
+    result = df.groupby(["a", [0, 1], "d"], as_index=as_index).value_counts()
     if as_index:
-        result = df.groupby(["a", [0, 1], "d"], as_index=as_index).value_counts()
         expected = Series(
             data=(1, 1),
             index=MultiIndex.from_tuples(
-                [(1, 0, 7, 3, 5, 9), (2, 1, 8, 4, 6, 10)],
+                expected_data,
                 names=expected_names,
             ),
         )
         tm.assert_series_equal(result, expected)
     else:
-        with pytest.raises(ValueError, match="cannot insert"):
-            df.groupby(["a", [0, 1], "d"], as_index=as_index).value_counts()
+        expected_data = [list(row) + [1] for row in expected_data]
+        expected_columns = list(expected_names)
+        expected_columns[1] = "level_1"
+        expected_columns.append("count")
+        expected = DataFrame(expected_data, columns=expected_columns)
+        tm.assert_frame_equal(result, expected)
+
+
+@pytest.mark.parametrize(
+    "normalize, expected_label",
+    [
+        (False, "count"),
+        (True, "proportion"),
+    ],
+)
+def test_result_label_duplicates(normalize, expected_label):
+    # Test for result column label duplicating an input column label
+    gb = DataFrame([[1, 2, 3]], columns=["a", "b", expected_label]).groupby(
+        "a", as_index=False
+    )
+    msg = f"Column label '{expected_label}' is duplicate of result column"
+    with pytest.raises(ValueError, match=msg):
+        gb.value_counts(normalize=normalize)
 
 
 def test_ambiguous_grouping():