From 1a1884904fd9855f146a120f6c7e44c740118113 Mon Sep 17 00:00:00 2001
From: Patrick Hoefler <61934744+phofl@users.noreply.github.com>
Date: Thu, 13 Feb 2025 12:49:17 +0100
Subject: [PATCH] Fix CI for changes in from_map

---
 distributed/shuffle/tests/test_shuffle.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/distributed/shuffle/tests/test_shuffle.py b/distributed/shuffle/tests/test_shuffle.py
index 7b888feb7bc..0a7d85a8243 100644
--- a/distributed/shuffle/tests/test_shuffle.py
+++ b/distributed/shuffle/tests/test_shuffle.py
@@ -2445,7 +2445,7 @@ def make_partition(i):
         return pd.DataFrame({"a": np.random.random(10), "b": np.random.random(10)})
 
     with dask.config.set({"dataframe.convert-string": False}):
-        ddf = dd.from_map(make_partition, range(5))
+        ddf = dd.from_map(make_partition, range(5), meta={"a": float, "b": float})
     with dask.config.set({"dataframe.shuffle.method": "p2p"}):
         out = ddf.shuffle(on="a", ignore_index=True)
     result, expected = c.compute([ddf, out])
@@ -2498,7 +2498,7 @@ def make_partition(i):
             )
         return pd.DataFrame({"a": np.random.random(10), "b": np.random.random(10)})
 
-    ddf = dd.from_map(make_partition, range(50))
+    ddf = dd.from_map(make_partition, range(50), meta={"a": float, "b": float})
     with dask.config.set({"dataframe.shuffle.method": "p2p"}):
         out = ddf.shuffle(on="a", ignore_index=True)
 
@@ -2522,7 +2522,7 @@ def make_partition(i):
             return pd.DataFrame({"a": np.random.random(10), "b": ["a"] * 10})
         return pd.DataFrame({"a": np.random.random(10), "b": np.random.random(10)})
 
-    ddf = dd.from_map(make_partition, range(50))
+    ddf = dd.from_map(make_partition, range(50), meta={"a": float, "b": float})
     with dask.config.set({"dataframe.shuffle.method": "p2p"}):
         out = ddf.shuffle(on="a", ignore_index=True)
 
@@ -2829,7 +2829,9 @@ def make_partition(partition_id, size):
         df["b"] = df["a"] % 23
         return df
 
-    df = dd.from_map(make_partition, np.arange(19), args=(250,))
+    df = dd.from_map(
+        make_partition, np.arange(19), args=(250,), meta={"a": int, "b": int}
+    )
 
     with dask.config.set(
         {"dataframe.shuffle.method": "p2p", "distributed.p2p.storage.disk": disk}