Repository: dask-contrib/dask-expr
Pull request #951 (2024-03-07 17:57:38)

Add dtype and columns type check for shuffle


Patch:
diff --git a/dask_expr/_collection.py b/dask_expr/_collection.py
index 5a2c440..2211551 100644
--- a/dask_expr/_collection.py
+++ b/dask_expr/_collection.py
@@ -49,6 +49,7 @@ from dask.utils import (
     IndexCallable,
     M,
     derived_from,
+    get_default_shuffle_method,
     get_meta_library,
     key_split,
     maybe_pluralize,
@@ -782,6 +783,20 @@ Expr={expr}"""
         elif pd.api.types.is_list_like(on) and not is_dask_collection(on):
             on = list(on)
 
+        method = shuffle_method or get_default_shuffle_method()
+        if method == "p2p":
+            from distributed.shuffle._arrow import check_dtype_support
+
+            check_dtype_support(self._meta)
+
+            if any(not isinstance(c, str) for c in self._meta.columns):
+                unsupported = {
+                    c: type(c) for c in self._meta.columns if not isinstance(c, str)
+                }
+                raise TypeError(
+                    f"p2p requires all column names to be str, found: {unsupported}",
+                )
+
         # Returned shuffled result
         return new_collection(
             RearrangeByColumn(
@@ -4476,14 +4491,15 @@ def from_array(arr, chunksize=50_000, columns=None, meta=None):
 
     from dask_expr.io.io import FromArray
 
-    return new_collection(
-        FromArray(
-            arr,
-            chunksize=chunksize,
-            original_columns=columns,
-            meta=meta,
-        )
+    result = FromArray(
+        arr,
+        chunksize=chunksize,
+        original_columns=columns,
+        meta=meta,
     )
+    if pyarrow_strings_enabled() and arr.dtype.kind in "OU":
+        result = expr.ArrowStringConversion(result)
+    return new_collection(result)
 
 
 def from_graph(*args, **kwargs):


Test patch:
diff --git a/dask_expr/io/tests/test_io.py b/dask_expr/io/tests/test_io.py
index 459bb4b..b967559 100644
--- a/dask_expr/io/tests/test_io.py
+++ b/dask_expr/io/tests/test_io.py
@@ -394,6 +394,14 @@ def test_from_array():
             assert expr.columns == ["a"]
 
 
+@pytest.mark.parametrize("dtype", [object, str])
+def test_from_array_string_conersion(dtype):
+    arr = np.array(["a", "b", "c", "d"], dtype=dtype)
+    result = from_array(arr, chunksize=2)
+    assert result.dtype == "string"
+    assert result.compute().dtype == "string"
+
+
 def test_from_dask_array():
     import dask.array as da
 
diff --git a/dask_expr/tests/test_distributed.py b/dask_expr/tests/test_distributed.py
index cb0d5fd..b8d0081 100644
--- a/dask_expr/tests/test_distributed.py
+++ b/dask_expr/tests/test_distributed.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import numpy as np
 import pytest
 
 from dask_expr import Repartition, from_pandas, map_partitions, merge
@@ -356,6 +357,19 @@ async def test_future_in_map_partitions(c, s, a, b):
     pd.testing.assert_frame_equal(result, expected)
 
 
+@gen_cluster(client=True)
+async def test_shuffle_consistency_checks(c, s, a, b):
+    pdf = pd.DataFrame({"x": [1, 2, 3]})
+    df = from_pandas(pdf, npartitions=2)
+    df2 = df.astype(np.csingle)
+    with pytest.raises(TypeError, match="p2p does not support data of type"):
+        df2.shuffle("x")
+
+    df.columns = [1]
+    with pytest.raises(TypeError, match="p2p requires all column names to be str"):
+        df.shuffle(1)
+
+
 @gen_cluster(client=True)
 async def test_merge_indicator(c, s, a, b):
     data = {