Repository: dask-contrib/dask-expr
Pull request #924 (2024-03-02 16:16:56)

Sort if split_out=1 in value_counts
ref #914


Patch:
diff --git a/dask_expr/_collection.py b/dask_expr/_collection.py
index c63f8e1..0e8a201 100644
--- a/dask_expr/_collection.py
+++ b/dask_expr/_collection.py
@@ -3842,6 +3842,9 @@ class Series(FrameBase):
                 split_out = 1
             else:
                 split_out = True
+        if split_out == 1 and split_out is not True and sort is None:
+            sort = True
+
         length = None
         if (split_out > 1 or split_out is True) and normalize:
             frame = self if not dropna else self.dropna()
diff --git a/dask_expr/_expr.py b/dask_expr/_expr.py
index 291178b..fe0c3b6 100644
--- a/dask_expr/_expr.py
+++ b/dask_expr/_expr.py
@@ -2653,8 +2653,10 @@ class Partitions(Expr):
         return (self.frame._name, self.partitions[index])
 
     def _simplify_down(self):
+        from dask_expr import SetIndexBlockwise
+
         if isinstance(self.frame, Blockwise) and not isinstance(
-            self.frame, (BlockwiseIO, Fused)
+            self.frame, (BlockwiseIO, Fused, SetIndexBlockwise)
         ):
             operands = [
                 Partitions(op, self.partitions)
diff --git a/dask_expr/_version.py b/dask_expr/_version.py
index f7b60f2..d8fe571 100644
--- a/dask_expr/_version.py
+++ b/dask_expr/_version.py
@@ -50,8 +50,8 @@ def get_config() -> VersioneerConfig:
     cfg = VersioneerConfig()
     cfg.VCS = "git"
     cfg.style = "pep440"
-    cfg.tag_prefix = ""
-    cfg.parentdir_prefix = "dask-"
+    cfg.tag_prefix = "v"
+    cfg.parentdir_prefix = "dask_expr-"
     cfg.versionfile_source = "dask_expr/_version.py"
     cfg.verbose = False
     return cfg
diff --git a/pyproject.toml b/pyproject.toml
index 4fd8bac..b6ac865 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -61,7 +61,7 @@ style = "pep440"
 versionfile_source = "dask_expr/_version.py"
 versionfile_build = "dask_expr/_version.py"
 tag_prefix = "v"
-parentdir_prefix = "dask-"
+parentdir_prefix = "dask_expr-"
 
 
 [tool.pytest.ini_options]


Test patch:
diff --git a/dask_expr/tests/test_collection.py b/dask_expr/tests/test_collection.py
index a1b9f1d..a0e9440 100644
--- a/dask_expr/tests/test_collection.py
+++ b/dask_expr/tests/test_collection.py
@@ -129,6 +129,11 @@ def test_setitem(pdf, df):
     assert_eq(df, pdf)
 
 
+def test_set_index_blockwise_culling(df, pdf):
+    result = df.set_index("x").partitions[0]
+    assert_eq(result, pdf.set_index("x").head(10))
+
+
 @xfail_gpu("https://github.com/rapidsai/cudf/issues/10271")
 def test_explode():
     with dask.config.set({"dataframe.convert-string": False}):
diff --git a/dask_expr/tests/test_reductions.py b/dask_expr/tests/test_reductions.py
index f56e0a5..77b987d 100644
--- a/dask_expr/tests/test_reductions.py
+++ b/dask_expr/tests/test_reductions.py
@@ -105,6 +105,15 @@ def test_value_counts(pdf, df, split_every, split_out):
     )
 
 
+def test_value_counts_sort(pdf):
+    pdf = pdf.iloc[4:]
+    df = from_pandas(pdf, npartitions=10)
+    pd.testing.assert_series_equal(
+        df.y.value_counts(split_out=1).compute(),
+        pdf.y.value_counts(),
+    )
+
+
 @pytest.mark.parametrize("split_every", [None, 5])
 @pytest.mark.parametrize("split_out", [1, True])
 def test_unique(pdf, df, split_every, split_out):