Avoid in-memory broadcasting when converting to_dask_dataframe

[Illviljan]

Turns out that there's a call to .set_dims that forces a broadcast on the numpy coordinates.

• Closes Notebook crashes after calling .to_dask_dataframe  #6811
• Tests added, see Add benchmarks for to_dataframe and to_dask_dataframe #7474.
• User visible changes (including notable bug fixes) are documented in whats-new.rst

Debugging script:

import dask.array as da
import xarray as xr
import numpy as np

chunks = 5000

# I have to restart the pc if running with this:
# dim1_sz = 100_000
# dim2_sz = 100_000

# Does not crash when using the following constants, >5 gig RAM increase though:
dim1_sz = 40_000
dim2_sz = 40_000

x = da.random.random((dim1_sz, dim2_sz), chunks=chunks)

ds = xr.Dataset(
    {
        "x": xr.DataArray(
            data=x,
            dims=["dim1", "dim2"],
            coords={"dim1": np.arange(0, dim1_sz), "dim2": np.arange(0, dim2_sz)},
        )
    }
)

# with dask.config.set(**{"array.slicing.split_large_chunks": True}):
df = ds.to_dask_dataframe()
print(df)

[dcherian]

Suggested change

dask_array_raveled = ravel_chunks(dask_array)

Unfortunately we can't do this, at least not by default.

We could ask dask to add this behaviour as an opt-in kwarg for dask.dataframe.from_array

[Illviljan]

How come?

If we go back to using .reshape(-1) or .ravel() we will continue getting this warning:

PerformanceWarning: Reshaping is producing a large chunk. To accept the large
chunk and silence this warning, set the option
with dask.config.set(**{'array.slicing.split_large_chunks': False}):
    array.reshape(shape)

To avoid creating the large chunks, set the option
with dask.config.set(**{'array.slicing.split_large_chunks': True}):
    array.reshape(shape)Explictly passing ``limit`` to ``reshape`` will also silence this warning
array.reshape(shape, limit='128 MiB')
  exec_fun(compile(ast_code, filename, 'exec'), ns_globals, ns_locals)

[dcherian]

reshape/ravel have an implied order. With this change the ordering of rows in the output dataframe depends on the chunking of the input array, which would be confusing as default behaviour

I think the warning is fine. Users can override with the dask context manager as suggested in the warning.

[Illviljan]

Ok, I'll undo it, it wasn't necessary for the real fix anyway.

For comparison, here's the df.visualization():

With reshape:

with reshape with context {'array.slicing.split_large_chunks': True}:

With ravel_chunks:

[dcherian]

Yup, it's a big improvement. If you're dying to add it someplace, polyfit would be a good candidate (and very impactful PR).

[Illviljan]

I like these kinds of improvements :)

With ravel_chunks:

       before           after         ratio
     [3ee7b5a6]       [e549724e]
-            983M             183M     0.19  pandas.ToDataFrameDask.peakmem_to_dataframe
-         2.76±0s      7.76±0.08ms     0.00  pandas.ToDataFrameDask.time_to_dataframe

With reshape

        before           after         ratio
     [3ee7b5a6]       [02a4e97f]
-            983M             183M     0.19  pandas.ToDataFrameDask.peakmem_to_dataframe
-         2.78±0s       9.20±0.1ms     0.00  pandas.ToDataFrameDask.time_to_dataframe

[dcherian]

Thanks. Great PR!