Fix value_counts shuffling strategy (#11698)

dask · Jan 24, 2025 · 39b918f · 39b918f
1 parent d8eb38e
commit 39b918f
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/dask/dataframe/dask_expr/_reductions.py b/dask/dataframe/dask_expr/_reductions.py
@@ -1427,6 +1427,7 @@ class ValueCounts(ReductionConstantDim):
     reduction_chunk = M.value_counts
     reduction_aggregate = methods.value_counts_aggregate
     reduction_combine = methods.value_counts_combine
+    split_by = None
 
     @functools.cached_property
     def _meta(self):
@@ -1441,8 +1442,8 @@ def aggregate(cls, inputs, **kwargs):
             return func(_concat(inputs), observed=True, **kwargs)
 
     @property
-    def split_by(self):
-        return self.frame._meta.name
+    def shuffle_by_index(self):
+        return True
 
     @property
     def chunk_kwargs(self):

diff --git a/dask/dataframe/dask_expr/tests/test_reductions.py b/dask/dataframe/dask_expr/tests/test_reductions.py
@@ -558,3 +558,16 @@ def test_reductions_timestamps_display():
     data = pd.to_datetime(["2024-10-02 12:00:00", "2024-10-02 14:00:00"])
     df = from_pandas(pd.DataFrame({"valid_time": data}))
     assert df["valid_time"].min().__repr__()
+
+
+def test_value_counts_shuffle_properly():
+    pdf = pd.DataFrame(
+        {
+            "A": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
+            "B": [11, 12, 13, 4, 15, 6, 17, 8, 19, 10],
+        }
+    )
+    df = from_pandas(pdf, npartitions=2)
+    result = (df["A"] == df["B"]).value_counts()
+    expected = (pdf["A"] == pdf["B"]).value_counts()
+    assert_eq(result, expected)