Repository: dask/dask
Pull request #11931 (2025-05-05 16:19:19)

Fix ``to_dask_array`` for single partition
Closes https://github.com/dask/dask/issues/11928 

github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/dask/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

      9 files  ±0        9 suites  ±0   3h 25m 1s ⏱️ + 4m 0s
 18 021 tests +1   16 807 ✅ +1   1 214 💤 ±0  0 ❌ ±0 
161 263 runs  +9  149 149 ✅ +8  12 114 💤 +1  0 ❌ ±0 

Results for commit f8c809e4. ± Comparison against base commit 4914f2f5.

[test-results]:data:application/gzip;base64,H4sIAHftGGgC/02MywrDIBBFfyW47mKMmoz9mWJ8gDSPYnQV+u81EqaB2ZxzmXOwEGe/s2enHx3bS8wEriST47ZW5L0AXlUd8zlzhJ74tRdrTzkgjH/5jp/2ySW5YOJcHZDwKW3pMqmsLT3wfhAXU1rqeiQp3drN3dKN72W7LUvMFVhAi6C9lGoMSgvrjFJKIEg/TdaA4w5G5KjZ9wdFRN6OGAEAAA==


Patch:
diff --git a/dask/dataframe/dask_expr/_collection.py b/dask/dataframe/dask_expr/_collection.py
index 91fea7840..0e16f3ed9 100644
--- a/dask/dataframe/dask_expr/_collection.py
+++ b/dask/dataframe/dask_expr/_collection.py
@@ -1404,7 +1404,7 @@ Expr={expr}"""
         A Dask Array
         """
         if lengths is True:
-            lengths = tuple(self.map_partitions(len, enforce_metadata=False).compute())
+            lengths = tuple(self.map_partitions(len).compute())
 
         arr = self.values
 
diff --git a/dask/dataframe/dask_expr/io/parquet.py b/dask/dataframe/dask_expr/io/parquet.py
index e5e517fde..4563d238c 100644
--- a/dask/dataframe/dask_expr/io/parquet.py
+++ b/dask/dataframe/dask_expr/io/parquet.py
@@ -1408,34 +1408,6 @@ class ReadParquetFSSpec(ReadParquet):
         dataset_info["all_columns"] = all_columns
         dataset_info["calculate_divisions"] = self.calculate_divisions
 
-        self._dataset_info_cache = dataset_info
-        return dataset_info
-
-    def _filtered_task(self, name: Key, index: int) -> Task:
-        tsk = Task(name, self._io_func, self._plan["parts"][index], _data_producer=True)
-        if self._series:
-            return Task(name, operator.getitem, tsk, self.columns[0])
-        return tsk
-
-    @property
-    def _io_func(self):
-        if self._plan["empty"]:
-            return identity
-        dataset_info = self._dataset_info
-        return ParquetFunctionWrapper(
-            self.engine,
-            dataset_info["fs"],
-            dataset_info["base_meta"],
-            self.columns,
-            dataset_info["index"],
-            dataset_info["kwargs"]["dtype_backend"],
-            {},  # All kwargs should now be in `common_kwargs`
-            self._plan["common_kwargs"],
-        )
-
-    @cached_property
-    def _plan(self):
-        dataset_info = self._dataset_info
         dataset_token = tokenize(dataset_info)
         if dataset_token not in _cached_plan:
             parts, stats, common_kwargs = self.engine._construct_collection_plan(
@@ -1459,7 +1431,7 @@ class ReadParquetFSSpec(ReadParquet):
             if len(divisions) < 2:
                 # empty dataframe - just use meta
                 divisions = (None, None)
-                parts = [self._meta]
+                parts = [meta]
                 empty = True
 
             _control_cached_plan(dataset_token)
@@ -1470,7 +1442,35 @@ class ReadParquetFSSpec(ReadParquet):
                 "divisions": divisions,
                 "common_kwargs": common_kwargs,
             }
-        return _cached_plan[dataset_token]
+        dataset_info["plan"] = _cached_plan[dataset_token]
+        self._dataset_info_cache = dataset_info
+        return dataset_info
+
+    def _filtered_task(self, name: Key, index: int) -> Task:
+        tsk = Task(name, self._io_func, self._plan["parts"][index], _data_producer=True)
+        if self._series:
+            return Task(name, operator.getitem, tsk, self.columns[0])
+        return tsk
+
+    @property
+    def _io_func(self):
+        if self._plan["empty"]:
+            return identity
+        dataset_info = self._dataset_info
+        return ParquetFunctionWrapper(
+            self.engine,
+            dataset_info["fs"],
+            dataset_info["base_meta"],
+            self.columns,
+            dataset_info["index"],
+            dataset_info["kwargs"]["dtype_backend"],
+            {},  # All kwargs should now be in `common_kwargs`
+            self._plan["common_kwargs"],
+        )
+
+    @property
+    def _plan(self):
+        return self._dataset_info["plan"]
 
     def _get_lengths(self) -> tuple | None:
         """Return known partition lengths using parquet statistics"""


Test patch:
diff --git a/dask/dataframe/dask_expr/io/tests/test_parquet.py b/dask/dataframe/dask_expr/io/tests/test_parquet.py
index 1568bcf02..35d0db215 100644
--- a/dask/dataframe/dask_expr/io/tests/test_parquet.py
+++ b/dask/dataframe/dask_expr/io/tests/test_parquet.py
@@ -42,7 +42,8 @@ def parquet_file(tmpdir):
         pytest.param(
             "arrow",
             marks=pytest.mark.xfail(not PYARROW_GE_1500, reason="requires 15.0.0"),
-        )
+        ),
+        "fsspec",
     ]
 )
 def filesystem(request):
@@ -609,3 +610,29 @@ def test_read_parquet_index_projection(tmpdir):
     expected = expected.assign(dts=expected.index - expected.tsprv)
 
     assert_eq(result.dts.min(), expected.dts.min())
+
+
+@pytest.mark.filterwarnings(
+    # The stats collection encounters the async client and falls back to a
+    # sync scheduler. Irrelevant for this test.
+    "ignore:.*Client.*asynchronous:UserWarning",
+)
+def test_ensure_plan_computed_during_optimization(tmpdir, filesystem):
+
+    pytest.importorskip("distributed")
+    from distributed.utils_test import gen_cluster
+
+    @gen_cluster(
+        client=True,
+        clean_kwargs={"threads": False},
+        config={"admin.async-client-fallback": "sync"},
+    )
+    async def _test(c, s, *workers):
+        # https://github.com/dask/dask/issues/11932
+
+        df = read_parquet(_make_file(tmpdir), filesystem=filesystem)
+        df = df[df.a == 1]
+        res = await c.compute(df.size)
+        assert res == 5
+
+    _test()
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
index acf7111ec..fc054c7cb 100644
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -2849,6 +2849,17 @@ def test_to_dask_array(meta, as_frame, lengths):
     assert result.chunks == expected_chunks
 
 
+def test_to_dask_array_single_partition():
+    s = pd.Series([1, 2, 3, 4, 5], name="foo", dtype="i4")
+    a = dd.from_pandas(s)
+
+    result = a.to_dask_array(lengths=True)
+    assert isinstance(result, da.Array)
+    expected_chunks = (s.shape,)
+    assert result.chunks == expected_chunks
+    da.assert_eq(s.values, result)
+
+
 def test_apply():
     df = pd.DataFrame({"x": [1, 2, 3, 4], "y": [10, 20, 30, 40]})
     ddf = dd.from_pandas(df, npartitions=2)