dataset.to_zarr(store=zarr_store, compute=True) results in memory spike which isn't cleared

What happened?

I am using rioxarray which in turn is using xarray. We use to_zarr from dataset which is an xarray object.
When using dataset to write to zarr store using the to_zarr method. A spike in memory is observed which is not freed post the execution of the upload.

What did you expect to happen?

When using dataset to write to zarr store using the to_zarr method. A spike in memory is observed, which is then freed post the completion of execution.
The input to it is a stacked multiband raster created out of sentinel tile.

Minimal Complete Verifiable Example

import os

import xarray
from fsspec.mapping import FSMap
import rioxarray
from fsspec.implementations.local import LocalFileSystem
from memory_profiler import profile

@profile
def _update_zarr_store_with_mosaic(
    path_to_mosaic_file: str,
    date: str,
    index: int,
) -> bool:
    local_fs: LocalFileSystem = LocalFileSystem()
    parent_path = "/Users/punit/work/atlas"
    zarr_store_path = os.path.join(parent_path, "testdata/zarr_store")
    zarr_store: FSMap = local_fs.get_mapper(root=zarr_store_path)
    data = rioxarray.open_rasterio(path_to_mosaic_file, chunks=True)

    data = data.expand_dims({"time": 1})
    data = data.assign_coords(time=[date])
    data.attrs = {}
    dataset = data.to_dataset(name="data")
    try:
        if index == 0:
            # during first push, we need to create time dimension hence we don't call append dim
            zarr_out = dataset.to_zarr(store=zarr_store)
        else:
            # during consequent pushes to zarr store, we want new data along the time dimension
            zarr_out = dataset.to_zarr(store=zarr_store, append_dim="time")
    except Exception:
        return False
    return True

if __name__ == "__main__":
    date_to_mosaic_file_path_dict = {
        "2024-02-27": "path_to_file_1",
        "2024-03-03": "path_to_file_2",
        "2024-03-08": "path_to_file_3",
    }
    xarray.show_versions()
    for idx, date in enumerate(sorted(date_to_mosaic_file_path_dict.keys())):
        _update_zarr_store_with_mosaic(date_to_mosaic_file_path_dict[date], date, idx)

MVCE confirmation

• Minimal example — the example is as focused as reasonably possible to demonstrate the underlying issue in xarray.
• Complete example — the example is self-contained, including all data and the text of any traceback.
• Verifiable example — the example copy & pastes into an IPython prompt or Binder notebook, returning the result.
• New issue — a search of GitHub Issues suggests this is not a duplicate.
• Recent environment — the issue occurs with the latest version of xarray and its dependencies.

Relevant log output

INSTALLED VERSIONS
------------------
commit: None
python: 3.11.6 (main, Nov  2 2023, 04:39:43) [Clang 14.0.3 (clang-1403.0.22.14.1)]
python-bits: 64
OS: Darwin
OS-release: 22.4.0
machine: arm64
processor: arm
byteorder: little
LC_ALL: None
LANG: None
LOCALE: (None, 'UTF-8')
libhdf5: None
libnetcdf: None

xarray: 2024.2.0
pandas: 2.2.1
numpy: 1.24.1
scipy: None
netCDF4: None
pydap: None
h5netcdf: None
h5py: None
Nio: None
zarr: 2.13.3
cftime: None
nc_time_axis: None
iris: None
bottleneck: None
dask: 2022.12.1
distributed: None
matplotlib: 3.8.3
cartopy: None
seaborn: None
numbagg: None
fsspec: 2023.6.0
cupy: None
pint: None
sparse: None
flox: None
numpy_groupies: None
setuptools: 69.1.1
pip: 23.3.1
conda: None
pytest: 7.2.0
mypy: None
IPython: None
sphinx: None
Filename: /Users/punit/work/atlas/sample.py

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
     9    134.3 MiB    134.3 MiB           1   @profile
    10                                         def _update_zarr_store_with_mosaic(
    11                                             path_to_mosaic_file: str,
    12                                             date: str,
    13                                             index: int,
    14                                         ) -> bool:
    15    134.3 MiB      0.0 MiB           1       local_fs: LocalFileSystem = LocalFileSystem()
    16    134.3 MiB      0.0 MiB           1       parent_path = "/Users/punit/work/atlas"
    17    134.3 MiB      0.0 MiB           1       zarr_store_path = os.path.join(parent_path, "testdata/zarr_store")
    18    134.3 MiB      0.0 MiB           1       zarr_store: FSMap = local_fs.get_mapper(root=zarr_store_path)
    19    142.8 MiB      8.5 MiB           1       data = rioxarray.open_rasterio(path_to_mosaic_file, chunks=True)
    20                                         
    21    142.8 MiB      0.0 MiB           1       data = data.expand_dims({"time": 1})
    22    142.8 MiB      0.1 MiB           1       data = data.assign_coords(time=[date])
    23    142.8 MiB      0.0 MiB           1       data.attrs = {}
    24    142.8 MiB      0.0 MiB           1       dataset = data.to_dataset(name="data")
    25    142.8 MiB      0.0 MiB           1       try:
    26    142.8 MiB      0.0 MiB           1           if index == 0:
    27                                                     # during first push, we need to create time dimension hence we don't call append dim
    28    142.8 MiB      0.0 MiB           1               zarr_out = dataset.to_zarr(store=zarr_store)
    29                                                 else:
    30                                                     # during consequent pushes to zarr store, we want new data along the time dimension
    31                                                     zarr_out = dataset.to_zarr(store=zarr_store, append_dim="time")
    32    142.8 MiB      0.0 MiB           1       except Exception:
    33    142.8 MiB      0.0 MiB           1           return False
    34                                             return True


