Repository: dask/dask
Pull request #11653 (2025-01-09 11:51:53)

Rename data_producer and add flag to dataframe io stuff
- [ ] Closes #xxxx
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`

@fjetter mentioned that data_producer conflicts with potential kwargs to the function, so adding the underscore should reduce the likelihood of that happening significantly

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._

     15 files  ±0       15 suites  ±0   4h 31m 55s ⏱️ + 3m 36s
 17 151 tests ±0   15 956 ✅ ±0   1 195 💤 ±0  0 ❌ ±0 
211 304 runs  ±0  194 128 ✅  - 1  17 176 💤 +1  0 ❌ ±0 

Results for commit 9c977f18. ± Comparison against base commit ae3fc2ed.

[test-results]:data:application/gzip;base64,H4sIACPBf2cC/02Myw6DIBBFf6Vh3YXDAOP0Z4xBSEh9NAirpv9eJYouz7m55yt8GN0qXg/Qz4dYc0gXDTn2KSzzzgaL2tZUdgINJ3drtracWJtLvsNnl8D12Pk+jJtrqnAxLvEwMc97WgJgow6uaVYg2yqPNAGZ093She9lu0xTSBsItkzkoQVG4waUiD0pwy01ymvpCKQD1JrF7w+Oq48hGgEAAA==


Patch:
diff --git a/dask/_task_spec.py b/dask/_task_spec.py
index fced8892c..1f25c2ee3 100644
--- a/dask/_task_spec.py
+++ b/dask/_task_spec.py
@@ -475,7 +475,7 @@ class GraphNode:
             outkey,
             (Dict({k: Alias(k) for k in external_deps}) if external_deps else {}),
             {},
-            data_producer=any(t.data_producer for t in tasks),
+            _data_producer=any(t.data_producer for t in tasks),
         )
 
 
@@ -630,7 +630,7 @@ class Task(GraphNode):
         func: Callable,
         /,
         *args: Any,
-        data_producer: bool = False,
+        _data_producer: bool = False,
         _dependencies: set | frozenset | None = None,
         **kwargs: Any,
     ):
@@ -656,7 +656,7 @@ class Task(GraphNode):
         self._is_coro = None
         self._token = None
         self._repr = None
-        self._data_producer = data_producer
+        self._data_producer = _data_producer
 
     @property
     def data_producer(self) -> bool:
@@ -774,7 +774,7 @@ class Task(GraphNode):
                 key or self.key,
                 self.func,
                 *new_args,
-                data_producer=self.data_producer,
+                _data_producer=self.data_producer,
                 **new_kwargs,
             )
         elif key is None or key == self.key:
@@ -785,7 +785,7 @@ class Task(GraphNode):
                 key,
                 self.func,
                 *self.args,
-                data_producer=self.data_producer,
+                _data_producer=self.data_producer,
                 **self.kwargs,
             )
 
diff --git a/dask/array/core.py b/dask/array/core.py
index 3576e660d..fb15a7f79 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -312,7 +312,7 @@ def graph_from_arraylike(
             ArraySliceDep(chunks),
             out_ind,
             numblocks={},
-            data_producer=True,
+            _data_producer=True,
             **kwargs,
         )
         return HighLevelGraph.from_collections(name, layer)
@@ -330,7 +330,7 @@ def graph_from_arraylike(
             ArraySliceDep(chunks),
             out_ind,
             numblocks={},
-            data_producer=True,
+            _data_producer=True,
             **kwargs,
         )
 
diff --git a/dask/blockwise.py b/dask/blockwise.py
index ea19d0382..8abc18d60 100644
--- a/dask/blockwise.py
+++ b/dask/blockwise.py
@@ -236,7 +236,7 @@ def blockwise(
     concatenate=None,
     new_axes=None,
     dependencies=(),
-    data_producer=False,
+    _data_producer=False,
     **kwargs,
 ):
     """Create a Blockwise symbolic mutable mapping
@@ -451,7 +451,7 @@ def blockwise(
         inputs_indices.extend((None,) * len(new_keys))
 
     indices = [(k, v) for k, v in zip(inputs, inputs_indices)]
-    task = Task(output, func, *task_args, data_producer=data_producer, **kwargs)
+    task = Task(output, func, *task_args, _data_producer=_data_producer, **kwargs)
     subgraph = Blockwise(
         output,
         output_indices,
@@ -801,7 +801,7 @@ class Blockwise(Layer):
                 chunks.bind,
                 self.task,
                 TaskRef(blockwise_token(len(indices))),
-                data_producer=self.task.data_producer,
+                _data_producer=self.task.data_producer,
             )
             indices.append((TaskRef(bind_to), None))
         else:
diff --git a/dask/dataframe/dask_expr/io/_delayed.py b/dask/dataframe/dask_expr/io/_delayed.py
index 651f23ee8..f2fb8eeac 100644
--- a/dask/dataframe/dask_expr/io/_delayed.py
+++ b/dask/dataframe/dask_expr/io/_delayed.py
@@ -65,6 +65,7 @@ class FromDelayed(PartitionsFiltered, BlockwiseIO):
                 name,
                 functools.partial(check_meta, meta=self._meta, funcname="from_delayed"),
                 TaskRef((self.delayed_container._name, index)),
+                _data_producer=True,
             )
         else:
             return Alias((self.delayed_container._name, index))  # type: ignore
diff --git a/dask/dataframe/dask_expr/io/io.py b/dask/dataframe/dask_expr/io/io.py
index 7394559e6..5c85c31e9 100644
--- a/dask/dataframe/dask_expr/io/io.py
+++ b/dask/dataframe/dask_expr/io/io.py
@@ -7,7 +7,7 @@ import operator
 import numpy as np
 import pyarrow as pa
 
-from dask._task_spec import List, Task
+from dask._task_spec import DataNode, List, Task
 from dask.dataframe import methods
 from dask.dataframe._pyarrow import to_pyarrow_string
 from dask.dataframe.core import apply_and_enforce, is_dataframe_like
@@ -140,7 +140,10 @@ class FusedIO(BlockwiseIO):
         bucket = self._fusion_buckets[index]
         # FIXME: This will likely require a wrapper
         return Task(
-            name, methods.concat, List(*(expr._filtered_task(name, i) for i in bucket))
+            name,
+            methods.concat,
+            List(*(expr._filtered_task(name, i) for i in bucket)),
+            _data_producer=True,
         )
 
     @functools.cached_property
@@ -216,6 +219,7 @@ class FusedParquetIO(FusedIO):
             columns,
             schema,
             **to_pandas_kwargs,
+            _data_producer=True,
         )
 
 
@@ -295,8 +299,11 @@ class FromMap(PartitionsFiltered, BlockwiseIO):
                 *vals,
                 *self.args,
                 **self.apply_kwargs,
+                _data_producer=True,
             )
-        return Task(name, self.func, *vals, *self.args, **self.apply_kwargs)
+        return Task(
+            name, self.func, *vals, *self.args, **self.apply_kwargs, _data_producer=True
+        )
 
 
 class FromMapProjectable(FromMap):
@@ -509,14 +516,16 @@ class FromPandas(PartitionsFiltered, BlockwiseIO):
     def _locations(self):
         return self._divisions_and_locations[1]
 
-    def _filtered_task(self, name: Key, index: int) -> Task:
+    def _filtered_task(self, name: Key, index: int) -> DataNode:  # type: ignore
         start, stop = self._locations()[index : index + 2]
         part = self.frame.iloc[start:stop]
         if self.pyarrow_strings_enabled:
             part = to_pyarrow_string(part)
         if self.operand("columns") is not None:
-            return part[self.columns[0]] if self._series else part[self.columns]
-        return part
+            return DataNode(
+                name, part[self.columns[0]] if self._series else part[self.columns]
+            )
+        return DataNode(name, part)
 
     def __str__(self):
         if self._absorb_projections and self.operand("columns"):
@@ -662,9 +671,22 @@ class FromArray(PartitionsFiltered, BlockwiseIO):
 
         if is_series_like(self._meta):
             return Task(
-                name, type(self._meta), data, idx, self._meta.dtype, self._meta.name
+                name,
+                type(self._meta),
+                data,
+                idx,
+                self._meta.dtype,
+                self._meta.name,
+                _data_producer=True,
             )
         else:
             if data.ndim == 2:
                 data = data[:, self._column_indices]
-            return Task(name, type(self._meta), data, idx, self._meta.columns)
+            return Task(
+                name,
+                type(self._meta),
+                data,
+                idx,
+                self._meta.columns,
+                _data_producer=True,
+            )
diff --git a/dask/dataframe/dask_expr/io/parquet.py b/dask/dataframe/dask_expr/io/parquet.py
index 5413a99ed..64070f0d9 100644
--- a/dask/dataframe/dask_expr/io/parquet.py
+++ b/dask/dataframe/dask_expr/io/parquet.py
@@ -1181,6 +1181,7 @@ class ReadParquetPyarrowFS(ReadParquet):
             arrow_to_pandas=self.arrow_to_pandas,
             dtype_backend=self.kwargs.get("dtype_backend"),
             pyarrow_strings_enabled=self.pyarrow_strings_enabled,
+            _data_producer=True,
         )
 
     @staticmethod
@@ -1401,7 +1402,7 @@ class ReadParquetFSSpec(ReadParquet):
         return dataset_info
 
     def _filtered_task(self, name: Key, index: int) -> Task:
-        tsk = Task(name, self._io_func, self._plan["parts"][index])
+        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


Test patch:
diff --git a/dask/dataframe/dask_expr/io/tests/test_from_pandas.py b/dask/dataframe/dask_expr/io/tests/test_from_pandas.py
index 5144320bf..dab466b66 100644
--- a/dask/dataframe/dask_expr/io/tests/test_from_pandas.py
+++ b/dask/dataframe/dask_expr/io/tests/test_from_pandas.py
@@ -25,6 +25,7 @@ def test_from_pandas(pdf, sort):
     assert df.npartitions == 2
     assert df.divisions == (0, 3, 5)
     assert_eq(df, pdf, sort_results=sort)
+    assert all(tsk.data_producer for tsk in df.dask.values())
 
 
 def test_from_pandas_noargs(pdf):
diff --git a/dask/dataframe/dask_expr/io/tests/test_io.py b/dask/dataframe/dask_expr/io/tests/test_io.py
index 6866e7d6f..42761ce81 100644
--- a/dask/dataframe/dask_expr/io/tests/test_io.py
+++ b/dask/dataframe/dask_expr/io/tests/test_io.py
@@ -100,6 +100,11 @@ def test_io_fusion(tmpdir, fmt):
     # All tasks should be fused for each partition
     assert len(df2.dask) == df2.npartitions
     assert_eq(df2, df[["a", "b"]] + 1)
+    assert all(
+        tsk.data_producer
+        for k, tsk in df.dask.items()
+        if not k[0].startswith("_to_string")
+    )
 
 
 def test_csv_integer_usecols(tmpdir):
@@ -226,6 +231,7 @@ def test_parquet_complex_filters(tmpdir):
 
     assert_eq(got, expect)
     assert_eq(got.optimize(), expect)
+    assert all(tsk.data_producer for tsk in df.dask.values())
 
 
 @pytest.mark.parametrize("optimize", [True, False])
@@ -309,6 +315,17 @@ def test_from_map(tmpdir, meta, label, allow_projection, enforce_metadata):
     assert_eq(df["a"], pdf["a"], check_index=False)
     assert_eq(df[["a"]], pdf[["a"]], check_index=False)
     assert_eq(df[["a", "b"]], pdf[["a", "b"]], check_index=False)
+    assert all(
+        tsk.data_producer
+        for key, tsk in df.__dask_graph__().items()
+        if not key[0].startswith("_to_string")
+    )
+    if allow_projection:
+        assert all(
+            tsk.data_producer
+            for key, tsk in df[["a"]].optimize(fuse=False).__dask_graph__().items()
+            if not key[0].startswith("_to_string")
+        )
 
     if label:
         if pyarrow_strings_enabled():
diff --git a/dask/dataframe/dask_expr/io/tests/test_parquet.py b/dask/dataframe/dask_expr/io/tests/test_parquet.py
index 50648f923..bc39fb963 100644
--- a/dask/dataframe/dask_expr/io/tests/test_parquet.py
+++ b/dask/dataframe/dask_expr/io/tests/test_parquet.py
@@ -133,6 +133,8 @@ def test_pyarrow_filesystem(parquet_file):
     df_pa = read_parquet(parquet_file, filesystem=filesystem)
     df = read_parquet(parquet_file)
     assert assert_eq(df, df_pa)
+    assert all(tsk.data_producer for tsk in df.dask.values())
+    assert all(tsk.data_producer for tsk in df.optimize().dask.values())
 
 
 @pytest.mark.xfail(not PYARROW_GE_1500, reason="requires 15.0.0")
diff --git a/dask/tests/test_task_spec.py b/dask/tests/test_task_spec.py
index ec7c1c01e..6cd2d215c 100644
--- a/dask/tests/test_task_spec.py
+++ b/dask/tests/test_task_spec.py
@@ -24,7 +24,6 @@ from dask._task_spec import (
     parse_input,
     resolve_aliases,
 )
-from dask.base import tokenize
 from dask.core import keys_in_tasks, reverse_dict
 from dask.sizeof import sizeof
 from dask.tokenize import tokenize
@@ -970,8 +969,8 @@ def test_block_io_fusion():
 def test_data_producer():
     tasks = [
         io := DataNode("foo", 1),
-        second := Task("second", func, io.ref(), data_producer=True),
-        third := Task("third", func, io.ref(), data_producer=True),
+        second := Task("second", func, io.ref(), _data_producer=True),
+        third := Task("third", func, io.ref(), _data_producer=True),
         fourth := Task("fourth", func, second.ref()),
         fifth := Task("fifth", func, third.ref()),
         Task("sixth", func, fourth.ref(), fifth.ref()),