Simplify missing value handling in xarray.corr (#6025)

Author: shoyer

xarray/core/computation.py

@@ -1346,25 +1346,10 @@ 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()
+    da_a = da_a.where(valid_values)
+    da_b = da_b.where(valid_values)
     valid_count = valid_values.sum(dim) - ddof
 
-    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:
-            # ensure consistent return dtype
-            return da.astype(float)
-
-    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)
     demeaned_da_b = da_b - da_b.mean(dim=dim)