Repository: pydata/xarray

Issue:
Improve performance of xarray.corr() on big datasets
**Is your feature request related to a problem? Please describe.**

I calculated correlation coefficients based on datasets with sizes between 90-180 GB using xarray and Dask distributed and experienced very low performance for the `xarray.corr()` function. By observing the Dask dashboard it seemed that during the calculation the whole datasets are loaded from disk several times which, given the size of my datasets, became for some of the calculations a major "performance bottleneck".

**Describe the solution you'd like**

The problem became so annoying that I implemented my own function to calculate the correlation coefficient (thanks @willirath!), which is considerably more performant (especially for the big datasets!), because it only touches the full data once. I have [uploaded a Jupyter notebook](https://gist.github.com/kathoef/2fbdfd19f29a03aed561e0f5f56d445a) that shows the equivalence of the `xarray.corr()` function and my implementation (using an "unaligned data with nan values"-example, which is what `xarray.corr()` covers) and an example based on Dask arrays, which demonstrates the performance problems that I have stated above, and also that the `xarray.corr()` function is not fully lazy. (Which I assume is actually not very desirable?)

At the moment, I think, in terms of improving big data performance, a considerable improvement could be achieved by removing the `if not valid_values.all()` clause [here](https://github.com/pydata/xarray/blob/v0.16.2/xarray/core/computation.py#L1313), because that seems to determine that a call of `xarray.corr()` is not fully lazy and causes the first (of several?) full touches of the datasets? I haven't checked what's going on afterwards, but maybe that is already a useful starting point? :thinking:

Discussion:
max-sixty: Hi @AndrewWilliams3142 , thanks for another PR!

This looks good. Could we add a test like https://github.com/pydata/xarray/pull/4559/files#diff-74d2dc289aa601b2de094fb3a3b687fd65963401b51b95cc5e0afcd06cc4cb82R45? And maybe reference this or #4559 as an explanation of what's going on.
AndrewWilliams3142: Hi @dcherian , just thinking about your suggestion for using `map_blocks` on the actual `valid_values` check. I've tested this and was wondering if you could maybe point to where I'm going wrong? It does mask out some of the values in a lazy way, but not the *correct* ones.

```python3
da_a = xr.DataArray(
        np.array([[1, 2, 3, 4], [1, 0.1, 0.2, 0.3], [2, 3.2, 0.6, 1.8]]),
        dims=("space", "time"),
        coords=[
            ("space", ["IA", "IL", "IN"]),
            ("time", pd.date_range("2000-01-01", freq="1D", periods=4)),
        ],
    ).chunk({'time':1})

da_b = xr.DataArray(
        np.array([[0.2, 0.4, 0.6, 2], [15, 10, 5, 1], [1, 3.2, np.nan, 1.8]]),
        dims=("space", "time"),
        coords=[
            ("space", ["IA", "IL", "IN"]),
            ("time", pd.date_range("2000-01-01", freq="1D", periods=4)),
        ],
    ).chunk({'time':1})

print(da_a)
>>> <xarray.DataArray (space: 3, time: 4)>
array([[1. , 2. , 3. , 4. ],
       [1. , 0.1, 0.2, 0.3],
       [2. , 3.2, 0.6, 1.8]])
Coordinates:
  * space    (space) <U2 'IA' 'IL' 'IN'
  * time     (time) datetime64[ns] 2000-01-01 2000-01-02 2000-01-03 2000-01-04

print(da_b)
>>> <xarray.DataArray (space: 3, time: 4)>
array([[ 0.2,  0.4,  0.6,  2. ],
       [15. , 10. ,  5. ,  1. ],
       [ 1. ,  3.2,  nan,  1.8]])
Coordinates:
  * space    (space) <U2 'IA' 'IL' 'IN'
  * time     (time) datetime64[ns] 2000-01-01 2000-01-02 2000-01-03 2000-01-04


# Define function to use in map_blocks
def _get_valid_values(da, other):
    da1, da2 = xr.align(da, other, join="inner", copy=False)
    
    # 2. Ignore the nans
    missing_vals = np.logical_or(da1.isnull(), da2.isnull())
    
    if missing_vals.any():
        da = da.where(missing_vals)
        return da
    else:
        return da


# test
outp = da_a.map_blocks(_get_valid_values, args=[da_b])


print(outp.compute())
>>> <xarray.DataArray (space: 3, time: 4)>
array([[1. , 2. , nan, 4. ],
       [1. , 0.1, nan, 0.3],
       [2. , 3.2, 0.6, 1.8]])
Coordinates:
  * time     (time) datetime64[ns] 2000-01-01 2000-01-02 2000-01-03 2000-01-04
  * space    (space) object 'IA' 'IL' 'IN'
```
keewis: #4559 and #4668 fixed a similar issue using `map_blocks`, maybe you can use that as a reference?
dcherian: Well that was confusing! `if missing_vals.any():` will not be triggered if all the values in a block are valid. With `.chunk({"time": 1})`, some of the blocks are all valid.

I discovered this with `print` debugging and the single-threaded scheduler which loops over blocks in a for-loop
```
# Define function to use in map_blocks
def _get_valid_values(da, other):
    da1, da2 = xr.align(da, other, join="inner", copy=False)
    
    # 2. Ignore the nans
    missing_vals = np.logical_or(da1.isnull(), da2.isnull())
    print(missing_vals)
    
    if missing_vals.any():
        da = da.where(missing_vals)
    return da

da_a.map_blocks(_get_valid_values, args=[da_b]).compute(scheduler="sync")
```

For better performance, we should try a `dask.array.map_blocks` approach with `duck_array_ops.isnull` 
AndrewWilliams3142: Hey both, I've added a test to check that dask doesn't compute when calling either `xr.corr()` or `xr.cov()`, and also that the end result is still a `dask ` array. Let me know if there's anything I've missed though! thanks for the help :)

@dcherian, regarding the `apply_ufunc` approach, I might leave that for now but as you said it can always be a future PR
max-sixty: Please feel free to add a whatsnew @AndrewWilliams3142 
max-sixty: Thank you very much @AndrewWilliams3142 !

Patch:
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index 9ba9a7e9..b957ed4b 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -21,11 +21,14 @@ v0.18.3 (unreleased)
 
 New Features
 ~~~~~~~~~~~~
-
 - Allow assigning values to a subset of a dataset using positional or label-based
   indexing (:issue:`3015`, :pull:`5362`). By `Matthias Göbel <https://github.com/matzegoebel>`_.
 - Attempting to reduce a weighted object over missing dimensions now raises an error (:pull:`5362`).
   By `Mattia Almansi <https://github.com/malmans2>`_.
+- :py:func:`xarray.cov` and :py:func:`xarray.corr` now lazily check for missing
+  values if inputs are dask arrays (:issue:`4804`, :pull:`5284`).
+  By `Andrew Williams <https://github.com/AndrewWilliams3142>`_.
+
 
 Breaking changes
 ~~~~~~~~~~~~~~~~
diff --git a/xarray/core/computation.py b/xarray/core/computation.py
index 6010d502..07b7bca2 100644
--- a/xarray/core/computation.py
+++ b/xarray/core/computation.py
@@ -1352,12 +1352,23 @@ def _cov_corr(da_a, da_b, dim=None, ddof=0, method=None):
 
     # 2. Ignore the nans
     valid_values = da_a.notnull() & da_b.notnull()
+    valid_count = valid_values.sum(dim) - ddof
 
-    if not valid_values.all():
-        da_a = da_a.where(valid_values)
-        da_b = da_b.where(valid_values)
+    def _get_valid_values(da, other):
+        """
+        Function to lazily mask da_a and da_b
+        following a similar approach to
+        https://github.com/pydata/xarray/pull/4559
+        """
+        missing_vals = np.logical_or(da.isnull(), other.isnull())
+        if missing_vals.any():
+            da = da.where(~missing_vals)
+            return da
+        else:
+            return da
 
-    valid_count = valid_values.sum(dim) - ddof
+    da_a = da_a.map_blocks(_get_valid_values, args=[da_b])
+    da_b = da_b.map_blocks(_get_valid_values, args=[da_a])
 
     # 3. Detrend along the given dim
     demeaned_da_a = da_a - da_a.mean(dim=dim)


Test patch:
diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py
index cbfa61a4..b7ae1ca9 100644
--- a/xarray/tests/test_computation.py
+++ b/xarray/tests/test_computation.py
@@ -22,7 +22,7 @@ from xarray.core.computation import (
     unified_dim_sizes,
 )
 
-from . import has_dask, requires_dask
+from . import has_dask, raise_if_dask_computes, requires_dask
 
 dask = pytest.importorskip("dask")
 
@@ -1386,6 +1386,7 @@ def arrays_w_tuples():
         da.isel(time=range(2, 20)).rolling(time=3, center=True).mean(),
         xr.DataArray([[1, 2], [1, np.nan]], dims=["x", "time"]),
         xr.DataArray([[1, 2], [np.nan, np.nan]], dims=["x", "time"]),
+        xr.DataArray([[1, 2], [2, 1]], dims=["x", "time"]),
     ]
 
     array_tuples = [
@@ -1394,12 +1395,40 @@ def arrays_w_tuples():
         (arrays[1], arrays[1]),
         (arrays[2], arrays[2]),
         (arrays[2], arrays[3]),
+        (arrays[2], arrays[4]),
+        (arrays[4], arrays[2]),
         (arrays[3], arrays[3]),
+        (arrays[4], arrays[4]),
     ]
 
     return arrays, array_tuples
 
 
+@pytest.mark.parametrize("ddof", [0, 1])
+@pytest.mark.parametrize(
+    "da_a, da_b",
+    [
+        arrays_w_tuples()[1][3],
+        arrays_w_tuples()[1][4],
+        arrays_w_tuples()[1][5],
+        arrays_w_tuples()[1][6],
+        arrays_w_tuples()[1][7],
+        arrays_w_tuples()[1][8],
+    ],
+)
+@pytest.mark.parametrize("dim", [None, "x", "time"])
+def test_lazy_corrcov(da_a, da_b, dim, ddof):
+    # GH 5284
+    from dask import is_dask_collection
+
+    with raise_if_dask_computes():
+        cov = xr.cov(da_a.chunk(), da_b.chunk(), dim=dim, ddof=ddof)
+        assert is_dask_collection(cov)
+
+        corr = xr.corr(da_a.chunk(), da_b.chunk(), dim=dim)
+        assert is_dask_collection(corr)
+
+
 @pytest.mark.parametrize("ddof", [0, 1])
 @pytest.mark.parametrize(
     "da_a, da_b",