Repository: dask/dask

Issue:
`from_dask_array` materializes the graph unnecessarily
`dask.dataframe.from_dask_array` has many instances of `ensure_dict` and low-level dict-merging, which is both inefficient and will prevent any optimizations from being applied to those layers (without https://github.com/dask/dask/issues/8447).

Convert `from_dask_array` to use and merge HighLevelGraphs properly. xref https://github.com/dask/dask/issues/8306

https://github.com/dask/dask/blob/27a7928fb750d604dfbad0b3f4361aa64a2c97b4/dask/dataframe/io/io.py#L444-L487


Discussion:
@gjoseph92 would you be able to provide a little more detail on this in case a new contributor would like to pick it up?
I would probably not recommend this for a new contributor. After reading through the current code, I think I would approach this by rewriting it with a different approach: rather than generating graphs, you can fully express this operation using the high-level Array and DataFrame APIs, then `da.map_blocks`/`da.blockwise`/`dd.map_partitions` (not sure which yet) to combine the array with the new index and convert it to a DataFrame.

But this will require a decent understanding of blockwise, which may not be documented well enough yet for someone to be able to pick up.

Resolving patch:
diff --git a/dask/blockwise.py b/dask/blockwise.py
--- a/dask/blockwise.py
+++ b/dask/blockwise.py
@@ -35,7 +35,7 @@ class BlockwiseDep:
     ``Blockwise`` layer.
 
     All ``BlockwiseDep`` instances must define a ``numblocks``
-    attribute to speficy the number of blocks/partitions the
+    attribute to specify the number of blocks/partitions the
     object can support along each dimension. The object should
     also define a ``produces_tasks`` attribute to specify if
     any nested tasks will be passed to the Blockwise function.
@@ -129,7 +129,7 @@ class BlockwiseDepDict(BlockwiseDep):
     ...     }
     ... )
 
-    Construct a Blockwise Layer with ``dep`` speficied
+    Construct a Blockwise Layer with ``dep`` specified
     in the ``indices`` list:
 
     >>> layer = Blockwise(
diff --git a/dask/dataframe/io/io.py b/dask/dataframe/io/io.py
--- a/dask/dataframe/io/io.py
+++ b/dask/dataframe/io/io.py
@@ -5,10 +5,10 @@
 
 import numpy as np
 import pandas as pd
-from tlz import merge
 
 import dask.array as da
 from dask.base import tokenize
+from dask.blockwise import BlockwiseDepDict, blockwise
 from dask.dataframe.core import (
     DataFrame,
     Index,
@@ -25,7 +25,7 @@
 )
 from dask.delayed import delayed
 from dask.highlevelgraph import HighLevelGraph
-from dask.utils import M, _deprecated, ensure_dict
+from dask.utils import M, _deprecated
 
 lock = Lock()
 
@@ -407,6 +407,26 @@ def dataframe_from_ctable(x, slc, columns=None, categories=None, lock=lock):
     return result
 
 
+def _partition_from_array(data, index=None, initializer=None, **kwargs):
+    """Create a Dask partition for either a DataFrame or Series.
+
+    Designed to be used with :func:`dask.blockwise.blockwise`. ``data`` is the array
+    from which the partition will be created. ``index`` can be:
+
+    1. ``None``, in which case each partition has an independent RangeIndex
+    2. a `tuple` with two elements, the start and stop values for a RangeIndex for
+       this partition, which gives a continuously varying RangeIndex over the
+       whole Dask DataFrame
+    3. an instance of a ``pandas.Index`` or a subclass thereof
+
+    The ``kwargs`` _must_ contain an ``initializer`` key which is set by calling
+    ``type(meta)``.
+    """
+    if isinstance(index, tuple):
+        index = pd.RangeIndex(*index)
+    return initializer(data, index=index, **kwargs)
+
+
 def from_dask_array(x, columns=None, index=None, meta=None):
     """Create a Dask DataFrame from a Dask Array.
 
@@ -454,15 +474,14 @@ def from_dask_array(x, columns=None, index=None, meta=None):
     """
     meta = _meta_from_array(x, columns, index, meta=meta)
 
-    if x.ndim == 2 and len(x.chunks[1]) > 1:
-        x = x.rechunk({1: x.shape[1]})
-
-    name = "from-dask-array" + tokenize(x, columns)
-    to_merge = []
+    name = "from-dask-array-" + tokenize(x, columns)
+    graph_dependencies = [x]
+    arrays_and_indices = [x.name, "ij" if x.ndim == 2 else "i"]
+    numblocks = {x.name: x.numblocks}
 
     if index is not None:
-        if not isinstance(index, Index):
-            raise ValueError("'index' must be an instance of dask.dataframe.Index")
+        # An index is explicitly given by the caller, so we can pass it through to the
+        # initializer after a few checks.
         if index.npartitions != x.numblocks[0]:
             msg = (
                 "The index and array have different numbers of blocks. "
@@ -470,32 +489,48 @@ def from_dask_array(x, columns=None, index=None, meta=None):
             )
             raise ValueError(msg)
         divisions = index.divisions
-        to_merge.append(ensure_dict(index.dask))
-        index = index.__dask_keys__()
+        graph_dependencies.append(index)
+        arrays_and_indices.extend([index._name, "i"])
+        numblocks[index._name] = (index.npartitions,)
 
     elif np.isnan(sum(x.shape)):
+        # The shape of the incoming array is not known in at least one dimension. As
+        # such, we can't create an index for the entire output DataFrame and we set
+        # the divisions to None to represent that.
         divisions = [None] * (len(x.chunks[0]) + 1)
-        index = [None] * len(x.chunks[0])
     else:
+        # The shape of the incoming array is known and we don't have an explicit index.
+        # Create a mapping of chunk number in the incoming array to
+        # (start row, stop row) tuples. These tuples will be used to create a sequential
+        # RangeIndex later on that is continuous over the whole DataFrame.
         divisions = [0]
-        for c in x.chunks[0]:
-            divisions.append(divisions[-1] + c)
-        index = [
-            (np.arange, a, b, 1, "i8") for a, b in zip(divisions[:-1], divisions[1:])
-        ]
+        stop = 0
+        index_mapping = {}
+        for i, increment in enumerate(x.chunks[0]):
+            stop += increment
+            index_mapping[(i,)] = (divisions[i], stop)
+            divisions.append(stop)
         divisions[-1] -= 1
+        arrays_and_indices.extend([BlockwiseDepDict(mapping=index_mapping), "i"])
 
-    dsk = {}
-    for i, (chunk, ind) in enumerate(zip(x.__dask_keys__(), index)):
-        if x.ndim == 2:
-            chunk = chunk[0]
-        if is_series_like(meta):
-            dsk[name, i] = (type(meta), chunk, ind, x.dtype, meta.name)
-        else:
-            dsk[name, i] = (type(meta), chunk, ind, meta.columns)
+    if is_series_like(meta):
+        kwargs = {"dtype": x.dtype, "name": meta.name, "initializer": type(meta)}
+    else:
+        kwargs = {"columns": meta.columns, "initializer": type(meta)}
+
+    blk = blockwise(
+        _partition_from_array,
+        name,
+        "i",
+        *arrays_and_indices,
+        numblocks=numblocks,
+        concatenate=True,
+        # kwargs passed through to the DataFrame/Series initializer
+        **kwargs,
+    )
 
-    to_merge.extend([ensure_dict(x.dask), dsk])
-    return new_dd_object(merge(*to_merge), name, meta, divisions)
+    graph = HighLevelGraph.from_collections(name, blk, dependencies=graph_dependencies)
+    return new_dd_object(graph, name, meta, divisions)
 
 
 def _link(token, result):


Test patch:
diff --git a/dask/dataframe/io/tests/test_io.py b/dask/dataframe/io/tests/test_io.py
--- a/dask/dataframe/io/tests/test_io.py
+++ b/dask/dataframe/io/tests/test_io.py
@@ -13,6 +13,7 @@
 from dask.dataframe.utils import assert_eq, is_categorical_dtype
 from dask.delayed import Delayed, delayed
 from dask.utils import tmpfile
+from dask.utils_test import hlg_layer_topological
 
 ####################
 # Arrays and BColz #
@@ -354,36 +355,45 @@ def test_from_pandas_with_datetime_index():
 
 def test_DataFrame_from_dask_array():
     x = da.ones((10, 3), chunks=(4, 2))
-
+    pdf = pd.DataFrame(np.ones((10, 3)), columns=["a", "b", "c"])
     df = dd.from_dask_array(x, ["a", "b", "c"])
-    assert isinstance(df, dd.DataFrame)
-    tm.assert_index_equal(df.columns, pd.Index(["a", "b", "c"]))
-    assert list(df.divisions) == [0, 4, 8, 9]
-    assert (df.compute(scheduler="sync").values == x.compute(scheduler="sync")).all()
+    assert not hlg_layer_topological(df.dask, -1).is_materialized()
+    assert_eq(df, pdf)
 
     # dd.from_array should re-route to from_dask_array
     df2 = dd.from_array(x, columns=["a", "b", "c"])
-    assert isinstance(df, dd.DataFrame)
-    tm.assert_index_equal(df2.columns, df.columns)
-    assert df2.divisions == df.divisions
+    assert not hlg_layer_topological(df2.dask, -1).is_materialized()
+    assert_eq(df, df2)
+
+
+def test_DataFrame_from_dask_array_with_blockwise_ops():
+    x = da.ones((10, 3), chunks=(4, 2))
+    x *= 2
+    pdf = pd.DataFrame(np.ones((10, 3)) * 2, columns=["a", "b", "c"])
+    df = dd.from_dask_array(x, ["a", "b", "c"])
+    # None of the layers in this graph should be materialized, everything should
+    # be a HighLevelGraph still.
+    assert all(
+        not hlg_layer_topological(df.dask, i).is_materialized()
+        for i in range(len(df.dask.layers))
+    )
+    assert_eq(df, pdf)
 
 
 def test_Series_from_dask_array():
     x = da.ones(10, chunks=4)
+    pser = pd.Series(np.ones(10), name="a")
 
     ser = dd.from_dask_array(x, "a")
-    assert isinstance(ser, dd.Series)
-    assert ser.name == "a"
-    assert list(ser.divisions) == [0, 4, 8, 9]
-    assert (ser.compute(scheduler="sync").values == x.compute(scheduler="sync")).all()
+    assert_eq(ser, pser)
 
+    # Not passing a name should result in the name == None
+    pser = pd.Series(np.ones(10))
     ser = dd.from_dask_array(x)
-    assert isinstance(ser, dd.Series)
-    assert ser.name is None
+    assert_eq(ser, pser)
 
     # dd.from_array should re-route to from_dask_array
     ser2 = dd.from_array(x)
-    assert isinstance(ser2, dd.Series)
     assert_eq(ser, ser2)
 
 
@@ -398,91 +408,91 @@ def test_from_dask_array_index(as_frame):
 
 def test_from_dask_array_index_raises():
     x = da.random.uniform(size=(10,), chunks=(5,))
-    with pytest.raises(ValueError) as m:
+    with pytest.raises(ValueError, match="must be an instance"):
         dd.from_dask_array(x, index=pd.Index(np.arange(10)))
-    assert m.match("must be an instance")
 
     a = dd.from_pandas(pd.Series(range(12)), npartitions=2)
     b = dd.from_pandas(pd.Series(range(12)), npartitions=4)
-    with pytest.raises(ValueError) as m:
+    with pytest.raises(ValueError, match=".*index.*numbers of blocks.*4 != 2"):
         dd.from_dask_array(a.values, index=b.index)
 
-    assert m.match("index")
-    assert m.match("number")
-    assert m.match("blocks")
-    assert m.match("4 != 2")
-
 
-def test_from_dask_array_compat_numpy_array():
+def test_from_array_raises_more_than_2D():
     x = da.ones((3, 3, 3), chunks=2)
+    y = np.ones((3, 3, 3))
 
-    with pytest.raises(ValueError):
+    with pytest.raises(ValueError, match="more than 2D array"):
         dd.from_dask_array(x)  # dask
 
-    with pytest.raises(ValueError):
-        dd.from_array(x.compute())  # numpy
+    with pytest.raises(ValueError, match="more than 2D array"):
+        dd.from_array(y)  # numpy
+
 
+def test_from_dask_array_compat_numpy_array():
     x = da.ones((10, 3), chunks=(3, 3))
+    y = np.ones((10, 3))
     d1 = dd.from_dask_array(x)  # dask
-    assert isinstance(d1, dd.DataFrame)
-    assert (d1.compute().values == x.compute()).all()
-    tm.assert_index_equal(d1.columns, pd.Index([0, 1, 2]))
+    p1 = pd.DataFrame(y)
+    assert_eq(d1, p1)
 
-    d2 = dd.from_array(x.compute())  # numpy
-    assert isinstance(d1, dd.DataFrame)
-    assert (d2.compute().values == x.compute()).all()
-    tm.assert_index_equal(d2.columns, pd.Index([0, 1, 2]))
+    d2 = dd.from_array(y)  # numpy
+    assert_eq(d2, d1)
 
-    with pytest.raises(ValueError):
+
+def test_from_array_wrong_column_shape_error():
+    x = da.ones((10, 3), chunks=(3, 3))
+    with pytest.raises(ValueError, match="names must match width"):
         dd.from_dask_array(x, columns=["a"])  # dask
 
-    with pytest.raises(ValueError):
-        dd.from_array(x.compute(), columns=["a"])  # numpy
+    y = np.ones((10, 3))
+    with pytest.raises(ValueError, match="names must match width"):
+        dd.from_array(y, columns=["a"])  # numpy
 
+
+def test_from_array_with_column_names():
+    x = da.ones((10, 3), chunks=(3, 3))
+    y = np.ones((10, 3))
     d1 = dd.from_dask_array(x, columns=["a", "b", "c"])  # dask
-    assert isinstance(d1, dd.DataFrame)
-    assert (d1.compute().values == x.compute()).all()
-    tm.assert_index_equal(d1.columns, pd.Index(["a", "b", "c"]))
+    p1 = pd.DataFrame(y, columns=["a", "b", "c"])
+    assert_eq(d1, p1)
 
-    d2 = dd.from_array(x.compute(), columns=["a", "b", "c"])  # numpy
-    assert isinstance(d1, dd.DataFrame)
-    assert (d2.compute().values == x.compute()).all()
-    tm.assert_index_equal(d2.columns, pd.Index(["a", "b", "c"]))
+    d2 = dd.from_array(y, columns=["a", "b", "c"])  # numpy
+    assert_eq(d1, d2)
 
 
 def test_from_dask_array_compat_numpy_array_1d():
 
     x = da.ones(10, chunks=3)
+    y = np.ones(10)
     d1 = dd.from_dask_array(x)  # dask
-    assert isinstance(d1, dd.Series)
-    assert (d1.compute().values == x.compute()).all()
-    assert d1.name is None
+    p1 = pd.Series(y)
+    assert_eq(d1, p1)
+
+    d2 = dd.from_array(y)  # numpy
+    assert_eq(d2, d1)
 
-    d2 = dd.from_array(x.compute())  # numpy
-    assert isinstance(d1, dd.Series)
-    assert (d2.compute().values == x.compute()).all()
-    assert d2.name is None
 
+def test_from_array_1d_with_column_names():
+    x = da.ones(10, chunks=3)
+    y = np.ones(10)
     d1 = dd.from_dask_array(x, columns="name")  # dask
-    assert isinstance(d1, dd.Series)
-    assert (d1.compute().values == x.compute()).all()
-    assert d1.name == "name"
+    p1 = pd.Series(y, name="name")
+    assert_eq(d1, p1)
 
     d2 = dd.from_array(x.compute(), columns="name")  # numpy
-    assert isinstance(d1, dd.Series)
-    assert (d2.compute().values == x.compute()).all()
-    assert d2.name == "name"
+    assert_eq(d2, d1)
+
 
+def test_from_array_1d_list_of_columns_gives_dataframe():
+    x = da.ones(10, chunks=3)
+    y = np.ones(10)
     # passing list via columns results in DataFrame
     d1 = dd.from_dask_array(x, columns=["name"])  # dask
-    assert isinstance(d1, dd.DataFrame)
-    assert (d1.compute().values == x.compute()).all()
-    tm.assert_index_equal(d1.columns, pd.Index(["name"]))
+    p1 = pd.DataFrame(y, columns=["name"])
+    assert_eq(d1, p1)
 
-    d2 = dd.from_array(x.compute(), columns=["name"])  # numpy
-    assert isinstance(d1, dd.DataFrame)
-    assert (d2.compute().values == x.compute()).all()
-    tm.assert_index_equal(d2.columns, pd.Index(["name"]))
+    d2 = dd.from_array(y, columns=["name"])  # numpy
+    assert_eq(d2, d1)
 
 
 def test_from_dask_array_struct_dtype():
@@ -518,10 +528,12 @@ def test_from_dask_array_unknown_chunks():
     assert not df.known_divisions
     assert_eq(df, pd.DataFrame(dx.compute()), check_index=False)
 
-    # Unknown width
+
+def test_from_dask_array_unknown_width_error():
+    dsk = {("x", 0, 0): np.random.random((2, 3)), ("x", 1, 0): np.random.random((5, 3))}
     dx = da.Array(dsk, "x", ((np.nan, np.nan), (np.nan,)), np.float64)
-    with pytest.raises(ValueError):
-        df = dd.from_dask_array(dx)
+    with pytest.raises(ValueError, match="Shape along axis 1 must be known"):
+        dd.from_dask_array(dx)
 
 
 def test_to_bag():