Repository: dask/dask

Issue:
`to_parquet` fails with dataframe built from futures
<!-- Please include a self-contained copy-pastable example that generates the issue if possible.

Please be concise with code posted. See guidelines below on how to provide a good bug report:

- Craft Minimal Bug Reports http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
- Minimal Complete Verifiable Examples https://stackoverflow.com/help/mcve

Bug reports that follow these guidelines are easier to diagnose, and so are often handled much more quickly.
-->

**What happened**:
Trying to write a dask dataframe built from pandas futures using the `to_parquet` method fails.

**What you expected to happen**:
Writing the dataframe successfully (the `to_csv` method works).

**Minimal Complete Verifiable Example**:

```python
import dask.dataframe as dd
import pandas as pd
from dask.distributed import Client


def create_partition(i):
    return pd.DataFrame({'x': [i]})

client = Client()
parts = [client.submit(create_partition, i) for i in range(2)]
dd.from_delayed(parts).to_parquet('output')
```
<details>
    <summary>Stacktrace</summary>

```python
distributed.worker - WARNING - Compute Failed
Function:  check_meta
args:      ('create_partition-d32679d122607182bc4efd4c9af07f3a', Empty DataFrame
Columns: [x]
Index: [], 'from_delayed')
kwargs:    {}
Exception: ValueError('Metadata mismatch found in `from_delayed`.\n\nExpected partition of type `pandas.core.frame.DataFrame` but got `str`')

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/tmp/ipykernel_30101/3007309697.py in <module>
      9 client = Client()
     10 parts = [client.submit(create_partition, i) for i in range(2)]
---> 11 dd.from_delayed(parts).to_parquet('output')

~/miniconda3/envs/mlforecast/lib/python3.9/site-packages/dask/dataframe/core.py in to_parquet(self, path, *args, **kwargs)
   4538         from .io import to_parquet
   4539 
-> 4540         return to_parquet(self, path, *args, **kwargs)
   4541 
   4542     def to_orc(self, path, *args, **kwargs):

~/miniconda3/envs/mlforecast/lib/python3.9/site-packages/dask/dataframe/io/parquet/core.py in to_parquet(df, path, engine, compression, write_index, append, overwrite, ignore_divisions, partition_on, storage_options, custom_metadata, write_metadata_file, compute, compute_kwargs, schema, **kwargs)
    723     if compute:
    724         if write_metadata_file:
--> 725             return compute_as_if_collection(
    726                 DataFrame, graph, (final_name, 0), **compute_kwargs
    727             )

~/miniconda3/envs/mlforecast/lib/python3.9/site-packages/dask/base.py in compute_as_if_collection(cls, dsk, keys, scheduler, get, **kwargs)
    313     schedule = get_scheduler(scheduler=scheduler, cls=cls, get=get)
    314     dsk2 = optimization_function(cls)(dsk, keys, **kwargs)
--> 315     return schedule(dsk2, keys, **kwargs)
    316 
    317 

~/miniconda3/envs/mlforecast/lib/python3.9/site-packages/distributed/client.py in get(self, dsk, keys, workers, allow_other_workers, resources, sync, asynchronous, direct, retries, priority, fifo_timeout, actors, **kwargs)
   2687                     should_rejoin = False
   2688             try:
-> 2689                 results = self.gather(packed, asynchronous=asynchronous, direct=direct)
   2690             finally:
   2691                 for f in futures.values():

~/miniconda3/envs/mlforecast/lib/python3.9/site-packages/distributed/client.py in gather(self, futures, errors, direct, asynchronous)
   1964             else:
   1965                 local_worker = None
-> 1966             return self.sync(
   1967                 self._gather,
   1968                 futures,

~/miniconda3/envs/mlforecast/lib/python3.9/site-packages/distributed/client.py in sync(self, func, asynchronous, callback_timeout, *args, **kwargs)
    858             return future
    859         else:
--> 860             return sync(
    861                 self.loop, func, *args, callback_timeout=callback_timeout, **kwargs
    862             )

~/miniconda3/envs/mlforecast/lib/python3.9/site-packages/distributed/utils.py in sync(loop, func, callback_timeout, *args, **kwargs)
    324     if error[0]:
    325         typ, exc, tb = error[0]
--> 326         raise exc.with_traceback(tb)
    327     else:
    328         return result[0]

~/miniconda3/envs/mlforecast/lib/python3.9/site-packages/distributed/utils.py in f()
    307             if callback_timeout is not None:
    308                 future = asyncio.wait_for(future, callback_timeout)
--> 309             result[0] = yield future
    310         except Exception:
    311             error[0] = sys.exc_info()

~/miniconda3/envs/mlforecast/lib/python3.9/site-packages/tornado/gen.py in run(self)
    760 
    761                     try:
--> 762                         value = future.result()
    763                     except Exception:
    764                         exc_info = sys.exc_info()

~/miniconda3/envs/mlforecast/lib/python3.9/site-packages/distributed/client.py in _gather(self, futures, errors, direct, local_worker)
   1829                             exc = CancelledError(key)
   1830                         else:
-> 1831                             raise exception.with_traceback(traceback)
   1832                         raise exc
   1833                     if errors == "skip":

~/miniconda3/envs/mlforecast/lib/python3.9/site-packages/dask/dataframe/utils.py in check_meta()
    404         )
    405 
--> 406     raise ValueError(
    407         "Metadata mismatch found%s.\n\n"
    408         "%s" % ((" in `%s`" % funcname if funcname else ""), errmsg)

ValueError: Metadata mismatch found in `from_delayed`.

Expected partition of type `pandas.core.frame.DataFrame` but got `str`
```
</details>

