OceanDataStore 0.3.0__py3-none-any.whl

OceanDataStore/zarr.py

@@ -0,0 +1,993 @@
+"""
+zarr.py
+
+Description:
+This module defines the functions to send and update data
+to an object store.
+
+Authors:
+    - Ollie Tooth
+    - Tobias Ferreira
+    - Joao Morado
+"""
+# -- Import Python Modules -- #
+import glob
+import time
+import logging
+import warnings
+from typing import Optional
+
+import numpy as np
+import xarray as xr
+
+import dask
+from dask.distributed import Client, LocalCluster
+from dask.distributed.diagnostics.plugin import WorkerPlugin
+
+# -- Import OceanDataStore Modules -- #
+from .object_store import ObjectStoreS3
+
+from .exceptions import (
+    ObjectNotFound,
+    DimensionNotFound,
+    DimensionSizeError,
+    AppendDimensionError,
+    AppendDimensionSizeError,
+    ChunkSizeError,
+)
+
+# -- Define WorkerPlugin -- #
+class CaptureWarningsPlugin(WorkerPlugin):
+    def setup(self, worker):
+        # Used to catch UserWarnings when rechunking:
+        logging.captureWarnings(True)
+    def teardown(self, worker):
+        logging.captureWarnings(False)
+
+
+# -- Define timing context manager -- #
+class timer():
+    """
+    Timer context manager class to return time
+    taken to write variables & datasets to an
+    object store.
+
+    Parameters
+    ----------
+    action : str
+        Action to be performed. Options are 'send' or 'update'.
+    url : str
+        URL path to Zarr store or Icechunk repository.
+    var : Optional[str], default=None
+        Name of variable to be sent or updated to store.
+    """
+    def __init__(self, action: str, url: str, var: Optional[str] = None) -> None:
+        # Define class attributes:
+        if action == 'send':
+            if var is not None:
+                self.action = f'Sent {var} to'
+            else:
+                self.action = 'Sent dataset to'
+        elif action == 'replace':
+            if var is not None:
+                self.action = f'Updated {var} in'
+            else:
+                self.action = 'Updated'
+        elif action == 'append':
+            if var is not None:
+                self.action = f'Appended {var} to'
+            else:
+                self.action = 'Appended to'
+        else:
+            raise ValueError("Invalid action: must be 'send', 'replace' or 'append'.")
+        self.url = url
+
+    def __enter__(self):
+        self.t_start = time.time()
+
+    def __exit__(self, type, value, traceback):
+        self.t_end = time.time()
+        logging.info(
+            f"Completed: {self.action} store {self.url} in {(self.t_end - self.t_start):.2f} seconds"
+            )
+
+
+# -- Define OceanDataStore Core Functions -- #
+def _check_zarr_store(obj_store: ObjectStoreS3,
+                      url: str
+                      ) -> bool:
+    """
+    Check if a Zarr store exists at a specified URL path.
+
+    Parameters
+    ----------
+    obj_store
+        ObjectStoreS3 remote filesystem.
+    url
+        URL path to Zarr store.
+    
+    Returns
+    -------
+    bool
+        True if the store exists, False otherwise.
+    """
+
+    return obj_store.exists(url.replace("s3://", ""))
+
+
+def _check_zarr_compatibility(data: xr.DataArray | xr.Dataset,
+                              obj_store: ObjectStoreS3,
+                              url: str,
+                              append_dim: str = "time_counter",
+                              rechunk: Optional[dict] = None,
+                              version: int = 3,
+                              ) -> None:
+    """
+    Check compatibility of DataArray or Dataset to update existing
+    Zarr store in cloud object storage.
+
+    Parameters
+    ----------
+    data: xr.DataArray | xr.Dataset
+        DataArray or DataSet to update existing Zarr store with.
+    obj_store: ObjectStoreS3
+        ObjectStoreS3 remote filesystem.
+    url: str
+        URL path to Zarr store.
+    append_dim: bool, default="time_counter"
+        Dimension to append data to existing Zarr store.
+    rechunk: Optional[dict], default=None
+        Mapping to rechunk dimensions.
+    version: int, default=3
+        Zarr version to use.
+    """
+    # 1. Check if the store exists:
+    if not _check_zarr_store(obj_store=obj_store, path=url):
+        raise ObjectNotFound(object_name=url)
+    
+    # 2. Check Zarr store compatibility:
+    try:
+        ds_store = xr.open_zarr(store=url,
+                                storage_options=obj_store.get_remote_options(),
+                                zarr_format=version
+                                )
+    except Exception as e:
+        raise FileNotFoundError(f"zarr version {version} is not compatible with the store: {e}")
+    
+    # 3. Check if core dimensions exist & size are compatible:
+    dims_data = {dim : data.sizes[dim] for dim in data.dims if dim != append_dim}
+    for dim in dims_data:
+        if dim in ds_store.dims:
+            if dims_data[dim] != ds_store.sizes[dim]:
+                raise DimensionSizeError(dim=dim, size=dims_data[dim], expected_size=ds_store.sizes[dim])
+        else:
+            raise DimensionNotFound(dim=dim, object_name=url)
+
+    # 4. Check if append dimension values are compatible:
+    if (data[append_dim][0] < ds_store[append_dim][0]):
+        raise AppendDimensionError(dim=append_dim)
+    
+    # 5. Check if specified chunks are compatible:
+    if rechunk is not None:
+        for dim in rechunk:
+            if dim in ds_store.dims:
+                if rechunk[dim] != ds_store.chunks[dim][0]:
+                    raise ChunkSizeError(chunks=rechunk, store_chunks=ds_store.chunks)
+
+
+def _write_to_zarr(
+        data: xr.DataArray | xr.Dataset,
+        obj_store: ObjectStoreS3,
+        url: str,
+        version: int = 3,
+        ) -> None:
+    """
+    Write DataArray or Dataset to Zarr store in cloud
+    object storage.
+
+    Parameters
+    ----------
+    data: xr.DataArray | xr.Dataset
+        DataArray or DataSet to write to Zarr store.
+    obj_store: ObjectStoreS3
+        ObjectStoreS3 remote filesystem.
+    url: str
+        URL path to Zarr store.
+    version: int, default=3
+        Zarr version to use.
+    """
+    # === Verify Inputs === #
+    if not isinstance(data, (xr.DataArray, xr.Dataset)):
+        raise TypeError("data must be a DataArray or Dataset.")
+    if not isinstance(obj_store, ObjectStoreS3):
+        raise TypeError("obj_store must be an ObjectStoreS3 instance.")
+    if not isinstance(url, str):
+        raise TypeError("url must be a string.")
+    if not isinstance(version, int):
+        raise TypeError("version must be an integer.")
+
+    # Convert DataArrays to Datasets:
+    if isinstance(data, xr.DataArray):
+        var = data.name
+        data = data.to_dataset()
+    else:
+        var = None
+
+    # Write Dataset to Zarr store in Object Store:
+    if _check_zarr_store(obj_store=obj_store, path=url):
+        logging.info(f"Skipping Variable: Store already exists at {url}")
+
+    else:
+        with timer(action='send', url=url, var=var):
+            # Catch consolidated metadata warnings:
+            with warnings.catch_warnings():
+                warnings.simplefilter(action="ignore", category=UserWarning)
+                data.to_zarr(store=url,
+                             storage_options=obj_store.get_remote_options(),
+                             mode="w",
+                             zarr_format=version
+                             )
+
+
+def _append_to_zarr(data: xr.DataArray | xr.Dataset,
+                    obj_store: ObjectStoreS3,
+                    url: str,
+                    append_dim: str = "time_counter",
+                    version: int = 3,
+                    ) -> None:
+    """
+    Append DataArray or Dataset to existing Zarr store in
+    cloud object storage.
+
+    Parameters
+    ----------
+    data: xr.DataArray | xr.Dataset
+        DataArray or DataSet to append to existing Zarr store.
+    obj_store: ObjectStoreS3
+        ObjectStoreS3 remote filesystem.
+    url: str
+        URL path to Zarr store.
+    append_dim: str, default="time_counter"
+        Dimension to append data to existing Zarr store.
+    version: int, default=3
+        Zarr version to use.
+    """
+    with timer(action='append', url=url):
+        # Catch consolidated metadata warnings:
+        with warnings.catch_warnings():
+            warnings.simplefilter(action="ignore", category=UserWarning)
+            data.to_zarr(store=url,
+                         storage_options=obj_store.get_remote_options(),
+                         append_dim=append_dim,
+                         zarr_format=version
+                         )
+
+
+def _replace_in_zarr(data: xr.DataArray | xr.Dataset,
+                     obj_store: ObjectStoreS3,
+                     url: str,
+                     region: dict,
+                     version: int = 3,
+                     ) -> None:
+    """
+    Append DataArray or Dataset to existing Zarr store in
+    cloud object storage.
+
+    Parameters
+    ----------
+    data: xr.DataArray | xr.Dataset
+        DataArray or DataSet to append to existing Zarr store.
+    obj_store: ObjectStoreS3
+        ObjectStoreS3 remote filesystem.
+    url: str
+        URL path to Zarr store.
+    region: dict
+        Region of existing Zarr store to replace data.
+    version: int, default=3
+        Zarr version to use.
+    """
+    # Drop variables w/o append dimension:
+    append_dim = list(region.keys())[0]
+    drop_list = [var for var in data.variables if append_dim not in data[var].dims]
+    data = data.drop_vars(drop_list)
+
+    with timer(action='replace', url=url):
+        # Catch consolidated metadata warnings:
+        with warnings.catch_warnings():
+            warnings.simplefilter(action="ignore", category=UserWarning)
+            data.to_zarr(store=url,
+                         storage_options=obj_store.get_remote_options(),
+                         region=region,
+                         zarr_format=version
+                         )
+
+
+def _update_zarr_store(data: xr.DataArray | xr.Dataset,
+                       obj_store: ObjectStoreS3,
+                       url: str,
+                       append_dim: str = "time_counter",
+                       rechunk: Optional[dict] = None,
+                       version: int = 3,
+                       ) -> None:
+    """
+    Update an existing Zarr store in object storage by replacing
+    existing values and/or appending new values.
+
+    Parameters
+    ----------
+    data: xr.DataArray | xr.Dataset
+        DataArray or DataSet to update existing Zarr store with.
+    obj_store: ObjectStoreS3
+        ObjectStoreS3 remote filesystem.
+    url: str
+        URL path to Zarr store.
+    append_dim: bool, default="time_counter"
+        Dimension to append data to existing Zarr store.
+    rechunk: Optional[dict], default=None
+        Mapping to rechunk dimensions.
+    version: int, default=3
+        Zarr version to use.
+    """
+    # === Verify Inputs === #
+    if not isinstance(data, (xr.DataArray, xr.Dataset)):
+        raise TypeError("data must be a DataArray or Dataset.")
+    if not isinstance(obj_store, ObjectStoreS3):
+        raise TypeError("obj_store must be an ObjectStoreS3 instance.")
+    if not isinstance(url, str):
+        raise TypeError("url must be a string.")
+    if not isinstance(append_dim, str):
+        raise TypeError("append_dim must be a string.")
+    if rechunk is not None:
+        if not isinstance(rechunk, dict):
+            raise TypeError("rechunk must be a dictionary.")
+    if not isinstance(version, int):
+        raise TypeError("version must be an integer.")
+
+    # Convert DataArrays to Datasets:
+    if isinstance(data, xr.DataArray):
+        var = data.name
+        ds_source = data.to_dataset()
+    else:
+        var = None
+        ds_source = data
+
+    # Check source Dataset compatibility with existing store:
+    _check_zarr_compatibility(data=ds_source,
+                              obj_store=obj_store,
+                              url=url,
+                              append_dim=append_dim,
+                              rechunk=rechunk,
+                              version=version
+                              )
+    logging.info(f"Passed Compatibility Checks for store {url}")
+
+    # === Update existing variable in Zarr Store === #
+    # Extract source & target append dimension values:
+    ds_target = xr.open_zarr(store=url,
+                             storage_options=obj_store.get_remote_options(),
+                             zarr_format=version
+                             )
+
+    if (var in ds_target.data_vars) or (var is None):
+
+        # === Updating existing Zarr store === #
+        # Extract source & target append dimension values:
+        target_append_dim = ds_target[append_dim].values
+        source_append_dim = ds_source[append_dim].values
+
+        # Determine intersection between source & target append dimensions:
+        intersect_append_dim = np.intersect1d(source_append_dim, target_append_dim)
+
+        if intersect_append_dim.size != 0:
+            # == Intersection exists -> replace overlapping values in target store == #
+
+            # Ensure all overlapping values exist along target append dimension:
+            overlap_append_dim = (source_append_dim <= target_append_dim[-1]).sum()
+            if intersect_append_dim.size != overlap_append_dim:
+                raise AppendDimensionSizeError(dim=append_dim, size=overlap_append_dim, expected_size=intersect_append_dim.size)
+            
+            # Determine source and target append dimension indices of overlap:
+            target_ind_min = np.flatnonzero(target_append_dim == source_append_dim[0])[0]
+            target_ind_max = target_append_dim.size
+            source_ind_min = 0
+            source_ind_max = target_ind_max - target_ind_min
+            source_ind_size = source_append_dim.size
+
+            # 1. Replace overlapping values in target store:
+            logging.info(f"Updating {url} along {append_dim} from {target_append_dim[target_ind_min]} to {target_append_dim[target_ind_max - 1]}.")
+            _replace_in_zarr(data=ds_source.isel({append_dim : slice(source_ind_min, source_ind_max)}),
+                             obj_store=obj_store,
+                             url=url,
+                             region={append_dim : slice(target_ind_min, target_ind_max)},
+                             version=version,
+                             )
+
+            # 2. Append new values to target store:
+            if source_ind_size > source_ind_max:
+                logging.info(f"Appending to {url} along {append_dim} from {source_append_dim[source_ind_max]} to {source_append_dim[source_ind_size - 1]}.")
+                _append_to_zarr(data=ds_source.isel({append_dim : slice(source_ind_max, source_ind_size)}),
+                                obj_store=obj_store,
+                                url=url,
+                                append_dim=append_dim,
+                                version=version,
+                                )
+
+        else:
+            # == No intersection -> append all source values to target store == #
+            _append_to_zarr(data=ds_source,
+                            obj_store=obj_store,
+                            url=url,
+                            append_dim=append_dim,
+                            version=version,
+                            )
+    else:
+        # == Add new variable to Zarr Store == #
+        logging.info(f"Sending Variable {var}")
+        _write_to_zarr(data=ds_source,
+                       obj_store=obj_store,
+                       url=url,
+                       version=version,
+                       )
+
+
+def _preprocess_dataset(file: list[str] | str | xr.Dataset,
+                        rechunk: Optional[dict] = None,
+                        append_dim: str = "time_counter",
+                        update_coords: Optional[dict] = None,
+                        grid_filepath: Optional[str] = None,
+                        attrs: Optional[dict] = None,
+                        parallel: bool = False,
+                        ) -> xr.Dataset:
+    """
+    Preprocess the dataset to be sent to the object store.
+
+    Parameters
+    ----------
+    file: list | str | xarray.Dataset
+        Regular expression or list of filepaths to netCDF file(s).
+        Users can also pass a single xarray.Dataset directly.
+    rechunk: Optional[dict], default=None
+        Mapping to rechunk dimensions. If None, dask chunks
+        will be set to on-disk chunks.
+    append_dim: str, default='time_counter'
+        Name of the dimension to append multi-file datasets.
+    update_coords: Optional[dict], default=None
+        Mapping of coordinate variables to update using model
+        grid file. Keys are coordinate variable names in the
+        dataset to be sent, and values are the corresponding
+        variable names in the model grid file. If None, no 
+        coordinates will be updated.
+    grid_filepath: Optional[str], default=None
+        Filepath to the model grid file to update coordinate
+        variables. Required if update_coords is not None.
+    attrs: Optional[dict], default=None
+        Dictionary of attributes to add to the dataset.
+        If None, no attributes will be added.
+    parallel: bool, default=False
+        Whether to open and preprocess the dataset in parallel
+        using `dask.delayed`.
+
+    Returns
+    -------
+    xr.Dataset
+        Preprocessed (multifile) dataset with optionally
+        updated coordinates, chunksizes and attributes.
+
+    """
+    # == Verify Inputs == #
+    if not isinstance(file, (list, str, xr.Dataset)):
+        raise TypeError("filepaths must be a list, a string or an xarray Dataset.")
+    if isinstance(file, list):
+        for fpath in file:
+            if not isinstance(fpath, str):
+                raise TypeError("filepaths must be a list of strings.")
+            if not fpath.endswith('.nc'):
+                raise ValueError("Invalid file extension: only .nc files are supported.")
+    elif isinstance(file, str):
+        if not file.endswith('.nc'):
+            raise ValueError("Invalid file extension: only .nc files are supported.")
+    if rechunk is not None:
+        if not isinstance(rechunk, dict):
+            raise TypeError("rechunk must be a dictionary.")
+    if not isinstance(append_dim, str):
+        raise TypeError("append_dim must be a string.")
+    if update_coords is not None:
+        if not isinstance(update_coords, dict):
+            raise TypeError("update_coords must be a dictionary.")
+    if grid_filepath is not None:
+        if not isinstance(grid_filepath, str):
+            raise TypeError("grid_filepath must be a string.")
+    if attrs is not None:
+        if not isinstance(attrs, dict):
+            raise TypeError("attrs must be a dictionary.")
+    if not isinstance(parallel, bool):
+        raise TypeError("parallel must be a boolean.")
+
+    # === Load netCDF dataset === #
+    if rechunk is None:
+        # Default to dask chunks equal to on-disk chunks:
+        rechunk = {}
+
+    # File names from str / regular expression:
+    if isinstance(file, str):
+        if '*' in file:
+            filepaths = sorted(glob.glob(file))
+            if len(filepaths) == 0:
+                raise FileNotFoundError(f"No files found at {filepaths}")
+        else:
+            filepaths = [file]
+    # File names from list:
+    elif isinstance(file, list):
+        filepaths = file
+
+    # Use input dataset:
+    if isinstance(file, xr.Dataset):
+        ds_filepath = file
+        if rechunk is not None:
+            ds_filepath = ds_filepath.chunk(rechunk)
+    else:
+        # Open multi-file dataset:
+        if len(filepaths) > 1:
+            ds_filepath = xr.open_mfdataset(filepaths,
+                                            engine='h5netcdf',
+                                            chunks=rechunk,
+                                            parallel=parallel,
+                                            concat_dim=append_dim,
+                                            combine='nested',
+                                            data_vars='minimal',
+                                            coords='minimal',
+                                            compat='override'
+                                            )
+        else:
+            # Open single file dataset:
+            ds_filepath = xr.open_dataset(filepaths[0], chunks=rechunk)
+
+    # === Update coordinates using model grid file === #
+    if update_coords is not None:
+        if grid_filepath is None:
+            raise ValueError(
+                "grid_filepath must be specified to update coordinate variables."
+                )
+        else:
+            ds_grid = xr.open_dataset(grid_filepath)
+        # Update coordinate vars using model grid file:
+        for key in update_coords.keys():
+            coord_data = ds_grid[update_coords[key]].squeeze(drop=True)
+            # Rechunk dimensions to user specified chunks:
+            if rechunk is not None:
+                coord_chunks = {dim: rechunk[dim] for dim in coord_data.dims}
+                ds_filepath = ds_filepath.assign_coords(
+                    {key: coord_data.chunk(coord_chunks)}
+                    )
+            else:
+                ds_filepath = ds_filepath.assign_coords(
+                    {key: coord_data}
+                    )
+        logging.info('Completed: Updated coordinate variables.')
+
+    # === Update Attributes === #
+    if attrs is not None:
+        ds_filepath = ds_filepath.assign_attrs(attrs)
+
+    return ds_filepath
+
+
+def _send_to_zarr(
+        file: list[str] | str | xr.Dataset,
+        bucket: str,
+        object_prefix: str,
+        store_credentials_json: str,
+        variables: Optional[list[str]] = None,
+        append_dim: str = "time_counter",
+        grid_filepath: Optional[str] = None,
+        update_coords: Optional[dict] = None,
+        rechunk: Optional[dict] = None,
+        attrs: Optional[dict] = None,
+        parallel: bool = False,
+        zarr_version: int = 3
+        ) -> None:
+    """
+    Write data to new Zarr store in cloud object storage.
+
+    Parameters
+    ----------
+    file: list | str | xarray.Dataset
+        Regular expression or list of filepaths to netCDF file(s).
+        Users can also pass a single xarray.Dataset directly.
+    bucket: str
+        Name of the bucket in the object store. Bucket names can contain only
+        lowercase letters, numbers, dots (.), and hyphens (-).
+    object_prefix: str
+        Prefix to be added to the object names in the object store.
+    store_credentials_json: str
+        Path to the JSON file containing the object store credentials.
+    variables: list[str], optional
+        List of variables to send to Zarr stores.
+        If None, all variables will be sent.
+    append_dim: str, default='time_counter'
+        Name of the dimension to append multifile datasets.
+    grid_filepath: str, optional
+        Path to file containing model grid parameter.
+    update_coords: dict, optional
+        Dictionary of coordinate variables to update.
+    rechunk: dict, optional
+        Rechunk strategy dictionary.
+    attrs: dict, optional
+        Attributes to add to the dataset.
+    parallel: bool, default=False,
+        Whether to perform open and preprocess steps in parallel using
+        `dask.delayed`.
+    zarr_version: int, default=3
+        Zarr version to use.
+    """
+    # === Initialise Asynchronous Object Store === #
+    logging.info("Reading object store credentials from %s", store_credentials_json)
+    obj_store = ObjectStoreS3(anon=False,
+                              asynchronous=True,
+                              store_credentials_json=store_credentials_json
+                              )
+
+    # === Preprocess Data === #
+    ds_filepath = _preprocess_dataset(file=file,
+                                      rechunk=rechunk,
+                                      append_dim=append_dim,
+                                      update_coords=update_coords,
+                                      grid_filepath=grid_filepath,
+                                      attrs=attrs,
+                                      parallel=parallel,
+                                      )
+    if variables is None:
+        variables = list(ds_filepath.data_vars)
+
+    # === Send Dataset to Zarr store === #
+    # Write to Zarr store:
+    url = f"s3://{bucket}/{object_prefix}"
+    logging.info(f"Sending Dataset to {url}")
+    _write_to_zarr(data=ds_filepath[variables],
+                   obj_store=obj_store,
+                   url=url,
+                   version=zarr_version
+                   )
+
+    # Release resources to avoid memory leaks:
+    ds_filepath.close()
+
+def send_to_zarr(
+    file: list[str] | str | xr.Dataset,
+    bucket: str,
+    object_prefix: str,
+    store_credentials_json: str,
+    variables: Optional[list[str]] = None,
+    append_dim: str = "time_counter",
+    grid_filepath: Optional[str] = None,
+    update_coords: Optional[dict] = None,
+    rechunk: Optional[dict] = None,
+    attrs: Optional[dict] = None,
+    client : Optional[Client] = None,
+    dask_config_kwargs: Optional[dict] = None,
+    dask_cluster_kwargs: Optional[dict] = None,
+    zarr_version: int = 3
+    ) -> None:
+    """
+    Write data to new Zarr store in cloud object storage with
+    option of using dask.
+
+    Parameters
+    ----------
+    file: list | str | xarray.Dataset
+        Regular expression or list of filepaths to netCDF file(s).
+        Users can also pass a single xarray.Dataset directly.
+    bucket: str
+        Name of the bucket in the object store. Bucket names can contain only
+        lowercase letters, numbers, dots (.), and hyphens (-).
+    object_prefix: str
+        Prefix to be added to the object names in the object store.
+    store_credentials_json: str
+        Path to the JSON file containing the object store credentials.
+    variables: list[str], optional
+        List of variables to send. If None, all variables will be sent.
+    append_dim: str, default="time_counter"
+        Name of the append dimension, by default "time_counter".
+    grid_filepath: str, optional
+        Path to file containing model grid parameter.
+    update_coords: dict, optional
+        Dictionary of coordinate variables to update.
+    rechunk: dict, optional
+        Rechunk strategy dictionary, by default None.
+    attrs: dict, optional
+        Attributes to add to the dataset.
+    client: dask.distributed.Client, optional
+        Dask Distributed Client.
+    dask_config_kwargs: dict[str,str], optional
+        Dask configuration settings passed to dask.config.set().
+    dask_cluster_kwargs: dict, optional
+        Dask cluster configuration settings passed to LocalCluster().
+    zarr_version: int, default=3
+        Zarr version to use.
+    """
+    if dask_cluster_kwargs is not None:
+        # === Send to Zarr store with Dask === #
+        if dask_config_kwargs is not None:
+            dask.config.set(dask_config_kwargs)
+            logging.info("Updated dask configuration settings.")
+
+        # Create local dask cluster & client:
+        with LocalCluster(**dask_cluster_kwargs) as cluster, Client(cluster) as client:
+            logging.info(f"Created LocalCluster with {dask_cluster_kwargs['n_workers']} workers @ Client: {client.dashboard_link}")
+
+            # Catch UserWarnings when rechunking data:
+            client.register_worker_plugin(CaptureWarningsPlugin())
+
+            _send_to_zarr(file=file,
+                          bucket=bucket,
+                          object_prefix=object_prefix,
+                          store_credentials_json=store_credentials_json,
+                          variables=variables,
+                          append_dim=append_dim,
+                          grid_filepath=grid_filepath,
+                          update_coords=update_coords,
+                          rechunk=rechunk,
+                          attrs=attrs,
+                          parallel=True,
+                          zarr_version=zarr_version
+                         )
+
+            # --- Shutdown Store & Dask Cluster --- #
+            cluster.close()
+            client.shutdown()
+            logging.info("Dask Cluster has been shutdown.")
+    
+    elif client is not None:
+        logging.info(f"Using existing Dask Cluster @ Client: {client.dashboard_link}")
+
+        # Catch UserWarnings when rechunking data:
+        client.register_worker_plugin(CaptureWarningsPlugin())
+
+        _send_to_zarr(file=file,
+                      bucket=bucket,
+                      object_prefix=object_prefix,
+                      store_credentials_json=store_credentials_json,
+                      variables=variables,
+                      append_dim=append_dim,
+                      grid_filepath=grid_filepath,
+                      update_coords=update_coords,
+                      rechunk=rechunk,
+                      attrs=attrs,
+                      parallel=True,
+                      zarr_version=zarr_version
+                      )
+
+        # --- Shutdown Store & Dask Cluster --- #
+        cluster.close()
+        client.shutdown()
+        logging.info("Existing Dask Cluster has been shutdown.")
+    
+    else:
+        # === Send to Zarr store without Dask === #
+        _send_to_zarr(file=file,
+                      bucket=bucket,
+                      object_prefix=object_prefix,
+                      store_credentials_json=store_credentials_json,
+                      variables=variables,
+                      append_dim=append_dim,
+                      grid_filepath=grid_filepath,
+                      update_coords=update_coords,
+                      rechunk=rechunk,
+                      attrs=attrs,
+                      parallel=False,
+                      zarr_version=zarr_version
+                      )
+
+def _update_zarr(
+        file: list[str] | str | xr.Dataset,
+        bucket: str,
+        object_prefix: str,
+        store_credentials_json: str,
+        variables: Optional[list[str]] = None,
+        append_dim: str = "time_counter",
+        grid_filepath: Optional[str] = None,
+        update_coords: Optional[dict] = None,
+        rechunk: Optional[dict] = None,
+        attrs: Optional[dict] = None,
+        parallel: bool = False,
+        zarr_version: int = 3
+        ) -> None:
+    """
+    Update existing Zarr store in cloud object storage
+    by replacing and/or appending data.
+
+    Parameters
+    ----------
+    file: list | str
+        Regular expression or list of filepaths to netCDF file(s).
+        Users can also pass a single xarray.Dataset directly.
+    bucket: str
+        Name of the bucket in the object store. Bucket names can contain only
+        lowercase letters, numbers, dots (.), and hyphens (-).
+    object_prefix: str
+        Prefix to be added to the object names in the object store.
+    store_credentials_json: str
+        Path to the JSON file containing the object store credentials.
+    variables: list, optional
+        List of variables to send to Zarr stores.
+        If None, all variables will be sent.
+    append_dim: str, default='time_counter'
+        Name of the dimension to append multifile datasets.
+    grid_filepath: str, optional
+        Path to file containing model grid parameter.
+    update_coords: dict, optional
+        Dictionary of coordinate variables to update.
+    rechunk: dict, optional
+        Rechunk strategy dictionary.
+    attrs: dict, optional
+        Attributes to add to the dataset.
+    parallel: bool, default=False
+        Whether to perform open and preprocess steps in parallel using
+        `dask.delayed`.
+    zarr_version: int, default=3
+        Zarr version to use.
+    """
+    # === Initialise Asynchronous Object Store === #
+    logging.info("Reading object store credentials from %s", store_credentials_json)
+    obj_store = ObjectStoreS3(anon=False,
+                              asynchronous=True,
+                              store_credentials_json=store_credentials_json
+                              )
+
+    # === Preprocess Data === #
+    ds_filepath = _preprocess_dataset(file=file,
+                                      rechunk=rechunk,
+                                      append_dim=append_dim,
+                                      update_coords=update_coords,
+                                      grid_filepath=grid_filepath,
+                                      attrs=attrs,
+                                      parallel=parallel,
+                                      )
+
+    if variables is None:
+        variables = list(ds_filepath.data_vars)
+    # Consider variables with append dimension only:
+    variables = [var for var in variables if append_dim in ds_filepath[var].dims]
+
+    # === Update Existing Zarr store === #
+    # Write to Zarr store:
+    url = f"s3://{bucket}/{object_prefix}"
+    logging.info(f"Updating Dataset at {url}")
+    _update_zarr_store(data=ds_filepath[variables],
+                       obj_store=obj_store,
+                       url=url,
+                       append_dim=append_dim,
+                       rechunk=rechunk,
+                       version=zarr_version
+                       )
+
+    # Release resources to avoid memory leaks:
+    ds_filepath.close()
+
+
+def update_zarr(
+    file: list[str] | str | xr.Dataset,
+    bucket: str,
+    object_prefix: str,
+    store_credentials_json: str,
+    variables: Optional[list[str]] = None,
+    append_dim: str = "time_counter",
+    grid_filepath: Optional[str] = None,
+    update_coords: Optional[dict] = None,
+    rechunk: Optional[dict] = None,
+    attrs: Optional[dict] = None,
+    client : Optional[Client] = None,
+    dask_config_kwargs: Optional[dict] = None,
+    dask_cluster_kwargs: Optional[dict] = None,
+    zarr_version: int = 3
+    ) -> None:
+    """
+    Update data in existing Zarr store in cloud object
+    storage with option of using dask.
+
+    Parameters
+    ----------
+    file: list | str | xarray.Dataset
+        Regular expression or list of filepaths to netCDF file(s).
+        Users can also pass a single xarray.Dataset directly.
+    bucket: str
+        Name of the bucket in the object store. Bucket names can contain only
+        lowercase letters, numbers, dots (.), and hyphens (-).
+    object_prefix: str
+        Prefix to be added to the object names in the object store.
+    store_credentials_json: str
+        Path to the JSON file containing the object store credentials.
+    variables: list, optional
+        List of variables to send to Zarr stores.
+        If None, all variables will be sent.
+    append_dim: str, default='time_counter'
+        Name of the dimension to append multifile datasets.
+    grid_filepath: str, optional
+        Path to file containing model grid parameter.
+    update_coords: dict, optional
+        Dictionary of coordinate variables to update.
+    rechunk: dict, optional
+        Rechunk strategy dictionary.
+    attrs: dict, optional
+        Attributes to add to the dataset.
+    client: dask.distributed.Client, optional
+        Dask Distributed Client.
+    dask_config_kwargs: Dict[str,str], optional
+        Dask configuration settings passed to dask.config.set().
+    dask_cluster_kwargs: dict, optional
+        Dask cluster configuration settings passed to LocalCluster().
+    zarr_version: int, default=3
+        zarr version to use.
+    """
+    if dask_cluster_kwargs is not None:
+        # === Update Zarr store with Dask === #
+        if dask_config_kwargs is not None:
+            dask.config.set(dask_config_kwargs)
+            logging.info("Updated dask configuration settings.")
+
+        # Create local dask cluster & client:
+        with LocalCluster(**dask_cluster_kwargs) as cluster, Client(cluster) as client:
+            logging.info(f"Created LocalCluster with {dask_cluster_kwargs['n_workers']} workers @ Client: {client.dashboard_link}")
+
+            # Catch UserWarnings when rechunking data:
+            client.register_worker_plugin(CaptureWarningsPlugin())
+
+            _update_zarr(file=file,
+                         bucket=bucket,
+                         object_prefix=object_prefix,
+                         store_credentials_json=store_credentials_json,
+                         variables=variables,
+                         append_dim=append_dim,
+                         grid_filepath=grid_filepath,
+                         update_coords=update_coords,
+                         rechunk=rechunk,
+                         attrs=attrs,
+                         parallel=True,
+                         zarr_version=zarr_version
+                         )
+
+            # --- Shutdown Store & Dask Cluster --- #
+            cluster.close()
+            client.shutdown()
+            logging.info("Dask Cluster has been shutdown.")
+
+    elif client is not None:
+        logging.info(f"Using existing Dask Cluster @ Client: {client.dashboard_link}")
+
+        # Catch UserWarnings when rechunking data:
+        client.register_worker_plugin(CaptureWarningsPlugin())
+
+        _update_zarr(file=file,
+                     bucket=bucket,
+                     object_prefix=object_prefix,
+                     store_credentials_json=store_credentials_json,
+                     variables=variables,
+                     append_dim=append_dim,
+                     grid_filepath=grid_filepath,
+                     update_coords=update_coords,
+                     rechunk=rechunk,
+                     attrs=attrs,
+                     parallel=True,
+                     zarr_version=zarr_version
+                     )
+
+        # --- Shutdown Store & Dask Cluster --- #
+        cluster.close()
+        client.shutdown()
+        logging.info("Existing Dask Cluster has been shutdown.")
+    
+    else:
+        # === Update Zarr store without Dask === #
+        _update_zarr(file=file,
+                     bucket=bucket,
+                     object_prefix=object_prefix,
+                     store_credentials_json=store_credentials_json,
+                     variables=variables,
+                     append_dim=append_dim,
+                     grid_filepath=grid_filepath,
+                     update_coords=update_coords,
+                     rechunk=rechunk,
+                     attrs=attrs,
+                     parallel=False,
+                     zarr_version=zarr_version
+                     )