Concurrent data downloading and appending to zarr store

[guidocioni]

I'm trying to speed up the download and processing of many netcdf files by running concurrent processes.
The main operations are

1. Download and extract a gzipped netcdf file
2. Read it with xarray, subset over coordinates
3. Append to the time dimension in the zarr store

So I start with

dates = pd.date_range(DATE_START, DATE_END, freq="1D")
# Write the first file to make sure we start with something
download_file(dates[0])
# then proceed with the others
process_map(download_file, dates, max_workers=MAX_WORKERS, chunksize=CHUNKSIZE)

(process_map is from tqdm.contrib.concurrent)

The function download_file basically does (I'm keeping short not because I don't want to share it but because it may be easier to read with less details):

with fsspec.open(url, mode="rb") as f_remote:
    with gzip.GzipFile(fileobj=f_remote) as f_decompressed:
        with io.BytesIO(f_decompressed.read()) as f_buffer:
            ds = xr.open_dataset(f_buffer).squeeze()
            subset = ds.sel(x=slice(x_min, x_max), y=slice(y_min, y_max))
            subset = subset.expand_dims("time").drop_vars('projection')
            if not os.path.exists(OUTPUT_ZARR):
                subset.to_zarr(OUTPUT_ZARR, mode="w", consolidated=True)
            else:  # Append to the existing Zarr dataset
                subset.to_zarr(OUTPUT_ZARR, mode="a", append_dim="time", consolidated=True)

This works when I run the script, but when I try to open the dataset with xarray I get

ValueError: conflicting sizes for dimension 
'time': length 6899 on 'time' and length 6897 on 
{'time': 'IMS_Surface_Values', 'y': 'IMS_Surface_Values', 'x': 'IMS_Surface_Values'}

If I look into the folder I can indeed see that the variable IMS_Surface_Values has chunks going from 0.0.0 to 6896.1.1 while time goes from 0 to 6898.
Is this because I'm writing concurrently or does the problem lie somewhere else?

Surely there must be a smarter way of downloading and appending to a zarr in parallel.

[rabernat]

Thanks for sharing!

It seems like there is a logical problem with what you're trying to do here.

Appending over the time dimension has to be done in the correct order. You want the first file to be first, then append the second file, then the third, etc. so that the Zarr dataset has the correct temporal order. Because it is sequential, this operation can't be parallelized! Your process_map function makes no guarantees about the order in which the files will be processed. If you want to use the append approach, you just have to loop over all of the files sequentially.

However, there is another approach you can use here that is parallelizable. Since you know the shape and coordinates of the final dataset in advance, you can pre-allocate the entire domain and then write to each time slice in parallel. In this mode of operation, you use region instead of append_dim, and the writes can happen in any order. (But you do have to guarantee that the writes are aligned with chunk boundaries, so choose your chunks explicity.) You can see a simple example of this pattern in the Arraylake Docs. We also used it in this blog post.

Finally, when doing complex distributed writes to Zarr, it can be really useful to have a store that supports transactional updates and rollbacks in case something goes wrong. Have a look at our new Icechunk project if you're curious about that.

[guidocioni]

Yep, I had the impression this was the issue...I just didn't find any relevant example online.
I think the region approach looks promising. I do indeed know the dimensions beforehand (even though the x, y dimensions may change if I change the subset region and it would be good to have something dynamic instead than hard coded) so I can try to do that and then write on the region.

[guidocioni]

Cannot believe it was THAT easy.
Thanks to the suggestions of @rabernat I was able to create the following example which does exactly what it's supposed to do by only writing the region specific to a certain time (which is automatically detected by xarray).
The result is that I'm able to download and process about 20 years of daily data in about 7 minutes, and this could be improved by using even more workers (but then the network will probably become the bottleneck).
I'm leaving this here in case someone has the same problem.

import xarray as xr
import fsspec
import gzip
import io
import pyproj
import pandas as pd
from tqdm.contrib.concurrent import process_map
import dask.array as da
import logging

# Define lat/lon box boundaries
lat_min, lat_max = 65, 25  # Replace with your values
lon_min, lon_max = -40, 40  # Replace with your values

# Create a transformer for polar stereographic projection
transformer = pyproj.Transformer.from_crs(
    "EPSG:4326",
    "+proj=stere +lat_0=90 +lat_ts=60 +lon_0=-80 +k=1 +x_0=0 +y_0=0 +a=6378137 +b=6356257 +units=m +no_defs",
    always_xy=True,
)

# Transform the lat/lon boundaries to x/y boundaries
x_min, y_min = transformer.transform(lon_min, lat_min)
x_max, y_max = transformer.transform(lon_max, lat_max)

RESOLUTION="4km"
OUTPUT_ZARR = f"/Users/guidocioni/Downloads/data_{RESOLUTION}.zarr"
DATE_START="2005-01-01"
DATE_END="2024-12-03"
MAX_WORKERS=10
CHUNKSIZE=10

# Helper function to generate the URL for a given date
def get_url_for_date(date):
    if RESOLUTION == "1km":
        version = "v1.3"
    else:
        if date < pd.to_datetime("2014-12-03"):
            version = "v1.2"
        else:
            version = "v1.3"
    return f"https://noaadata.apps.nsidc.org/NOAA/G02156/netcdf/{RESOLUTION}/{date.year}/ims{date.year}{date.day_of_year:03d}_{RESOLUTION}_{version}.nc.gz"


# Parallel version
def main():
    # Process each year
    dates = pd.date_range(DATE_START, DATE_END, freq="1D")

    # Initialize the zarr store    
    url = get_url_for_date(dates[0])
    with fsspec.open(url, mode="rb") as f_remote:
        with gzip.GzipFile(fileobj=f_remote) as f_decompressed:
            with io.BytesIO(f_decompressed.read()) as f_buffer:
                ds = xr.open_dataset(f_buffer).squeeze()
                subset = ds.sel(x=slice(x_min, x_max), y=slice(y_min, y_max)).drop_vars('projection').expand_dims("time")
                # Initialize the Zarr store with the correct dimensions
                shape = (len(dates), subset.sizes['y'], subset.sizes['x'])
                chunks = (1, subset.sizes['y'], subset.sizes['x'])  # Write 1 time step at a time
                template = xr.Dataset(
                    data_vars={
                        "IMS_Surface_Values": (
                            ('time', 'y', 'x'),
                            da.zeros(shape, chunks=chunks, dtype=subset.IMS_Surface_Values.dtype)
                        )
                    },
                    coords={
                        'time': dates,
                        'y': subset['y'],
                        'x': subset['x']
                    }
                )
                template.to_zarr(OUTPUT_ZARR, mode="w", compute=False)
                # Now append first time step that we already downloaded
                subset.to_zarr(OUTPUT_ZARR, region="auto")

    process_map(download_time_step, dates[1:], max_workers=MAX_WORKERS, chunksize=CHUNKSIZE)

def download_time_step(date):
    url = get_url_for_date(date)
    try:
        with fsspec.open(url, mode="rb") as f_remote:
            with gzip.GzipFile(fileobj=f_remote) as f_decompressed:
                with io.BytesIO(f_decompressed.read()) as f_buffer:
                    ds = xr.open_dataset(f_buffer).squeeze()
                    subset = ds.sel(x=slice(x_min, x_max), y=slice(y_min, y_max)).drop_vars('projection').expand_dims("time")
                    subset.to_zarr(OUTPUT_ZARR, region="auto")
    except Exception as e:
        logging.error(
            f"Could not process {url}: {type(e).__name__} at line {e.__traceback__.tb_lineno} of {__file__}: {e}")

if __name__ == "__main__":
    main()

[guidocioni]

Hey @rabernat just curiosity.
Would it be possible to add timesteps to this zarr store in the future by using a similar concurrent approach, or will this not be possible as the time array needs to be defined at the beginning and cannot be modified afterwards?

Of course the serial approach (appending one instant after another sequentially) would still be available but this could become slow in case there are many timesteps to append.

Thanks again

[rabernat]

Yes absolutely, as long as there is only one process which is extending the array at a time. You could also resize the array one time and then add new data using the region approach.

[guidocioni]

Yes, so my question is how do you dynamically resize the time dimension of the zarr store once it has been created and written to disk? :)
Are there any special parameters to define an unlimited dimension?

[rabernat]

AFAIK Xarray cannot do this. You need to use the Zarr API directly, e.g. zarr.Array.resize on all of the variables you want to extend. You will also need to manually write the new time coordinate data. Not particularly user-friendly. We should explore whether Xarray could provide a better API for this as it's a fairly common scenario.

It could be instructive to see how Xarray implements append: https://github.com/pydata/xarray/blob/96e0ff7d70c605a1505ff89a2d62b5c4138b0305/xarray/backends/zarr.py#L1097-L1107