**Anything else we need to know?**:
I believe this was introduced in #7968

**Environment**:
- Dask version: 2021.9.1
- Python version: 3.9.7
- Operating System: Ubuntu
- Install method (conda, pip, source): conda



Discussion:
Thanks for reporting this @jmoralez. I am able to reproduce, but I do wonder if passing futures rather than delayed objects to `dd.from_delayed` was ever intentionally supported. Maybe @jrbourbeau or @jcrist will know the answer to that question.
Thank you, I originally got that from https://stackoverflow.com/questions/48792589/how-do-i-convert-a-list-of-pandas-futures-to-a-dask-dataframe. Although it does feel strange, is there another way to create a dataframe from futures?
Yeah, I don't think we ever intentionally supported implicitly converting futures to delayed objects (and I'm not sure if we'd want to). Explicitly calling `delayed` on each future does work though:

```python
import dask
import dask.dataframe as dd
import pandas as pd
from dask.distributed import Client


def create_partition(i):
    return pd.DataFrame({'x': [i]})

client = Client()
parts = [client.submit(create_partition, i) for i in range(2)]
dd.from_delayed([dask.delayed(p) for p in parts]).to_parquet('output')
```

Note that if your code looks like the above, you're better off dropping the `submit` calls and relying on `dask.delayed` to manage your `create_partition` calls.

```python
import dask
import dask.dataframe as dd
import pandas as pd
from dask.distributed import Client

@dask.delayed
def create_partition(i):
    return pd.DataFrame({'x': [i]})

client = Client()
parts = [create_partition(i) for i in range(2)]
dd.from_delayed(parts).to_parquet('output')
```

