Repository: pydata/xarray

Issue:
Push method with limit bigger than the chunksize of Dask array is not correctly working
### What happened?

After upgrading to Xarray 2024.11.0, a colleague from work reported to me that the forward fill method was not properly working when the limit parameter was used, causing some values to get NaN instead of the previous value.

I investigated what was causing the issue, and I found that the problem is in the method used as binop on the cumreduction for calculating the valid positions to push, it is not correctly restarting the counter because the cumreduction calls twice this function, one for adding the last value of the previous chunk to the actual one, and the second time for combining the last value of two chunks, this last call it's not working as expected because it is not able to detect if there was a restart or not of the accumulation, making that the number of nans accumulates indefinitely.

I have not found a way to correctly detect if there was a restart on the counter or not when the cumreduction calls the binop again, mainly because the function does not receive any parameter indicating the shape of the chunk, so it is hard to determine if the counter suffered a restart or not just seeing it's last accumulated value, for that reason, I propose to replace the direct use of the cumreduction (only for the valid positions) with some Dask function as shown on the attached images, unfortunately, this generates more tasks than before.

Before:
![image](https://github.com/user-attachments/assets/0d3e7b77-9085-4339-b2b7-8f1b266b0149)
After:
![image](https://github.com/user-attachments/assets/712ce82d-9cd6-4fc6-8d0e-89df81a1cb47)


### What did you expect to happen?

I expect that the results of the push method will match the in-memory version as shown in the minimal complete verifiable example.

### Minimal Complete Verifiable Example

```Python
from xarray.core.dask_array_ops import push
from bottleneck import push as push_bottleneck

v = np.array([np.nan, np.nan, np.nan, 2, np.nan, np.nan, np.nan, 9, np.nan, np.nan, np.nan])

arr = da.from_array(v, chunks=3)

r1 = push(arr, 4, axis=0, method="sequential")
r2 = push_bottleneck(v, 4, axis=0)

print(r1.compute())
print(r2)
assert np.array_equal(r1.compute(), r2, equal_nan=True)
```


### MVCE confirmation

- [X] Minimal example — the example is as focused as reasonably possible to demonstrate the underlying issue in xarray.
- [X] Complete example — the example is self-contained, including all data and the text of any traceback.
- [X] Verifiable example — the example copy & pastes into an IPython prompt or [Binder notebook](https://mybinder.org/v2/gh/pydata/xarray/main?urlpath=lab/tree/doc/examples/blank_template.ipynb), returning the result.
- [X] New issue — a search of GitHub Issues suggests this is not a duplicate.
- [X] Recent environment — the issue occurs with the latest version of xarray and its dependencies.

### Relevant log output

```Python
Output of the minimal complete verifiable example:

dask      [nan nan nan  2.  2.  2. nan  9.  9. nan nan]
in-memory [nan nan nan  2.  2.  2.  2.  9.  9.  9.  9.]

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Cell In[34], line 13
     11 print(r1.compute())
     12 print(r2)
---> 13 assert np.array_equal(r1.compute(), r2, equal_nan=True)

AssertionError:
```


### Anything else we need to know?

_No response_

### Environment

<details>

INSTALLED VERSIONS
------------------
commit: 5279bd15708e21ea582e56e4f27fc3e9663a8509
python: 3.11.10 | packaged by conda-forge | (main, Sep 10 2024, 10:53:25) [MSC v.1940 64 bit (AMD64)]
python-bits: 64
OS: Windows
OS-release: 10
machine: AMD64
processor: Intel64 Family 6 Model 165 Stepping 2, GenuineIntel
byteorder: little
LC_ALL: None
LANG: None
LOCALE: ('Spanish_Venezuela', '1252')
libhdf5: 1.14.3
libnetcdf: 4.9.2

xarray: 2025.1.2.dev1+g5279bd15.d20250111
pandas: 2.2.2
numpy: 2.0.2
scipy: 1.14.1
netCDF4: 1.7.1.post2
pydap: None
h5netcdf: None
h5py: None
zarr: 2.18.3
cftime: 1.6.4
nc_time_axis: None
iris: None
bottleneck: 1.4.2
dask: 2024.9.0
distributed: 2024.9.0
matplotlib: None
cartopy: None
seaborn: None
numbagg: 0.8.2
fsspec: 2024.9.0
cupy: None
pint: None
sparse: None
flox: None
numpy_groupies: None
setuptools: 73.0.1
pip: 24.2
conda: None
pytest: 8.3.3
mypy: 1.11.2
IPython: 8.27.0
sphinx: None


</details>


Discussion:

Patch:
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index 3810c806..76ed4c43 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -95,6 +95,9 @@ Deprecations
 
 Bug fixes
 ~~~~~~~~~
+
+- Fix :py:meth:`DataArray.ffill`, :py:meth:`DataArray.bfill`, :py:meth:`Dataset.ffill` and :py:meth:`Dataset.bfill` when the limit is bigger than the chunksize (:issue:`9939`).
+  By `Joseph Nowak <https://github.com/josephnowak>`_.
 - Fix issues related to Pandas v3 ("us" vs. "ns" for python datetime, copy on write) and handling of 0d-numpy arrays in datetime/timedelta decoding (:pull:`9953`).
   By `Kai Mühlbauer <https://github.com/kmuehlbauer>`_.
 - Remove dask-expr from CI runs, add "pyarrow" dask dependency to windows CI runs, fix related tests (:issue:`9962`, :pull:`9971`).
diff --git a/xarray/core/dask_array_ops.py b/xarray/core/dask_array_ops.py
index 2dca3853..aff67667 100644
--- a/xarray/core/dask_array_ops.py
+++ b/xarray/core/dask_array_ops.py
@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import math
-from functools import partial
 
 from xarray.core import dtypes, nputils
 
@@ -92,31 +91,6 @@ def _dtype_push(a, axis, dtype=None):
     return _push(a, axis=axis)
 
 
-def _reset_cumsum(a, axis, dtype=None):
-    import numpy as np
-
-    cumsum = np.cumsum(a, axis=axis)
-    reset_points = np.maximum.accumulate(np.where(a == 0, cumsum, 0), axis=axis)
-    return cumsum - reset_points
-
-
-def _last_reset_cumsum(a, axis, keepdims=None):
-    import numpy as np
-
-    # Take the last cumulative sum taking into account the reset
-    # This is useful for blelloch method
-    return np.take(_reset_cumsum(a, axis=axis), axis=axis, indices=[-1])
-
-
-def _combine_reset_cumsum(a, b, axis):
-    import numpy as np
-
-    # It is going to sum the previous result until the first
-    # non nan value
-    bitmask = np.cumprod(b != 0, axis=axis)
-    return np.where(bitmask, b + a, b)
-
-
 def push(array, n, axis, method="blelloch"):
     """
     Dask-aware bottleneck.push
@@ -145,16 +119,18 @@ def push(array, n, axis, method="blelloch"):
     )
 
     if n is not None and 0 < n < array.shape[axis] - 1:
-        valid_positions = da.reductions.cumreduction(
-            func=_reset_cumsum,
-            binop=partial(_combine_reset_cumsum, axis=axis),
-            ident=0,
-            x=da.isnan(array, dtype=int),
-            axis=axis,
-            dtype=int,
-            method=method,
-            preop=_last_reset_cumsum,
-        )
-        pushed_array = da.where(valid_positions <= n, pushed_array, np.nan)
+        # The idea is to calculate a cumulative sum of a bitmask
+        # created from the isnan method, but every time a False is found the sum
+        # must be restarted, and the final result indicates the amount of contiguous
+        # nan values found in the original array on every position
+        nan_bitmask = da.isnan(array, dtype=int)
+        cumsum_nan = nan_bitmask.cumsum(axis=axis, method=method)
+        valid_positions = da.where(nan_bitmask == 0, cumsum_nan, np.nan)
+        valid_positions = push(valid_positions, None, axis, method=method)
+        # All the NaNs at the beginning are converted to 0
+        valid_positions = da.nan_to_num(valid_positions)
+        valid_positions = cumsum_nan - valid_positions
+        valid_positions = valid_positions <= n
+        pushed_array = da.where(valid_positions, pushed_array, np.nan)
 
     return pushed_array


Test patch:
diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py
index bc6e16ff..0237252f 100644
--- a/xarray/tests/test_backends.py
+++ b/xarray/tests/test_backends.py
@@ -6134,453 +6134,424 @@ def test_zarr_closing_internal_zip_store():
 @requires_zarr
 @pytest.mark.usefixtures("default_zarr_format")
 class TestZarrRegionAuto:
-    def test_zarr_region_auto_all(self, tmp_path):
-        x = np.arange(0, 50, 10)
-        y = np.arange(0, 20, 2)
-        data = np.ones((5, 10))
-        ds = xr.Dataset(
-            {
-                "test": xr.DataArray(
-                    data,
-                    dims=("x", "y"),
-                    coords={"x": x, "y": y},
-                )
-            }
-        )
-        ds.to_zarr(tmp_path / "test.zarr")
-
-        ds_region = 1 + ds.isel(x=slice(2, 4), y=slice(6, 8))
-        ds_region.to_zarr(tmp_path / "test.zarr", region="auto")
-
-        ds_updated = xr.open_zarr(tmp_path / "test.zarr")
+    """These are separated out since we should not need to test this logic with every store."""
 
-        expected = ds.copy()
-        expected["test"][2:4, 6:8] += 1
-        assert_identical(ds_updated, expected)
+    @contextlib.contextmanager
+    def create_zarr_target(self):
+        with create_tmp_file(suffix=".zarr") as tmp:
+            yield tmp
 
-    def test_zarr_region_auto_mixed(self, tmp_path):
+    @contextlib.contextmanager
+    def create(self):
         x = np.arange(0, 50, 10)
         y = np.arange(0, 20, 2)
         data = np.ones((5, 10))
         ds = xr.Dataset(
-            {
-                "test": xr.DataArray(
-                    data,
-                    dims=("x", "y"),
-                    coords={"x": x, "y": y},
-                )
-            }
+            {"test": xr.DataArray(data, dims=("x", "y"), coords={"x": x, "y": y})}
         )
-        ds.to_zarr(tmp_path / "test.zarr")
+        with self.create_zarr_target() as target:
+            self.save(target, ds)
+            yield target, ds
 
-        ds_region = 1 + ds.isel(x=slice(2, 4), y=slice(6, 8))
-        ds_region.to_zarr(
-            tmp_path / "test.zarr", region={"x": "auto", "y": slice(6, 8)}
-        )
+    def save(self, target, ds, **kwargs):
+        ds.to_zarr(target, **kwargs)
 
-        ds_updated = xr.open_zarr(tmp_path / "test.zarr")
+    @pytest.mark.parametrize(
+        "region",
+        [
+            pytest.param("auto", id="full-auto"),
+            pytest.param({"x": "auto", "y": slice(6, 8)}, id="mixed-auto"),
+        ],
+    )
+    def test_zarr_region_auto(self, region):
+        with self.create() as (target, ds):
+            ds_region = 1 + ds.isel(x=slice(2, 4), y=slice(6, 8))
+            self.save(target, ds_region, region=region)
+            ds_updated = xr.open_zarr(target)
+
+            expected = ds.copy()
+            expected["test"][2:4, 6:8] += 1
+            assert_identical(ds_updated, expected)
+
+    def test_zarr_region_auto_noncontiguous(self):
+        with self.create() as (target, ds):
+            with pytest.raises(ValueError):
+                self.save(target, ds.isel(x=[0, 2, 3], y=[5, 6]), region="auto")
 
-        expected = ds.copy()
-        expected["test"][2:4, 6:8] += 1
-        assert_identical(ds_updated, expected)
+            dsnew = ds.copy()
+            dsnew["x"] = dsnew.x + 5
+            with pytest.raises(KeyError):
+                self.save(target, dsnew, region="auto")
 
-    def test_zarr_region_auto_noncontiguous(self, tmp_path):
-        x = np.arange(0, 50, 10)
-        y = np.arange(0, 20, 2)
-        data = np.ones((5, 10))
-        ds = xr.Dataset(
-            {
-                "test": xr.DataArray(
-                    data,
-                    dims=("x", "y"),
-                    coords={"x": x, "y": y},
-                )
-            }
-        )
-        ds.to_zarr(tmp_path / "test.zarr")
+    def test_zarr_region_index_write(self, tmp_path):
+        region: Mapping[str, slice] | Literal["auto"]
+        region_slice = dict(x=slice(2, 4), y=slice(6, 8))
 
-        ds_region = 1 + ds.isel(x=[0, 2, 3], y=[5, 6])
-        with pytest.raises(ValueError):
-            ds_region.to_zarr(tmp_path / "test.zarr", region={"x": "auto", "y": "auto"})
+        with self.create() as (target, ds):
+            ds_region = 1 + ds.isel(region_slice)
+            for region in [region_slice, "auto"]:  # type: ignore[assignment]
+                with patch.object(
+                    ZarrStore,
+                    "set_variables",
+                    side_effect=ZarrStore.set_variables,
+                    autospec=True,
+                ) as mock:
+                    self.save(target, ds_region, region=region, mode="r+")
+
+                    # should write the data vars but never the index vars with auto mode
+                    for call in mock.call_args_list:
+                        written_variables = call.args[1].keys()
+                        assert "test" in written_variables
+                        assert "x" not in written_variables
+                        assert "y" not in written_variables
+
+    def test_zarr_region_append(self):
+        with self.create() as (target, ds):
+            x_new = np.arange(40, 70, 10)
+            data_new = np.ones((3, 10))
+            ds_new = xr.Dataset(
+                {
+                    "test": xr.DataArray(
+                        data_new,
+                        dims=("x", "y"),
+                        coords={"x": x_new, "y": ds.y},
+                    )
+                }
+            )
 
-    def test_zarr_region_auto_new_coord_vals(self, tmp_path):
-        x = np.arange(0, 50, 10)
-        y = np.arange(0, 20, 2)
-        data = np.ones((5, 10))
-        ds = xr.Dataset(
-            {
-                "test": xr.DataArray(
-                    data,
-                    dims=("x", "y"),
-                    coords={"x": x, "y": y},
-                )
-            }
-        )
-        ds.to_zarr(tmp_path / "test.zarr")
+            # Now it is valid to use auto region detection with the append mode,
+            # but it is still unsafe to modify dimensions or metadata using the region
+            # parameter.
+            with pytest.raises(KeyError):
+                self.save(target, ds_new, mode="a", append_dim="x", region="auto")
 
-        x = np.arange(5, 55, 10)
-        y = np.arange(0, 20, 2)
-        data = np.ones((5, 10))
-        ds = xr.Dataset(
-            {
-                "test": xr.DataArray(
-                    data,
-                    dims=("x", "y"),
-                    coords={"x": x, "y": y},
-                )
-            }
-        )
+    def test_zarr_region(self):
+        with self.create() as (target, ds):
+            ds_transposed = ds.transpose("y", "x")
+            ds_region = 1 + ds_transposed.isel(x=[0], y=[0])
+            self.save(target, ds_region, region={"x": slice(0, 1), "y": slice(0, 1)})
 
-        ds_region = 1 + ds.isel(x=slice(2, 4), y=slice(6, 8))
-        with pytest.raises(KeyError):
-            ds_region.to_zarr(tmp_path / "test.zarr", region={"x": "auto", "y": "auto"})
+            # Write without region
+            self.save(target, ds_transposed, mode="r+")
 
-    def test_zarr_region_index_write(self, tmp_path):
-        x = np.arange(0, 50, 10)
-        y = np.arange(0, 20, 2)
-        data = np.ones((5, 10))
-        ds = xr.Dataset(
-            {
-                "test": xr.DataArray(
-                    data,
-                    dims=("x", "y"),
-                    coords={"x": x, "y": y},
-                )
-            }
+    @requires_dask
+    def test_zarr_region_chunk_partial(self):
+        """
+        Check that writing to partial chunks with `region` fails, assuming `safe_chunks=False`.
+        """
+        ds = (
+            xr.DataArray(np.arange(120).reshape(4, 3, -1), dims=list("abc"))
+            .rename("var1")
+            .to_dataset()
         )
 
-        region_slice = dict(x=slice(2, 4), y=slice(6, 8))
-        ds_region = 1 + ds.isel(region_slice)
+        with self.create_zarr_target() as target:
+            self.save(target, ds.chunk(5), compute=False, mode="w")
+            with pytest.raises(ValueError):
+                for r in range(ds.sizes["a"]):
+                    self.save(
+                        target, ds.chunk(3).isel(a=[r]), region=dict(a=slice(r, r + 1))
+                    )
 
-        ds.to_zarr(tmp_path / "test.zarr")
+    @requires_dask
+    def test_zarr_append_chunk_partial(self):
+        t_coords = np.array([np.datetime64("2020-01-01").astype("datetime64[ns]")])
+        data = np.ones((10, 10))
 
-        region: Mapping[str, slice] | Literal["auto"]
-        for region in [region_slice, "auto"]:  # type: ignore[assignment]
-            with patch.object(
-                ZarrStore,
-                "set_variables",
-                side_effect=ZarrStore.set_variables,
-                autospec=True,
-            ) as mock:
-                ds_region.to_zarr(tmp_path / "test.zarr", region=region, mode="r+")
+        da = xr.DataArray(
+            data.reshape((-1, 10, 10)),
+            dims=["time", "x", "y"],
+            coords={"time": t_coords},
+            name="foo",
+        )
+        new_time = np.array([np.datetime64("2021-01-01").astype("datetime64[ns]")])
+        da2 = xr.DataArray(
+            data.reshape((-1, 10, 10)),
+            dims=["time", "x", "y"],
+            coords={"time": new_time},
+            name="foo",
+        )
 
-                # should write the data vars but never the index vars with auto mode
-                for call in mock.call_args_list:
-                    written_variables = call.args[1].keys()
-                    assert "test" in written_variables
-                    assert "x" not in written_variables
-                    assert "y" not in written_variables
+        with self.create_zarr_target() as target:
+            self.save(target, da, mode="w", encoding={"foo": {"chunks": (5, 5, 1)}})
 
-    def test_zarr_region_append(self, tmp_path):
-        x = np.arange(0, 50, 10)
-        y = np.arange(0, 20, 2)
-        data = np.ones((5, 10))
-        ds = xr.Dataset(
-            {
-                "test": xr.DataArray(
-                    data,
-                    dims=("x", "y"),
-                    coords={"x": x, "y": y},
+            with pytest.raises(ValueError, match="encoding was provided"):
+                self.save(
+                    target,
+                    da2,
+                    append_dim="time",
+                    mode="a",
+                    encoding={"foo": {"chunks": (1, 1, 1)}},
                 )
-            }
-        )
-        ds.to_zarr(tmp_path / "test.zarr")
 
-        x_new = np.arange(40, 70, 10)
-        data_new = np.ones((3, 10))
-        ds_new = xr.Dataset(
-            {
-                "test": xr.DataArray(
-                    data_new,
-                    dims=("x", "y"),
-                    coords={"x": x_new, "y": y},
+            # chunking with dask sidesteps the encoding check, so we need a different check
+            with pytest.raises(ValueError, match="Specified zarr chunks"):
+                self.save(
+                    target,
+                    da2.chunk({"x": 1, "y": 1, "time": 1}),
+                    append_dim="time",
+                    mode="a",
                 )
-            }
-        )
 
-        # Now it is valid to use auto region detection with the append mode,
-        # but it is still unsafe to modify dimensions or metadata using the region
-        # parameter.
-        with pytest.raises(KeyError):
-            ds_new.to_zarr(
-                tmp_path / "test.zarr", mode="a", append_dim="x", region="auto"
-            )
+    @requires_dask
+    def test_zarr_region_chunk_partial_offset(self):
+        # https://github.com/pydata/xarray/pull/8459#issuecomment-1819417545
+        with self.create_zarr_target() as store:
+            data = np.ones((30,))
+            da = xr.DataArray(
+                data, dims=["x"], coords={"x": range(30)}, name="foo"
+            ).chunk(x=10)
+            self.save(store, da, compute=False)
 
+            self.save(store, da.isel(x=slice(10)).chunk(x=(10,)), region="auto")
 
-@requires_zarr
-@pytest.mark.usefixtures("default_zarr_format")
-def test_zarr_region(tmp_path):
-    x = np.arange(0, 50, 10)
-    y = np.arange(0, 20, 2)
-    data = np.ones((5, 10))
-    ds = xr.Dataset(
-        {
-            "test": xr.DataArray(
-                data,
-                dims=("x", "y"),
-                coords={"x": x, "y": y},
+            self.save(
+                store,
+                da.isel(x=slice(5, 25)).chunk(x=(10, 10)),
+                safe_chunks=False,
+                region="auto",
             )
-        }
-    )
-    ds.to_zarr(tmp_path / "test.zarr")
 
-    ds_transposed = ds.transpose("y", "x")
-
-    ds_region = 1 + ds_transposed.isel(x=[0], y=[0])
-    ds_region.to_zarr(
-        tmp_path / "test.zarr", region={"x": slice(0, 1), "y": slice(0, 1)}
-    )
+            with pytest.raises(ValueError):
+                self.save(
+                    store, da.isel(x=slice(5, 25)).chunk(x=(10, 10)), region="auto"
+                )
 
-    # Write without region
-    ds_transposed.to_zarr(tmp_path / "test.zarr", mode="r+")
+    @requires_dask
+    def test_zarr_safe_chunk_append_dim(self):
+        with self.create_zarr_target() as store:
+            data = np.ones((20,))
+            da = xr.DataArray(
+                data, dims=["x"], coords={"x": range(20)}, name="foo"
+            ).chunk(x=5)
 
+            self.save(store, da.isel(x=slice(0, 7)), safe_chunks=True, mode="w")
+            with pytest.raises(ValueError):
+                # If the first chunk is smaller than the border size then raise an error
+                self.save(
+                    store,
+                    da.isel(x=slice(7, 11)).chunk(x=(2, 2)),
+                    append_dim="x",
+                    safe_chunks=True,
+                )
 
-@requires_zarr
-@requires_dask
-@pytest.mark.usefixtures("default_zarr_format")
-def test_zarr_region_chunk_partial(tmp_path):
-    """
-    Check that writing to partial chunks with `region` fails, assuming `safe_chunks=False`.
-    """
-    ds = (
-        xr.DataArray(np.arange(120).reshape(4, 3, -1), dims=list("abc"))
-        .rename("var1")
-        .to_dataset()
-    )
+            self.save(store, da.isel(x=slice(0, 7)), safe_chunks=True, mode="w")
+            # If the first chunk is of the size of the border size then it is valid
+            self.save(
+                store,
+                da.isel(x=slice(7, 11)).chunk(x=(3, 1)),
+                safe_chunks=True,
+                append_dim="x",
+            )
+            assert xr.open_zarr(store)["foo"].equals(da.isel(x=slice(0, 11)))
 
-    ds.chunk(5).to_zarr(tmp_path / "foo.zarr", compute=False, mode="w")
-    with pytest.raises(ValueError):
-        for r in range(ds.sizes["a"]):
-            ds.chunk(3).isel(a=[r]).to_zarr(
-                tmp_path / "foo.zarr", region=dict(a=slice(r, r + 1))
+            self.save(store, da.isel(x=slice(0, 7)), safe_chunks=True, mode="w")
+            # If the first chunk is of the size of the border size + N * zchunk then it is valid
+            self.save(
+                store,
+                da.isel(x=slice(7, 17)).chunk(x=(8, 2)),
+                safe_chunks=True,
+                append_dim="x",
             )
+            assert xr.open_zarr(store)["foo"].equals(da.isel(x=slice(0, 17)))
 
+            self.save(store, da.isel(x=slice(0, 7)), safe_chunks=True, mode="w")
+            with pytest.raises(ValueError):
+                # If the first chunk is valid but the other are not then raise an error
+                self.save(
+                    store,
+                    da.isel(x=slice(7, 14)).chunk(x=(3, 3, 1)),
+                    append_dim="x",
+                    safe_chunks=True,
+                )
 
-@requires_zarr
-@requires_dask
-@pytest.mark.usefixtures("default_zarr_format")
-def test_zarr_append_chunk_partial(tmp_path):
-    t_coords = np.array([np.datetime64("2020-01-01").astype("datetime64[ns]")])
-    data = np.ones((10, 10))
-
-    da = xr.DataArray(
-        data.reshape((-1, 10, 10)),
-        dims=["time", "x", "y"],
-        coords={"time": t_coords},
-        name="foo",
-    )
-    da.to_zarr(tmp_path / "foo.zarr", mode="w", encoding={"foo": {"chunks": (5, 5, 1)}})
+            self.save(store, da.isel(x=slice(0, 7)), safe_chunks=True, mode="w")
+            with pytest.raises(ValueError):
+                # If the first chunk have a size bigger than the border size but not enough
+                # to complete the size of the next chunk then an error must be raised
+                self.save(
+                    store,
+                    da.isel(x=slice(7, 14)).chunk(x=(4, 3)),
+                    append_dim="x",
+                    safe_chunks=True,
+                )
 
-    new_time = np.array([np.datetime64("2021-01-01").astype("datetime64[ns]")])
+            self.save(store, da.isel(x=slice(0, 7)), safe_chunks=True, mode="w")
+            # Append with a single chunk it's totally valid,
+            # and it does not matter the size of the chunk
+            self.save(
+                store,
+                da.isel(x=slice(7, 19)).chunk(x=-1),
+                append_dim="x",
+                safe_chunks=True,
+            )
+            assert xr.open_zarr(store)["foo"].equals(da.isel(x=slice(0, 19)))
 
-    da2 = xr.DataArray(
-        data.reshape((-1, 10, 10)),
-        dims=["time", "x", "y"],
-        coords={"time": new_time},
-        name="foo",
-    )
-    with pytest.raises(ValueError, match="encoding was provided"):
-        da2.to_zarr(
-            tmp_path / "foo.zarr",
-            append_dim="time",
-            mode="a",
-            encoding={"foo": {"chunks": (1, 1, 1)}},
-        )
+    @requires_dask
+    @pytest.mark.parametrize("mode", ["r+", "a"])
+    def test_zarr_safe_chunk_region(self, mode: Literal["r+", "a"]):
+        with self.create_zarr_target() as store:
+            arr = xr.DataArray(
+                list(range(11)), dims=["a"], coords={"a": list(range(11))}, name="foo"
+            ).chunk(a=3)
+            self.save(store, arr, mode="w")
 
-    # chunking with dask sidesteps the encoding check, so we need a different check
-    with pytest.raises(ValueError, match="Specified zarr chunks"):
-        da2.chunk({"x": 1, "y": 1, "time": 1}).to_zarr(
-            tmp_path / "foo.zarr", append_dim="time", mode="a"
-        )
+            with pytest.raises(ValueError):
+                # There are two Dask chunks on the same Zarr chunk,
+                # which means that it is unsafe in any mode
+                self.save(
+                    store,
+                    arr.isel(a=slice(0, 3)).chunk(a=(2, 1)),
+                    region="auto",
+                    mode=mode,
+                )
 
+            with pytest.raises(ValueError):
+                # the first chunk is covering the border size, but it is not
+                # completely covering the second chunk, which means that it is
+                # unsafe in any mode
+                self.save(
+                    store,
+                    arr.isel(a=slice(1, 5)).chunk(a=(3, 1)),
+                    region="auto",
+                    mode=mode,
+                )
 
-@requires_zarr
-@requires_dask
-@pytest.mark.usefixtures("default_zarr_format")
-def test_zarr_region_chunk_partial_offset(tmp_path):
-    # https://github.com/pydata/xarray/pull/8459#issuecomment-1819417545
-    store = tmp_path / "foo.zarr"
-    data = np.ones((30,))
-    da = xr.DataArray(data, dims=["x"], coords={"x": range(30)}, name="foo").chunk(x=10)
-    da.to_zarr(store, compute=False)
+            with pytest.raises(ValueError):
+                # The first chunk is safe but the other two chunks are overlapping with
+                # the same Zarr chunk
+                self.save(
+                    store,
+                    arr.isel(a=slice(0, 5)).chunk(a=(3, 1, 1)),
+                    region="auto",
+                    mode=mode,
+                )
 
-    da.isel(x=slice(10)).chunk(x=(10,)).to_zarr(store, region="auto")
+            # Fully update two contiguous chunks is safe in any mode
+            self.save(store, arr.isel(a=slice(3, 9)), region="auto", mode=mode)
 
-    da.isel(x=slice(5, 25)).chunk(x=(10, 10)).to_zarr(
-        store, safe_chunks=False, region="auto"
-    )
+            # The last chunk is considered full based on their current size (2)
+            self.save(store, arr.isel(a=slice(9, 11)), region="auto", mode=mode)
+            self.save(
+                store, arr.isel(a=slice(6, None)).chunk(a=-1), region="auto", mode=mode
+            )
 
-    with pytest.raises(ValueError):
-        da.isel(x=slice(5, 25)).chunk(x=(10, 10)).to_zarr(store, region="auto")
+            # Write the last chunk of a region partially is safe in "a" mode
+            self.save(store, arr.isel(a=slice(3, 8)), region="auto", mode="a")
+            with pytest.raises(ValueError):
+                # with "r+" mode it is invalid to write partial chunk
+                self.save(store, arr.isel(a=slice(3, 8)), region="auto", mode="r+")
 
+            # This is safe with mode "a", the border size is covered by the first chunk of Dask
+            self.save(
+                store, arr.isel(a=slice(1, 4)).chunk(a=(2, 1)), region="auto", mode="a"
+            )
+            with pytest.raises(ValueError):
+                # This is considered unsafe in mode "r+" because it is writing in a partial chunk
+                self.save(
+                    store,
+                    arr.isel(a=slice(1, 4)).chunk(a=(2, 1)),
+                    region="auto",
+                    mode="r+",
+                )
 
-@requires_zarr
-@requires_dask
-@pytest.mark.usefixtures("default_zarr_format")
-def test_zarr_safe_chunk_append_dim(tmp_path):
-    store = tmp_path / "foo.zarr"
-    data = np.ones((20,))
-    da = xr.DataArray(data, dims=["x"], coords={"x": range(20)}, name="foo").chunk(x=5)
-
-    da.isel(x=slice(0, 7)).to_zarr(store, safe_chunks=True, mode="w")
-    with pytest.raises(ValueError):
-        # If the first chunk is smaller than the border size then raise an error
-        da.isel(x=slice(7, 11)).chunk(x=(2, 2)).to_zarr(
-            store, append_dim="x", safe_chunks=True
-        )
+            # This is safe on mode "a" because there is a single dask chunk
+            self.save(
+                store, arr.isel(a=slice(1, 5)).chunk(a=(4,)), region="auto", mode="a"
+            )
+            with pytest.raises(ValueError):
+                # This is unsafe on mode "r+", because the Dask chunk is partially writing
+                # in the first chunk of Zarr
+                self.save(
+                    store,
+                    arr.isel(a=slice(1, 5)).chunk(a=(4,)),
+                    region="auto",
+                    mode="r+",
+                )
 
-    da.isel(x=slice(0, 7)).to_zarr(store, safe_chunks=True, mode="w")
-    # If the first chunk is of the size of the border size then it is valid
-    da.isel(x=slice(7, 11)).chunk(x=(3, 1)).to_zarr(
-        store, safe_chunks=True, append_dim="x"
-    )
-    assert xr.open_zarr(store)["foo"].equals(da.isel(x=slice(0, 11)))
+            # The first chunk is completely covering the first Zarr chunk
+            # and the last chunk is a partial one
+            self.save(
+                store, arr.isel(a=slice(0, 5)).chunk(a=(3, 2)), region="auto", mode="a"
+            )
 
-    da.isel(x=slice(0, 7)).to_zarr(store, safe_chunks=True, mode="w")
-    # If the first chunk is of the size of the border size + N * zchunk then it is valid
-    da.isel(x=slice(7, 17)).chunk(x=(8, 2)).to_zarr(
-        store, safe_chunks=True, append_dim="x"
-    )
-    assert xr.open_zarr(store)["foo"].equals(da.isel(x=slice(0, 17)))
+            with pytest.raises(ValueError):
+                # The last chunk is partial, so it is considered unsafe on mode "r+"
+                self.save(
+                    store,
+                    arr.isel(a=slice(0, 5)).chunk(a=(3, 2)),
+                    region="auto",
+                    mode="r+",
+                )
 
-    da.isel(x=slice(0, 7)).to_zarr(store, safe_chunks=True, mode="w")
-    with pytest.raises(ValueError):
-        # If the first chunk is valid but the other are not then raise an error
-        da.isel(x=slice(7, 14)).chunk(x=(3, 3, 1)).to_zarr(
-            store, append_dim="x", safe_chunks=True
-        )
+            # The first chunk is covering the border size (2 elements)
+            # and also the second chunk (3 elements), so it is valid
+            self.save(
+                store, arr.isel(a=slice(1, 8)).chunk(a=(5, 2)), region="auto", mode="a"
+            )
 
-    da.isel(x=slice(0, 7)).to_zarr(store, safe_chunks=True, mode="w")
-    with pytest.raises(ValueError):
-        # If the first chunk have a size bigger than the border size but not enough
-        # to complete the size of the next chunk then an error must be raised
-        da.isel(x=slice(7, 14)).chunk(x=(4, 3)).to_zarr(
-            store, append_dim="x", safe_chunks=True
-        )
+            with pytest.raises(ValueError):
+                # The first chunk is not fully covering the first zarr chunk
+                self.save(
+                    store,
+                    arr.isel(a=slice(1, 8)).chunk(a=(5, 2)),
+                    region="auto",
+                    mode="r+",
+                )
 
-    da.isel(x=slice(0, 7)).to_zarr(store, safe_chunks=True, mode="w")
-    # Append with a single chunk it's totally valid,
-    # and it does not matter the size of the chunk
-    da.isel(x=slice(7, 19)).chunk(x=-1).to_zarr(store, append_dim="x", safe_chunks=True)
-    assert xr.open_zarr(store)["foo"].equals(da.isel(x=slice(0, 19)))
+            with pytest.raises(ValueError):
+                # Validate that the border condition is not affecting the "r+" mode
+                self.save(store, arr.isel(a=slice(1, 9)), region="auto", mode="r+")
 
+            self.save(store, arr.isel(a=slice(10, 11)), region="auto", mode="a")
+            with pytest.raises(ValueError):
+                # Validate that even if we write with a single Dask chunk on the last Zarr
+                # chunk it is still unsafe if it is not fully covering it
+                # (the last Zarr chunk has size 2)
+                self.save(store, arr.isel(a=slice(10, 11)), region="auto", mode="r+")
 
-@requires_zarr
-@requires_dask
-@pytest.mark.usefixtures("default_zarr_format")
-def test_zarr_safe_chunk_region(tmp_path):
-    store = tmp_path / "foo.zarr"
-
-    arr = xr.DataArray(
-        list(range(11)), dims=["a"], coords={"a": list(range(11))}, name="foo"
-    ).chunk(a=3)
-    arr.to_zarr(store, mode="w")
-
-    modes: list[Literal["r+", "a"]] = ["r+", "a"]
-    for mode in modes:
-        with pytest.raises(ValueError):
-            # There are two Dask chunks on the same Zarr chunk,
-            # which means that it is unsafe in any mode
-            arr.isel(a=slice(0, 3)).chunk(a=(2, 1)).to_zarr(
-                store, region="auto", mode=mode
-            )
+            # Validate the same as the above test but in the beginning of the last chunk
+            self.save(store, arr.isel(a=slice(9, 10)), region="auto", mode="a")
+            with pytest.raises(ValueError):
+                self.save(store, arr.isel(a=slice(9, 10)), region="auto", mode="r+")
 
-        with pytest.raises(ValueError):
-            # the first chunk is covering the border size, but it is not
-            # completely covering the second chunk, which means that it is
-            # unsafe in any mode
-            arr.isel(a=slice(1, 5)).chunk(a=(3, 1)).to_zarr(
-                store, region="auto", mode=mode
+            self.save(
+                store, arr.isel(a=slice(7, None)).chunk(a=-1), region="auto", mode="a"
             )
+            with pytest.raises(ValueError):
+                # Test that even a Dask chunk that covers the last Zarr chunk can be unsafe
+                # if it is partial covering other Zarr chunks
+                self.save(
+                    store,
+                    arr.isel(a=slice(7, None)).chunk(a=-1),
+                    region="auto",
+                    mode="r+",
+                )
 
-        with pytest.raises(ValueError):
-            # The first chunk is safe but the other two chunks are overlapping with
-            # the same Zarr chunk
-            arr.isel(a=slice(0, 5)).chunk(a=(3, 1, 1)).to_zarr(
-                store, region="auto", mode=mode
-            )
+            with pytest.raises(ValueError):
+                # If the chunk is of size equal to the one in the Zarr encoding, but
+                # it is partially writing in the first chunk then raise an error
+                self.save(
+                    store,
+                    arr.isel(a=slice(8, None)).chunk(a=3),
+                    region="auto",
+                    mode="r+",
+                )
+
+            with pytest.raises(ValueError):
+                self.save(
+                    store, arr.isel(a=slice(5, -1)).chunk(a=5), region="auto", mode="r+"
+                )
 
-        # Fully update two contiguous chunks is safe in any mode
-        arr.isel(a=slice(3, 9)).to_zarr(store, region="auto", mode=mode)
-
-        # The last chunk is considered full based on their current size (2)
-        arr.isel(a=slice(9, 11)).to_zarr(store, region="auto", mode=mode)
-        arr.isel(a=slice(6, None)).chunk(a=-1).to_zarr(store, region="auto", mode=mode)
-
-    # Write the last chunk of a region partially is safe in "a" mode
-    arr.isel(a=slice(3, 8)).to_zarr(store, region="auto", mode="a")
-    with pytest.raises(ValueError):
-        # with "r+" mode it is invalid to write partial chunk
-        arr.isel(a=slice(3, 8)).to_zarr(store, region="auto", mode="r+")
-
-    # This is safe with mode "a", the border size is covered by the first chunk of Dask
-    arr.isel(a=slice(1, 4)).chunk(a=(2, 1)).to_zarr(store, region="auto", mode="a")
-    with pytest.raises(ValueError):
-        # This is considered unsafe in mode "r+" because it is writing in a partial chunk
-        arr.isel(a=slice(1, 4)).chunk(a=(2, 1)).to_zarr(store, region="auto", mode="r+")
-
-    # This is safe on mode "a" because there is a single dask chunk
-    arr.isel(a=slice(1, 5)).chunk(a=(4,)).to_zarr(store, region="auto", mode="a")
-    with pytest.raises(ValueError):
-        # This is unsafe on mode "r+", because the Dask chunk is partially writing
-        # in the first chunk of Zarr
-        arr.isel(a=slice(1, 5)).chunk(a=(4,)).to_zarr(store, region="auto", mode="r+")
-
-    # The first chunk is completely covering the first Zarr chunk
-    # and the last chunk is a partial one
-    arr.isel(a=slice(0, 5)).chunk(a=(3, 2)).to_zarr(store, region="auto", mode="a")
-
-    with pytest.raises(ValueError):
-        # The last chunk is partial, so it is considered unsafe on mode "r+"
-        arr.isel(a=slice(0, 5)).chunk(a=(3, 2)).to_zarr(store, region="auto", mode="r+")
-
-    # The first chunk is covering the border size (2 elements)
-    # and also the second chunk (3 elements), so it is valid
-    arr.isel(a=slice(1, 8)).chunk(a=(5, 2)).to_zarr(store, region="auto", mode="a")
-
-    with pytest.raises(ValueError):
-        # The first chunk is not fully covering the first zarr chunk
-        arr.isel(a=slice(1, 8)).chunk(a=(5, 2)).to_zarr(store, region="auto", mode="r+")
-
-    with pytest.raises(ValueError):
-        # Validate that the border condition is not affecting the "r+" mode
-        arr.isel(a=slice(1, 9)).to_zarr(store, region="auto", mode="r+")
-
-    arr.isel(a=slice(10, 11)).to_zarr(store, region="auto", mode="a")
-    with pytest.raises(ValueError):
-        # Validate that even if we write with a single Dask chunk on the last Zarr
-        # chunk it is still unsafe if it is not fully covering it
-        # (the last Zarr chunk has size 2)
-        arr.isel(a=slice(10, 11)).to_zarr(store, region="auto", mode="r+")
-
-    # Validate the same as the above test but in the beginning of the last chunk
-    arr.isel(a=slice(9, 10)).to_zarr(store, region="auto", mode="a")
-    with pytest.raises(ValueError):
-        arr.isel(a=slice(9, 10)).to_zarr(store, region="auto", mode="r+")
-
-    arr.isel(a=slice(7, None)).chunk(a=-1).to_zarr(store, region="auto", mode="a")
-    with pytest.raises(ValueError):
-        # Test that even a Dask chunk that covers the last Zarr chunk can be unsafe
-        # if it is partial covering other Zarr chunks
-        arr.isel(a=slice(7, None)).chunk(a=-1).to_zarr(store, region="auto", mode="r+")
-
-    with pytest.raises(ValueError):
-        # If the chunk is of size equal to the one in the Zarr encoding, but
-        # it is partially writing in the first chunk then raise an error
-        arr.isel(a=slice(8, None)).chunk(a=3).to_zarr(store, region="auto", mode="r+")
-
-    with pytest.raises(ValueError):
-        arr.isel(a=slice(5, -1)).chunk(a=5).to_zarr(store, region="auto", mode="r+")
-
-    # Test if the code is detecting the last chunk correctly
-    data = np.random.default_rng(0).random((2920, 25, 53))
-    ds = xr.Dataset({"temperature": (("time", "lat", "lon"), data)})
-    chunks = {"time": 1000, "lat": 25, "lon": 53}
-    ds.chunk(chunks).to_zarr(store, compute=False, mode="w")
-    region = {"time": slice(1000, 2000, 1)}
-    chunk = ds.isel(region)
-    chunk = chunk.chunk()
-    chunk.chunk().to_zarr(store, region=region)
+            # Test if the code is detecting the last chunk correctly
+            data = np.random.default_rng(0).random((2920, 25, 53))
+            ds = xr.Dataset({"temperature": (("time", "lat", "lon"), data)})
+            chunks = {"time": 1000, "lat": 25, "lon": 53}
+            self.save(store, ds.chunk(chunks), compute=False, mode="w")
+            region = {"time": slice(1000, 2000, 1)}
+            chunk = ds.isel(region)
+            chunk = chunk.chunk()
+            self.save(store, chunk.chunk(), region=region)
 
 
 @requires_h5netcdf
diff --git a/xarray/tests/test_duck_array_ops.py b/xarray/tests/test_duck_array_ops.py
index 9b8b50d9..0771e003 100644
--- a/xarray/tests/test_duck_array_ops.py
+++ b/xarray/tests/test_duck_array_ops.py
@@ -1025,31 +1025,40 @@ def test_least_squares(use_dask, skipna):
 @requires_dask
 @requires_bottleneck
 @pytest.mark.parametrize("method", ["sequential", "blelloch"])
-def test_push_dask(method):
+@pytest.mark.parametrize(
+    "arr",
+    [
+        [np.nan, 1, 2, 3, np.nan, np.nan, np.nan, np.nan, 4, 5, np.nan, 6],
+        [
+            np.nan,
+            np.nan,
+            np.nan,
+            2,
+            np.nan,
+            np.nan,
+            np.nan,
+            9,
+            np.nan,
+            np.nan,
+            np.nan,
+            np.nan,
+        ],
+    ],
+)
+def test_push_dask(method, arr):
     import bottleneck
-    import dask.array
+    import dask.array as da
 
-    array = np.array([np.nan, 1, 2, 3, np.nan, np.nan, np.nan, np.nan, 4, 5, np.nan, 6])
+    arr = np.array(arr)
+    chunks = list(range(1, 11)) + [(1, 2, 3, 2, 2, 1, 1)]
 
     for n in [None, 1, 2, 3, 4, 5, 11]:
-        expected = bottleneck.push(array, axis=0, n=n)
-        for c in range(1, 11):
+        expected = bottleneck.push(arr, axis=0, n=n)
+        for c in chunks:
             with raise_if_dask_computes():
-                actual = push(
-                    dask.array.from_array(array, chunks=c), axis=0, n=n, method=method
-                )
+                actual = push(da.from_array(arr, chunks=c), axis=0, n=n, method=method)
             np.testing.assert_equal(actual, expected)
 
-        # some chunks of size-1 with NaN
-        with raise_if_dask_computes():
-            actual = push(
-                dask.array.from_array(array, chunks=(1, 2, 3, 2, 2, 1, 1)),
-                axis=0,
-                n=n,
-                method=method,
-            )
-        np.testing.assert_equal(actual, expected)
-
 
 def test_extension_array_equality(categorical1, int1):
     int_duck_array = PandasExtensionArray(int1)