Repository: pydata/xarray

Issue:
Implement `skipna` in xr.quantile for speedup
`xr.quantile` uses `np.nanquantile` which is slower than `np.quantile` but only needed when ignoring nans is needed. Adding `skipna` as kwarg would lead to a speedup for many use-cases.


#### MCVE Code Sample

`np.quantile` is much faster than `np.nanquantile`
```python
control = xr.DataArray(np.random.random((50,256,192)),dims=['time','x','y'])
%time _ = control.quantile(dim='time',q=q)
CPU times: user 4.14 s, sys: 61.4 ms, total: 4.2 s
Wall time: 4.3 s


%time _ = np.quantile(control,q,axis=0)
CPU times: user 47.1 ms, sys: 4.27 ms, total: 51.4 ms
Wall time: 52.6 ms

%time _ = np.nanquantile(control,q,axis=0)
CPU times: user 3.18 s, sys: 21.4 ms, total: 3.2 s
Wall time: 3.22 s
```

#### Expected Output
faster xr.quantile:
```
%time _ = control.quantile(dim='time',q=q)
CPU times: user 4.95 s, sys: 34.3 ms, total: 4.98 s
Wall time: 5.88 s

%time _ = control.quantile(dim='time',q=q, skipna=False)
CPU times: user 85.3 ms, sys: 16.7 ms, total: 102 ms
Wall time: 127 ms

```

#### Problem Description

np.nanquantile not always needed


#### Versions

<details><summary>Output of `xr.show_versions()`</summary>
xr=0.15.1

</details>


Discussion:
pep8speaks: Hello @aaronspring! Thanks for opening this PR. We checked the lines you've touched for [PEP 8](https://www.python.org/dev/peps/pep-0008) issues, and found:



