Repository: dask/dask
Pull request #10493 (2023-09-06 08:58:01)

Fix `_partitions` `dtype` in `meta` for `DataFrame.set_index` and `DataFrame.sort_values`
- [x] Fixes issue with `DataFrame.set_index()` that occurred in dask/distributed#8157
- [x] Closes dask/distributed#8165
- [x] Tests added / passed
- [x] Passes `pre-commit run --all-files`


fjetter: I can confirm this fixes the reproducer in https://github.com/dask/distributed/issues/8165

@hendrikmakait can we put a test in using the reproducer of #8165? I think that would be more meaningful than the low level asserts on the metadata. This will likely need to live in `test_distributed.py` and use P2P
hendrikmakait: > [](https://github.com/dask/distributed/issues/8165)
> 
> @hendrikmakait can we put a test in using the reproducer of #8165? I think that would be more meaningful than the low level asserts on the metadata. This will likely need to live in `test_distributed.py` and use P2P

Agreed, we can put an additional regression test in distributed. Note that dask/distributed#8157 includes an early check for the dtype that raises instead of silently failing. 
phofl: This might make a new release necessary. I think set_index and sort_values is broken generally on the current released version
hendrikmakait: > This might make a new release necessary. I think set_index and sort_values is broken generally on the current released version

+1, from what I see from the regression tests, these don't work with P2P.
hendrikmakait: > There is also a `test_distributed.py` file in dask/dask that you could use for the test

I'd rather have the regression tests live together with the actual P2P implementation.
phofl: cc @jrbourbeau for the new release

Patch:
diff --git a/dask/dataframe/shuffle.py b/dask/dataframe/shuffle.py
index fd69a9f81..469835849 100644
--- a/dask/dataframe/shuffle.py
+++ b/dask/dataframe/shuffle.py
@@ -288,7 +288,6 @@ def set_partition(
     shuffle
     partd
     """
-    meta = df._meta._constructor_sliced([0])
     if isinstance(divisions, tuple):
         # pd.isna considers tuples to be scalars. Convert to a list.
         divisions = list(divisions)
@@ -310,6 +309,11 @@ def set_partition(
     else:
         divisions = df._meta._constructor_sliced(divisions, dtype=dtype)
 
+    meta = df._meta._constructor_sliced([0])
+    # Ensure that we have the same index as before to avoid alignment
+    # when calculating meta dtypes later on
+    meta.index = df._meta_nonempty.index[:1]
+
     if not isinstance(index, Series):
         partitions = df[index].map_partitions(
             set_partitions_pre, divisions=divisions, meta=meta
@@ -452,6 +456,9 @@ def rearrange_by_divisions(
     if not duplicates:
         divisions = divisions.drop_duplicates()
     meta = df._meta._constructor_sliced([0])
+    # Ensure that we have the same index as before to avoid alignment
+    # when calculating meta dtypes later on
+    meta.index = df._meta_nonempty.index[:1]
     # Assign target output partitions to every row
     partitions = df[column].map_partitions(
         set_partitions_pre,


Test patch:
diff --git a/dask/dataframe/tests/test_shuffle.py b/dask/dataframe/tests/test_shuffle.py
index be28b7619..80b2bc431 100644
--- a/dask/dataframe/tests/test_shuffle.py
+++ b/dask/dataframe/tests/test_shuffle.py
@@ -1428,6 +1428,43 @@ def test_set_index_with_series_uses_fastpath():
     assert_eq(res, expected)
 
 
+def test_set_index_partitions_meta_dtype():
+    ddf = dd.from_pandas(
+        pd.DataFrame({"a": np.random.randint(0, 10, 100)}, index=np.random.random(100)),
+        npartitions=10,
+    )
+    # Disk-based shuffle doesn't use HLG layers at the moment, so we only test tasks
+    ddf = ddf.set_index("a", shuffle="tasks")
+    # Cull the HLG
+    dsk = ddf.__dask_graph__()
+
+    for layer in dsk.layers.values():
+        if isinstance(layer, dd.shuffle.SimpleShuffleLayer):
+            assert layer.meta_input["_partitions"].dtype == np.int64
+
+
+def test_sort_values_partitions_meta_dtype_with_divisions():
+    with dask.config.set({"dataframe.shuffle.method": "tasks"}):
+        ddf = dd.from_pandas(
+            pd.DataFrame(
+                {
+                    "a": np.random.randint(0, 10, 100),
+                    "b": np.random.randint(0, 10, 100),
+                },
+                index=np.random.random(100),
+            ),
+            npartitions=10,
+        )
+        # Disk-based shuffle doesn't use HLG layers at the moment, so we only test tasks
+        ddf = ddf.set_index("a", shuffle="tasks").sort_values("b")
+        # Cull the HLG
+        dsk = ddf.__dask_graph__()
+
+        for layer in dsk.layers.values():
+            if isinstance(layer, dd.shuffle.SimpleShuffleLayer):
+                assert layer.meta_input["_partitions"].dtype == np.int64
+
+
 @pytest.mark.parametrize("ascending", [True, False])
 @pytest.mark.parametrize("by", ["a", "b", ["a", "b"]])
 @pytest.mark.parametrize("nelem", [10, 500])