Repository: dask/dask

Issue:
Add optimize_graph flag to Bag `to_dataframe` function
This functionality is necessary to avoid the loss of annotations during graph optimization. Consider the following snippet:
```
with dask.annotate(resources={"GPU": 1}):
    result = bag.map_partitions(func)

df = result.to_dataframe()
```
In some cases `func` will not be executed on a worker with at least 1 GPU resource, because the resource annotation is lost during the graph optimization that occurs in `to_dataframe`.


Discussion:

Resolving patch:
diff --git a/dask/bag/core.py b/dask/bag/core.py
--- a/dask/bag/core.py
+++ b/dask/bag/core.py
@@ -1531,7 +1531,7 @@ def groupby(
             msg = "Shuffle must be 'disk' or 'tasks'"
             raise NotImplementedError(msg)
 
-    def to_dataframe(self, meta=None, columns=None):
+    def to_dataframe(self, meta=None, columns=None, optimize_graph=True):
         """Create Dask Dataframe from a Dask Bag.
 
         Bag should contain tuples, dict records, or scalars.
@@ -1559,6 +1559,10 @@ def to_dataframe(self, meta=None, columns=None):
             result (any names not found in the data will become all-NA
             columns).  Note that if ``meta`` is provided, column names will be
             taken from there and this parameter is invalid.
+        optimize_graph : bool, optional
+            If True [default], the graph is optimized before converting into
+            :class:`dask.dataframe.DataFrame`.
+
 
         Examples
         --------
@@ -1596,14 +1600,15 @@ def to_dataframe(self, meta=None, columns=None):
         # the empty frame
         cols = list(meta.columns)
         dtypes = meta.dtypes.to_dict()
-        name = "to_dataframe-" + tokenize(self, cols, dtypes)
-        dsk = self.__dask_optimize__(self.dask, self.__dask_keys__())
 
-        for i in range(self.npartitions):
-            dsk[(name, i)] = (to_dataframe, (self.name, i), cols, dtypes)
+        dfs = self.map_partitions(to_dataframe, cols, dtypes)
+        if optimize_graph:
+            dsk = self.__dask_optimize__(dfs.dask, dfs.__dask_keys__())
+        else:
+            dsk = dfs.dask
 
         divisions = [None] * (self.npartitions + 1)
-        return dd.DataFrame(dsk, name, meta, divisions)
+        return dd.DataFrame(dsk, dfs.name, meta, divisions)
 
     def to_delayed(self, optimize_graph=True):
         """Convert into a list of ``dask.delayed`` objects, one per partition.


Test patch:
diff --git a/dask/bag/tests/test_bag.py b/dask/bag/tests/test_bag.py
--- a/dask/bag/tests/test_bag.py
+++ b/dask/bag/tests/test_bag.py
@@ -31,7 +31,7 @@
 from dask.bag.utils import assert_eq
 from dask.delayed import Delayed
 from dask.utils import filetexts, tmpdir, tmpfile
-from dask.utils_test import add, inc
+from dask.utils_test import add, hlg_layer_topological, inc
 
 dsk = {("x", 0): (range, 5), ("x", 1): (range, 5), ("x", 2): (range, 5)}
 
@@ -1587,6 +1587,13 @@ def test_dask_layers_to_delayed(optimize):
     # `da.Array.to_delayed` causes the layer name to not match the key.
     # Ensure the layer name is propagated between `Delayed` and `Item`.
     da = pytest.importorskip("dask.array")
+    i = db.Item.from_delayed(da.ones(1).to_delayed()[0])
+    name = i.key[0]
+    assert i.key[1:] == (0,)
+    assert i.dask.layers.keys() == {"delayed-" + name}
+    assert i.dask.dependencies == {"delayed-" + name: set()}
+    assert i.__dask_layers__() == ("delayed-" + name,)
+
     arr = da.ones(1) + 1
     delayed = arr.to_delayed(optimize_graph=optimize)[0]
     i = db.Item.from_delayed(delayed)
@@ -1605,3 +1612,39 @@ def test_dask_layers_to_delayed(optimize):
 
     with pytest.raises(ValueError, match="not in"):
         db.Item(arr.dask, (arr.name,), layer="foo")
+
+
+def test_to_dataframe_optimize_graph():
+    pytest.importorskip("dask.dataframe")
+    from dask.dataframe.utils import assert_eq as assert_eq_df
+
+    x = db.from_sequence(
+        [{"name": "test1", "v1": 1}, {"name": "test2", "v1": 2}], npartitions=2
+    )
+
+    # linear `map` tasks will be fused by graph optimization
+    with dask.annotate(foo=True):
+        y = x.map(lambda a: dict(**a, v2=a["v1"] + 1))
+        y = y.map(lambda a: dict(**a, v3=a["v2"] + 1))
+        y = y.map(lambda a: dict(**a, v4=a["v3"] + 1))
+
+    # verifying the maps are not fused yet
+    assert len(y.dask) == y.npartitions * 4
+
+    # with optimizations
+    d = y.to_dataframe()
+
+    # All the `map` tasks have been fused
+    assert len(d.dask) < len(y.dask)
+
+    # no optimizations
+    d2 = y.to_dataframe(optimize_graph=False)
+
+    # Graph hasn't been fused. It contains all the original tasks,
+    # plus one extra layer converting to DataFrame
+    assert len(d2.dask) == len(y.dask) + d.npartitions
+
+    # Annotations are still there
+    assert hlg_layer_topological(d2.dask, 1).annotations == {"foo": True}
+
+    assert_eq_df(d, d2)