* In the file [`xarray/core/dataset.py`](https://github.com/pydata/xarray/blob/dfd85f2b1ad98da412debd0eb84d66e7f7304206/xarray/core/dataset.py):

> [Line 5143:96](https://github.com/pydata/xarray/blob/dfd85f2b1ad98da412debd0eb84d66e7f7304206/xarray/core/dataset.py#L5143): [F821](https://duckduckgo.com/?q=pep8%20F821) undefined name 'skipna'




aaronspring: I am unsure whether I added enough tests. only one test that checks for different values with or without skipna.
max-sixty: This looks great @aaronspring ! Thanks!
 

Patch:
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index 24120270..2c30db99 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -45,7 +45,12 @@ New Features
   By `Julia Signell <https://github.com/jsignell>`_.
 - :py:meth:`Dataset.where` and :py:meth:`DataArray.where` accept a lambda as a
   first argument, which is then called on the input; replicating pandas' behavior.
-  By `Maximilian Roos <https://github.com/max-sixty>`_
+  By `Maximilian Roos <https://github.com/max-sixty>`_.
+- Implement ``skipna`` in :py:meth:`Dataset.quantile`, :py:meth:`DataArray.quantile`,
+  :py:meth:`core.groupby.DatasetGroupBy.quantile`, :py:meth:`core.groupby.DataArrayGroupBy.quantile`
+  (:issue:`3843`, :pull:`3844`) 
+  By `Aaron Spring <https://github.com/aaronspring>`_.
+
 
 Bug fixes
 ~~~~~~~~~
diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py
index 4e80ef22..7fcb42bf 100644
--- a/xarray/core/dataarray.py
+++ b/xarray/core/dataarray.py
@@ -2939,6 +2939,7 @@ class DataArray(AbstractArray, DataWithCoords):
         dim: Union[Hashable, Sequence[Hashable], None] = None,
         interpolation: str = "linear",
         keep_attrs: bool = None,
+        skipna: bool = True,
     ) -> "DataArray":
         """Compute the qth quantile of the data along the specified dimension.
 
@@ -2966,6 +2967,8 @@ class DataArray(AbstractArray, DataWithCoords):
             If True, the dataset's attributes (`attrs`) will be copied from
             the original object to the new one.  If False (default), the new
             object will be returned without attributes.
+        skipna : bool, optional
+            Whether to skip missing values when aggregating.
 
         Returns
         -------
@@ -2978,7 +2981,7 @@ class DataArray(AbstractArray, DataWithCoords):
 
         See Also
         --------
-        numpy.nanquantile, pandas.Series.quantile, Dataset.quantile
+        numpy.nanquantile, numpy.quantile, pandas.Series.quantile, Dataset.quantile
 
         Examples
         --------
@@ -3015,7 +3018,11 @@ class DataArray(AbstractArray, DataWithCoords):
         """
 
         ds = self._to_temp_dataset().quantile(
-            q, dim=dim, keep_attrs=keep_attrs, interpolation=interpolation
+            q,
+            dim=dim,
+            keep_attrs=keep_attrs,
+            interpolation=interpolation,
+            skipna=skipna,
         )
         return self._from_temp_dataset(ds)
 
diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py
index 52940e98..f286236d 100644
--- a/xarray/core/dataset.py
+++ b/xarray/core/dataset.py
@@ -5140,7 +5140,13 @@ class Dataset(Mapping, ImplementsDatasetReduce, DataWithCoords):
         return aligned_self.isel(**indices)
 
     def quantile(
-        self, q, dim=None, interpolation="linear", numeric_only=False, keep_attrs=None
+        self,
+        q,
+        dim=None,
+        interpolation="linear",
+        numeric_only=False,
+        keep_attrs=None,
+        skipna=True,
     ):
         """Compute the qth quantile of the data along the specified dimension.
 
@@ -5171,6 +5177,8 @@ class Dataset(Mapping, ImplementsDatasetReduce, DataWithCoords):
             object will be returned without attributes.
         numeric_only : bool, optional
             If True, only apply ``func`` to variables with a numeric dtype.
+        skipna : bool, optional
+            Whether to skip missing values when aggregating.
 
         Returns
         -------
@@ -5183,7 +5191,7 @@ class Dataset(Mapping, ImplementsDatasetReduce, DataWithCoords):
 
         See Also
         --------
-        numpy.nanquantile, pandas.Series.quantile, DataArray.quantile
+        numpy.nanquantile, numpy.quantile, pandas.Series.quantile, DataArray.quantile
 
         Examples
         --------
@@ -5258,6 +5266,7 @@ class Dataset(Mapping, ImplementsDatasetReduce, DataWithCoords):
                             dim=reduce_dims,
                             interpolation=interpolation,
                             keep_attrs=keep_attrs,
+                            skipna=skipna,
                         )
 
             else:
diff --git a/xarray/core/groupby.py b/xarray/core/groupby.py
index f2a9ebac..4223d9dc 100644
--- a/xarray/core/groupby.py
+++ b/xarray/core/groupby.py
@@ -558,7 +558,9 @@ class GroupBy(SupportsArithmetic):
         out = ops.fillna(self, value)
         return out
 
-    def quantile(self, q, dim=None, interpolation="linear", keep_attrs=None):
+    def quantile(
+        self, q, dim=None, interpolation="linear", keep_attrs=None, skipna=True
+    ):
         """Compute the qth quantile over each array in the groups and
         concatenate them together into a new array.
 
@@ -582,6 +584,8 @@ class GroupBy(SupportsArithmetic):
                 * higher: ``j``.
                 * nearest: ``i`` or ``j``, whichever is nearest.
                 * midpoint: ``(i + j) / 2``.
+        skipna : bool, optional
+            Whether to skip missing values when aggregating.
 
         Returns
         -------
@@ -595,7 +599,7 @@ class GroupBy(SupportsArithmetic):
 
         See Also
         --------
-        numpy.nanquantile, pandas.Series.quantile, Dataset.quantile,
+        numpy.nanquantile, numpy.quantile, pandas.Series.quantile, Dataset.quantile,
         DataArray.quantile
 
         Examples
@@ -656,6 +660,7 @@ class GroupBy(SupportsArithmetic):
             dim=dim,
             interpolation=interpolation,
             keep_attrs=keep_attrs,
+            skipna=skipna,
         )
 
         return out
