Repository: dask/dask

Issue:
Dataframe correlation function uses too much memory
Whenever I am calculating corr on a dask data frame with a larger but not enormous number of columns (200-2000) my worker memory is blowing up and getting  a MemoryError at one of two lines:

```
from /dask/dataframe/core.py

3885     x = np.where(keep, mat[:, None, :], np.nan)
3886     sums = np.nansum(x, 0)
```

This is reliably reproduced  with the code below, exact breaking point will obviously vary with worker,  I'm currently using 4 thread, 24GB workers.

```
arr = da.random.randint(100, size=(10_000, 500), chunks=1_000)
df  = dd.from_array(arr, columns=[f'col_{i}' for i in range(arr.shape[1])])
df.to_parquet(path)
df = dd.read_parquet(path)
corr = df.corr().compute()
```

Workers are reporting " distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 22.55 GB -- Worker memory limit: 24.00 GB"

Dask will break with memory errors at sizes small enough to process the entire calculation in memory with Pandas. I will need to scale to 10's of millions of rows and 1-3 thousand columns so staying in memory doesn't look to a be a good option.


Discussion:
Thanks for opening up this issue @jhb9! 

I'm able to reproduce similar worker memory warnings with the following example:

<details>
<summary> Example details </summary>


```python
import numpy as np
import dask.array as da
import dask.dataframe as dd
from dask.distributed import Client

client = Client(processes=False,
                n_workers=1,
                threads_per_worker=4,
                memory_limit='2GB')

# Generate sample DataFrame and write to disk
np.random.seed(3)
a = da.random.randint(100, size=(10_000, 500), chunks=1_000)
ddf = dd.from_array(a, columns=[f'col_{i}' for i in range(a.shape[1])])
ddf.to_hdf('test_data.hdf', 'dataframe', mode='w')

# Read in DataFrame
ddf = dd.read_hdf('test_data.hdf', 'dataframe', chunksize=1_000)

# Compute correlation matrix
correlation = ddf.corr()
corr = correlation.compute()
```
</details>

<br>

Running this on my laptop gives warning messages like 

```
distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 2.57 GB -- Worker memory limit: 2.00 GB
```

I tried specifying `split_every` which could help free up memory on the workers by performing a tree-reduction on the correlation calculation dask graph. But unfortunately I'm still get the same worker memory warning. 

Looking at the the `cov_corr_chunk()` function (in `dask/dataframe/core.py`) that is called on each partition of the DataFrame, it looks like there are NumPy arrays created that are of shape `(n_rows, n_cols, n_cols)` for each DataFrame partition. Perhaps this could be the memory bottleneck here. 

Maybe reducing the partition size of the DataFrame and setting `split_every=2` with the synchronous scheduler could help?


cc @jcrist @mrocklin 
> it looks like there are NumPy arrays created that are of shape (n_rows, n_cols, n_cols) for each DataFrame partition. Perhaps this could be the memory bottleneck herecd

That makes sense to me.  My guess is that there is a way to implement this function in smaller memory.  
I'm going to mark this as good first issue, not necessarily because it's easy (it's actually quite an interesting problem) but rather because it requires little knowledge of dask itself.  The goal is to reimplement the `cov_corr_chunk` function in `dask/dataframe/core.py` in a way that produces the same result, but uses less memory.  This probably requires some good skill with numpy or pandas.

Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -3893,18 +3893,27 @@ def cov_corr(df, min_periods=None, corr=False, scalar=False, split_every=False):
 
 
 def cov_corr_chunk(df, corr=False):
-    """Chunk part of a covariance or correlation computation"""
-    mat = df.astype('float64', copy=False).values
-    mask = np.isfinite(mat)
-    keep = np.bitwise_and(mask[:, None, :], mask[:, :, None])
-
-    x = np.where(keep, mat[:, None, :], np.nan)
-    sums = np.nansum(x, 0)
-    counts = keep.astype('int').sum(0)
+    """Chunk part of a covariance or correlation computation
+    """
+    shape = (df.shape[1], df.shape[1])
+    sums = np.zeros(shape)
+    counts = np.zeros(shape)
+    df = df.astype('float64', copy=False)
+    for idx, col in enumerate(df):
+        mask = df[col].notnull()
+        sums[idx] = df[mask].sum().values
+        counts[idx] = df[mask].count().values
     cov = df.cov().values
     dtype = [('sum', sums.dtype), ('count', counts.dtype), ('cov', cov.dtype)]
     if corr:
-        m = np.nansum((x - sums / np.where(counts, counts, np.nan)) ** 2, 0)
+        mu = (sums / counts).T
+        m = np.zeros(shape)
+        mask = df.isnull().values
+        for idx, x in enumerate(df):
+            mu_discrepancy = np.subtract.outer(df[x], mu[idx]) ** 2
+            mu_discrepancy[mask] = np.nan
+            m[idx] = np.nansum(mu_discrepancy, axis=0)
+        m = m.T
         dtype.append(('m', m.dtype))
 
     out = np.empty(counts.shape, dtype=dtype)


Test patch: