add dask and a few other methods

pydata · dcherian · Nov 26, 2024 · Nov 18, 2024 · Nov 18, 2024 · Nov 18, 2024
commit f6074d2fa3b9c2d3900cda25c4cb322f2c698bd1
diff --git a/xarray/tests/test_duck_array_wrapping.py b/xarray/tests/test_duck_array_wrapping.py
@@ -6,22 +6,35 @@
 
 # Don't run cupy in CI because it requires a GPU
 NAMESPACE_ARRAYS = {
-    "jax.numpy": {
+    "cupy": {
         "attrs": {
             "array": "ndarray",
             "constructor": "asarray",
         },
+        "xfails": {"quantile": "no nanquantile"},
+    },
+    "dask.array": {
+        "attrs": {
+            "array": "Array",
+            "constructor": "from_array",
+        },
         "xfails": {
-            "rolling": "no sliding_window_view",
-            "rolling_mean": "no sliding_window_view",
+            "argsort": "no argsort",
+            "conjugate": "conj but no conjugate",
+            "searchsorted": "dask.array.searchsorted but no Array.searchsorted",
         },
     },
-    "cupy": {
+    "jax.numpy": {
         "attrs": {
             "array": "ndarray",
             "constructor": "asarray",
         },
-        "xfails": {"quantile": "no nanquantile"},
+        "xfails": {
+            "rolling_construct": "no sliding_window_view",
+            "rolling_reduce": "no sliding_window_view",
+            "cumulative_construct": "no sliding_window_view",
+            "cumulative_reduce": "no sliding_window_view",
+        },
     },
     "pint": {
         "attrs": {
@@ -42,7 +55,8 @@
             "isin": "returns a bool",
             "isnull": "returns a bool",
             "notnull": "returns a bool",
-            "rolling_mean": "no dispatch for numbagg/bottleneck",
+            "rolling_reduce": "no dispatch for numbagg/bottleneck",
+            "cumulative_reduce": "no dispatch for numbagg/bottleneck",
             "searchsorted": "returns an int",
             "weighted": "no tensordot",
         },
@@ -59,8 +73,12 @@
             "count": "dense output",
             "dot": "fails on some platforms/versions",
             "isin": "no isin",
-            "rolling": "no sliding_window_view",
-            "rolling_mean": "no sliding_window_view",
+            "rolling_construct": "no sliding_window_view",
+            "rolling_reduce": "no sliding_window_view",
+            "cumulative_construct": "no sliding_window_view",
+            "cumulative_reduce": "no sliding_window_view",
+            "coarsen_construct": "pad constant_values must be fill_value",
+            "coarsen_reduce": "pad constant_values must be fill_value",
             "weighted": "fill_value error",
             "coarsen": "pad constant_values must be fill_value",
             "quantile": "no non skipping version",
@@ -119,7 +137,7 @@ def setUp(self, request, namespace):
 
     def test_apply_ufunc(self):
         func = lambda x: x + 1
-        result = xr.apply_ufunc(func, self.x1)
+        result = xr.apply_ufunc(func, self.x1, dask="parallelized")
         assert isinstance(result.data, self.Array)
 
     def test_align(self):
@@ -248,26 +266,51 @@ def test_groupby(self):
         result = self.x.groupby("x").mean()
         assert isinstance(result.data, self.Array)
 
-    def test_rolling(self):
+    def test_groupby_bins(self):
+        result = self.x.groupby_bins("x", bins=[0, 2, 4, 6]).mean()
+        assert isinstance(result.data, self.Array)
+
+    def test_rolling_iter(self):
         result = self.x.rolling(x=3)
         elem = next(iter(result))[1]
         assert isinstance(elem.data, self.Array)
 
+    def test_rolling_construct(self):
+        result = self.x.rolling(x=3).construct(x="window")
+        assert isinstance(result.data, self.Array)
+
     @pytest.mark.parametrize("skipna", [True, False])
-    def test_rolling_mean(self, skipna):
+    def test_rolling_reduce(self, skipna):
         result = self.x.rolling(x=3).mean(skipna=skipna)
         assert isinstance(result.data, self.Array)
 
     @pytest.mark.xfail(reason="rolling_exp uses numbagg")
-    def test_rolling_exp(self):
+    def test_rolling_exp_reduce(self):
         result = self.x.rolling_exp(x=3).mean()
         assert isinstance(result.data, self.Array)
 
+    def test_cumulative_iter(self):
+        result = self.x.cumulative("x")
+        elem = next(iter(result))[1]
+        assert isinstance(elem.data, self.Array)
+
+    def test_cumulative_construct(self):
+        result = self.x.cumulative("x").construct(x="window")
+        assert isinstance(result.data, self.Array)
+
+    def test_cumulative_reduce(self):
+        result = self.x.cumulative("x").sum()
+        assert isinstance(result.data, self.Array)
+
     def test_weighted(self):
         result = self.x.weighted(self.x.fillna(0)).mean()
         assert isinstance(result.data, self.Array)
 
-    def test_coarsen(self):
+    def test_coarsen_construct(self):
+        result = self.x.coarsen(x=2, boundary="pad").construct(x=["a", "b"])
+        assert isinstance(result.data, self.Array)
+
+    def test_coarsen_reduce(self):
         result = self.x.coarsen(x=2, boundary="pad").mean()
         assert isinstance(result.data, self.Array)
 
@@ -391,6 +434,10 @@ def test_argsort(self):
         result = self.x.argsort()
         assert isinstance(result.data, self.Array)
 
+    def test_astype(self):
+        result = self.x.astype(int)
+        assert isinstance(result.data, self.Array)
+
     def test_clip(self):
         result = self.x.clip(min=2.0, max=4.0)
         assert isinstance(result.data, self.Array)