diff --git a/xarray/core/variable.py b/xarray/core/variable.py
index 62f9fde6..435edb6f 100644
--- a/xarray/core/variable.py
+++ b/xarray/core/variable.py
@@ -1678,7 +1678,9 @@ class Variable(
         """
         return self.broadcast_equals(other, equiv=equiv)
 
-    def quantile(self, q, dim=None, interpolation="linear", keep_attrs=None):
+    def quantile(
+        self, q, dim=None, interpolation="linear", keep_attrs=None, skipna=True
+    ):
         """Compute the qth quantile of the data along the specified dimension.
 
         Returns the qth quantiles(s) of the array elements.
@@ -1725,6 +1727,8 @@ class Variable(
 
         from .computation import apply_ufunc
 
+        _quantile_func = np.nanquantile if skipna else np.quantile
+
         if keep_attrs is None:
             keep_attrs = _get_keep_attrs(default=False)
 
@@ -1739,7 +1743,7 @@ class Variable(
 
         def _wrapper(npa, **kwargs):
             # move quantile axis to end. required for apply_ufunc
-            return np.moveaxis(np.nanquantile(npa, **kwargs), 0, -1)
+            return np.moveaxis(_quantile_func(npa, **kwargs), 0, -1)
 
         axis = np.arange(-1, -1 * len(dim) - 1, -1)
         result = apply_ufunc(


Test patch:
diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py
index b8a9c5ed..33f1b403 100644
--- a/xarray/tests/test_dataarray.py
+++ b/xarray/tests/test_dataarray.py
@@ -2368,13 +2368,15 @@ class TestDataArray:
         with pytest.raises(TypeError):
             orig.mean(out=np.ones(orig.shape))
 
+    @pytest.mark.parametrize("skipna", [True, False])
     @pytest.mark.parametrize("q", [0.25, [0.50], [0.25, 0.75]])
     @pytest.mark.parametrize(
         "axis, dim", zip([None, 0, [0], [0, 1]], [None, "x", ["x"], ["x", "y"]])
     )
-    def test_quantile(self, q, axis, dim):
-        actual = DataArray(self.va).quantile(q, dim=dim, keep_attrs=True)
-        expected = np.nanpercentile(self.dv.values, np.array(q) * 100, axis=axis)
+    def test_quantile(self, q, axis, dim, skipna):
+        actual = DataArray(self.va).quantile(q, dim=dim, keep_attrs=True, skipna=skipna)
+        _percentile_func = np.nanpercentile if skipna else np.percentile
+        expected = _percentile_func(self.dv.values, np.array(q) * 100, axis=axis)
         np.testing.assert_allclose(actual.values, expected)
         if is_scalar(q):
             assert "quantile" not in actual.dims
diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py
index 44ffafb2..d2e8c6b7 100644
--- a/xarray/tests/test_dataset.py
+++ b/xarray/tests/test_dataset.py
@@ -4697,12 +4697,13 @@ class TestDataset:
         )
         assert_identical(expected, actual)
 
+    @pytest.mark.parametrize("skipna", [True, False])
     @pytest.mark.parametrize("q", [0.25, [0.50], [0.25, 0.75]])
-    def test_quantile(self, q):
+    def test_quantile(self, q, skipna):
         ds = create_test_data(seed=123)
 
         for dim in [None, "dim1", ["dim1"]]:
-            ds_quantile = ds.quantile(q, dim=dim)
+            ds_quantile = ds.quantile(q, dim=dim, skipna=skipna)
             if is_scalar(q):
                 assert "quantile" not in ds_quantile.dims
             else:
@@ -4710,12 +4711,27 @@ class TestDataset:
 
             for var, dar in ds.data_vars.items():
                 assert var in ds_quantile
-                assert_identical(ds_quantile[var], dar.quantile(q, dim=dim))
+                assert_identical(
+                    ds_quantile[var], dar.quantile(q, dim=dim, skipna=skipna)
+                )
         dim = ["dim1", "dim2"]
-        ds_quantile = ds.quantile(q, dim=dim)
+        ds_quantile = ds.quantile(q, dim=dim, skipna=skipna)
         assert "dim3" in ds_quantile.dims
         assert all(d not in ds_quantile.dims for d in dim)
 
+    @pytest.mark.parametrize("skipna", [True, False])
+    def test_quantile_skipna(self, skipna):
+        q = 0.1
+        dim = "time"
+        ds = Dataset({"a": ([dim], np.arange(0, 11))})
+        ds = ds.where(ds >= 1)
+
+        result = ds.quantile(q=q, dim=dim, skipna=skipna)
+
+        value = 1.9 if skipna else np.nan
+        expected = Dataset({"a": value}, coords={"quantile": q})
+        assert_identical(result, expected)
+
     @requires_bottleneck
     def test_rank(self):
         ds = create_test_data(seed=1234)
diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py
index c86ecd01..c600f7a7 100644
--- a/xarray/tests/test_variable.py
+++ b/xarray/tests/test_variable.py
@@ -1511,14 +1511,16 @@ class TestVariable(VariableSubclassobjects):
         with pytest.warns(DeprecationWarning, match="allow_lazy is deprecated"):
             v.mean(dim="x", allow_lazy=False)
 
+    @pytest.mark.parametrize("skipna", [True, False])
     @pytest.mark.parametrize("q", [0.25, [0.50], [0.25, 0.75]])
     @pytest.mark.parametrize(
         "axis, dim", zip([None, 0, [0], [0, 1]], [None, "x", ["x"], ["x", "y"]])
     )
-    def test_quantile(self, q, axis, dim):
+    def test_quantile(self, q, axis, dim, skipna):
         v = Variable(["x", "y"], self.d)
-        actual = v.quantile(q, dim=dim)
-        expected = np.nanpercentile(self.d, np.array(q) * 100, axis=axis)
+        actual = v.quantile(q, dim=dim, skipna=skipna)
+        _percentile_func = np.nanpercentile if skipna else np.percentile
+        expected = _percentile_func(self.d, np.array(q) * 100, axis=axis)
         np.testing.assert_allclose(actual.values, expected)
 
     @requires_dask