Also note that you probably want to provide a `meta=` kwarg to `from_delayed`, otherwise the first partition will be computed eagerly to infer the schema (which _may_ be inefficient, or may not matter - depending on what you're doing).
Yes I currently call `.head()` on the first partition to get the meta. Thank you very much!
@jcrist I'm sorry about this but I didn't test the first snippet (where each future gets delayed) and it fails with the same error for me. I use other futures as arguments to the `create_partition` function so a solution like that one would be very nice. Is there another way?
Huh yeah, I am re-opening this issue.
Ah, sorry, I had assumed since the `from_delayed` call worked the df was fine, I hadn't actually tested the `to_parquet` call. I think this is related to #8173, I'll start on a fix for this today. For now though you can manually disable graph optimizations for the `to_parquet` call and that should fix things. Note that wrapping futures in `delayed` isn't necessary, that does happen automatically for `dd.from_delayed`.

```python
import dask
import dask.dataframe as dd
import pandas as pd
from dask.distributed import Client

def create_partition(i):
    return pd.DataFrame({'x': [i]})


if __name__ == "__main__":
    client = Client()
    parts = [client.submit(create_partition, i) for i in range(2)]

    df = dd.from_delayed(parts, meta={"x": int})
    # Skip computing automatically, then call `compute` manually with optimizations disabled
    df.to_parquet("output", compute=False).compute()
```

I've checked, and this does work (at least on main).
Hmmm, actually this isn't #8173. I can confirm that the above does work, but it's not related to graph optimization. I've amended the workaround above to optimize like normal. Looking deeper now.

Resolving patch:
diff --git a/dask/dataframe/io/parquet/core.py b/dask/dataframe/io/parquet/core.py
--- a/dask/dataframe/io/parquet/core.py
+++ b/dask/dataframe/io/parquet/core.py
@@ -718,17 +718,12 @@ def to_parquet(
     else:
         dsk[(final_name, 0)] = (lambda x: None, part_tasks)
 
-    graph = HighLevelGraph.from_collections(name, dsk, dependencies=[df])
+    graph = HighLevelGraph.from_collections(final_name, dsk, dependencies=[df])
 
     if compute:
-        if write_metadata_file:
-            return compute_as_if_collection(
-                DataFrame, graph, (final_name, 0), **compute_kwargs
-            )
-        else:
-            return compute_as_if_collection(
-                DataFrame, graph, part_tasks, **compute_kwargs
-            )
+        return compute_as_if_collection(
+            Scalar, graph, [(final_name, 0)], **compute_kwargs
+        )
     else:
         return Scalar(graph, final_name, "")
 
diff --git a/dask/highlevelgraph.py b/dask/highlevelgraph.py
--- a/dask/highlevelgraph.py
+++ b/dask/highlevelgraph.py
@@ -373,17 +373,25 @@ def __dask_distributed_pack__(
 
         # Find aliases not in `client_keys` and substitute all matching keys
         # with its Future
-        values = {
+        future_aliases = {
             k: v
             for k, v in dsk.items()
             if isinstance(v, Future) and k not in client_keys
         }
-        if values:
-            dsk = subs_multiple(dsk, values)
+        if future_aliases:
+            dsk = subs_multiple(dsk, future_aliases)
 
-        # Unpack remote data and record its dependencies
-        dsk = {k: unpack_remotedata(v, byte_keys=True) for k, v in dsk.items()}
-        unpacked_futures = set.union(*[v[1] for v in dsk.values()]) if dsk else set()
+        # Remove `Future` objects from graph and note any future dependencies
+        dsk2 = {}
+        fut_deps = {}
+        for k, v in dsk.items():
+            dsk2[k], futs = unpack_remotedata(v, byte_keys=True)
+            if futs:
+                fut_deps[k] = futs
+        dsk = dsk2
+
+        # Check that any collected futures are valid
+        unpacked_futures = set.union(*fut_deps.values()) if fut_deps else set()
         for future in unpacked_futures:
             if future.client is not client:
                 raise ValueError(
@@ -391,21 +399,26 @@ def __dask_distributed_pack__(
                 )
             if stringify(future.key) not in client.futures:
                 raise CancelledError(stringify(future.key))
-        unpacked_futures_deps = {}
-        for k, v in dsk.items():
-            if len(v[1]):
-                unpacked_futures_deps[k] = {f.key for f in v[1]}
-        dsk = {k: v[0] for k, v in dsk.items()}
 
         # Calculate dependencies without re-calculating already known dependencies
-        missing_keys = dsk.keys() - known_key_dependencies.keys()
-        dependencies = {
-            k: keys_in_tasks(all_hlg_keys, [dsk[k]], as_list=False)
+        # - Start with known dependencies
+        dependencies = known_key_dependencies.copy()
+        # - Remove aliases for any tasks that depend on both an alias and a future.
+        #   These can only be found in the known_key_dependencies cache, since
+        #   any dependencies computed in this method would have already had the
+        #   aliases removed.
+        if future_aliases:
+            alias_keys = set(future_aliases)
+            dependencies = {k: v - alias_keys for k, v in dependencies.items()}
+        # - Add in deps for any missing keys
+        missing_keys = dsk.keys() - dependencies.keys()
+        dependencies.update(
+            (k, keys_in_tasks(all_hlg_keys, [dsk[k]], as_list=False))
             for k in missing_keys
-        }
-        for k, v in unpacked_futures_deps.items():
-            dependencies[k] = set(dependencies.get(k, ())) | v
-        dependencies.update(known_key_dependencies)
+        )
+        # - Add in deps for any tasks that depend on futures
+        for k, futures in fut_deps.items():
+            dependencies[k].update(f.key for f in futures)
 
         # The scheduler expect all keys to be strings
         dependencies = {
@@ -1122,7 +1135,6 @@ def __dask_distributed_unpack__(hlg: dict) -> dict:
 
             # Unpack the annotations
             unpack_anno(anno, layer["annotations"], unpacked_layer["dsk"].keys())
-
         return {"dsk": dsk, "deps": deps, "annotations": anno}
 
     def __repr__(self) -> str:


Test patch:
diff --git a/dask/tests/test_distributed.py b/dask/tests/test_distributed.py
--- a/dask/tests/test_distributed.py
+++ b/dask/tests/test_distributed.py
@@ -322,12 +322,19 @@ def test_blockwise_array_creation(c, io, fuse):
     ["parquet-pyarrow", "parquet-fastparquet", "csv", "hdf"],
 )
 @pytest.mark.parametrize("fuse", [True, False, None])
-def test_blockwise_dataframe_io(c, tmpdir, io, fuse):
+@pytest.mark.parametrize("from_futures", [True, False])
+def test_blockwise_dataframe_io(c, tmpdir, io, fuse, from_futures):
     pd = pytest.importorskip("pandas")
     dd = pytest.importorskip("dask.dataframe")
 
     df = pd.DataFrame({"x": [1, 2, 3] * 5, "y": range(15)})
-    ddf0 = dd.from_pandas(df, npartitions=3)
+
+    if from_futures:
+        parts = [df.iloc[:5], df.iloc[5:10], df.iloc[10:15]]
+        futs = c.scatter(parts)
+        ddf0 = dd.from_delayed(futs, meta=parts[0])
+    else:
+        ddf0 = dd.from_pandas(df, npartitions=3)
 
     if io.startswith("parquet"):
         if io == "parquet-pyarrow":
@@ -631,3 +638,20 @@ async def test_annotation_pack_unpack(c, s, a, b):
     unpacked_hlg = HighLevelGraph.__dask_distributed_unpack__(packed_hlg)
     annotations = unpacked_hlg["annotations"]
     assert annotations == {"workers": {"n": ("alice",)}}
+
+
+@gen_cluster(client=True)
+async def test_pack_MaterializedLayer_handles_futures_in_graph_properly(c, s, a, b):
+    fut = c.submit(inc, 1)
+
+    hlg = HighLevelGraph(
+        {"l1": MaterializedLayer({"x": fut, "y": (inc, "x"), "z": (inc, "y")})},
+        {"l1": set()},
+    )
+    # fill hlg.key_dependencies cache. This excludes known futures, so only
+    # includes a subset of all dependencies. Previously if the cache was present
+    # the future dependencies would be missing when packed.
+    hlg.get_all_dependencies()
+    packed = hlg.__dask_distributed_pack__(c, ["z"], {})
+    unpacked = HighLevelGraph.__dask_distributed_unpack__(packed)
+    assert unpacked["deps"] == {"x": {fut.key}, "y": {fut.key}, "z": {"y"}}