Filename: /Users/punit/work/atlas/sample.py

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
     9    142.8 MiB    142.8 MiB           1   @profile
    10                                         def _update_zarr_store_with_mosaic(
    11                                             path_to_mosaic_file: str,
    12                                             date: str,
    13                                             index: int,
    14                                         ) -> bool:
    15    142.8 MiB      0.0 MiB           1       local_fs: LocalFileSystem = LocalFileSystem()
    16    142.8 MiB      0.0 MiB           1       parent_path = "/Users/punit/work/atlas"
    17    142.8 MiB      0.0 MiB           1       zarr_store_path = os.path.join(parent_path, "testdata/zarr_store")
    18    142.8 MiB      0.0 MiB           1       zarr_store: FSMap = local_fs.get_mapper(root=zarr_store_path)
    19    143.2 MiB      0.3 MiB           1       data = rioxarray.open_rasterio(path_to_mosaic_file, chunks=True)
    20                                         
    21    143.2 MiB      0.0 MiB           1       data = data.expand_dims({"time": 1})
    22    143.2 MiB      0.0 MiB           1       data = data.assign_coords(time=[date])
    23    143.2 MiB      0.0 MiB           1       data.attrs = {}
    24    143.2 MiB      0.0 MiB           1       dataset = data.to_dataset(name="data")
    25    143.2 MiB      0.0 MiB           1       try:
    26    143.2 MiB      0.0 MiB           1           if index == 0:
    27                                                     # during first push, we need to create time dimension hence we don't call append dim
    28                                                     zarr_out = dataset.to_zarr(store=zarr_store)
    29                                                 else:
    30                                                     # during consequent pushes to zarr store, we want new data along the time dimension
    31    511.9 MiB    368.7 MiB           1               zarr_out = dataset.to_zarr(store=zarr_store, append_dim="time")
    32                                             except Exception:
    33                                                 return False
    34    511.9 MiB      0.0 MiB           1       return True


Filename: /Users/punit/work/atlas/sample.py

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
     9    511.9 MiB    511.9 MiB           1   @profile
    10                                         def _update_zarr_store_with_mosaic(
    11                                             path_to_mosaic_file: str,
    12                                             date: str,
    13                                             index: int,
    14                                         ) -> bool:
    15    511.9 MiB      0.0 MiB           1       local_fs: LocalFileSystem = LocalFileSystem()
    16    511.9 MiB      0.0 MiB           1       parent_path = "/Users/punit/work/atlas"
    17    511.9 MiB      0.0 MiB           1       zarr_store_path = os.path.join(parent_path, "testdata/zarr_store")
    18    511.9 MiB      0.0 MiB           1       zarr_store: FSMap = local_fs.get_mapper(root=zarr_store_path)
    19    512.2 MiB      0.3 MiB           1       data = rioxarray.open_rasterio(path_to_mosaic_file, chunks=True)
    20                                         
    21    512.2 MiB      0.0 MiB           1       data = data.expand_dims({"time": 1})
    22    512.2 MiB      0.0 MiB           1       data = data.assign_coords(time=[date])
    23    512.2 MiB      0.0 MiB           1       data.attrs = {}
    24    512.2 MiB      0.0 MiB           1       dataset = data.to_dataset(name="data")
    25    512.2 MiB      0.0 MiB           1       try:
    26    512.2 MiB      0.0 MiB           1           if index == 0:
    27                                                     # during first push, we need to create time dimension hence we don't call append dim
    28                                                     zarr_out = dataset.to_zarr(store=zarr_store)
    29                                                 else:
    30                                                     # during consequent pushes to zarr store, we want new data along the time dimension
    31    543.0 MiB     30.9 MiB           1               zarr_out = dataset.to_zarr(store=zarr_store, append_dim="time")
    32                                             except Exception:
    33                                                 return False
    34    543.0 MiB     -0.0 MiB           1       return True

Anything else we need to know?

1. The log is an output of the code running with three files that is the dictionary contains three paths.
2. I can provide the input tiff files, but they are quite big (1.4, 1.1, 1.1 GB)respectively.

Environment

The entry point is Dataset.to_zarr with append_dim="time"; reproduce the second write using the supplied memory_profiler example and compare it with the first write. Trace the xarray and Dask write path using the reported dependency versions to isolate whether the retained memory is expected; done means a confirmed cause and a regression test or clearly documented behavior.