Parallel COG / TIFF Storage

[wietzesuijker]

We're exploring the Planetary Computer Hub for data processing. We have a requirement of a file format that's easy to interactively explore for quality checks etc through software like QGIS. COG / plain geotiff seems excellent for this, where with zarr it's harder to interactively explore. Please correct me if I'm wrong!!
Thanks, Tom Augspurger for suggesting zarr and xcog example. xcog could be part of the solution in combination with a vrt.

Based on rioxarray's documentation (link), we explored parallel COG / tif storage using dask.

Using rxr.open_rasterio() and ds.rio.to_raster(), we routed operations through the Planetary Hub's gateway cluster. However, the dask dashboard doesn't show the last part of the data writing step. Instead, the system seems idle for a while before completing the tiff writing.

For reference, our process resembles this example from the pc docs, with an active dask client on a gateway cluster. In our case, the data comes from a COG in blob storage. We're seeking a format to support dimensions - input as well as output - like {"band": 10, "x": 100_000, "y": 100_000}.

fs = AzureBlobFileSystem(kwargs)
file_url = fs.url(blob_name)
ds = rxr.open_rasterio(file_url, chunks=True, lock=False)

I've observed that data writing is parallelized when using rioxarray with the relevant dask args (dask.array.store https://github.com/corteva/rioxarray/blob/master/rioxarray/raster_writer.py#L294), though the writing to disk is handled by the JupyterHub server node, rather than being distributed across the gateway cluster workers. Our goal is to parallelize the COG chunk writing to Azure Blob Storage.
The code snippet below highlights the buffer processing data, seemingly in memory or over the network. This setup suggests a close dependency between the hub and the gateway cluster. Our intention is to leverage the gateway more effectively, minimizing potential network i/o costs and dependency on the hub.
Relevant code snippet from the example from the pc docs:

import io

with io.BytesIO() as buffer:
    ndvi.rio.to_raster(buffer, driver="COG")
    buffer.seek(0)
    blob_client = container_client.get_blob_client("ndvi-wb.tif")
    blob_client.upload_blob(buffer, overwrite=True)

While investigating, I found the "put-block-list" method here. However, I'm not sure how relevant this is for the use case above.

Any advice?

[TomAugspurger]

We're seeking a format to support dimensions - input as well as output - like {"band": 10, "x": 100_000, "y": 100_000}.

Can you clarify your desired output file(s)? Do you want a single COG or many (one per chunk)?

[wietzesuijker]

Preferably one file, however an alternative that's easily explorable in e.g. QGIS is fine as well (such as the vrt I mentioned) . The further pipeline uses dask and xarray, so this is mostly relevant for interactive quality checks on intermediate data.

[TomAugspurger]

Thanks.

It looks like for .rio.to_raster to work, it needs to lock around each write, with something like lock.distributed.Lock(). However, that seems to directly contradict your desire here:

Our goal is to parallelize the COG chunk writing to Azure Blob Storage.

I suspect that you can't have both parallel writes and a single output file.

If you do want parallel writes, I'd suggest using Zarr.

If you want single files, then perhaps you can parallelize things at some other level of your workload?

[wietzesuijker]

Thanks. Zarr could work if there's a neat way to interactively explore the data in QGIS / in combination with separate reference data (e.g. hi-res google earth or planet data). Would you know a way to achieve that?
*Just came across this library. We'll give it a spin https://github.com/pyvista/pyvista-xarray

[TomAugspurger]

I think GDAL (and so QGIS) can read Zarr these days, but I'm not sure how well things work. I don't that Zarr has a standardized and broadly used version of overviews yet.

[wietzesuijker]

Thanks, I don't think QGIS supports it intuitively with the latest version of GDAL. Added an issue to the qgis repo as this seems beyond PC.

[wietzesuijker]

Thanks.

It looks like for .rio.to_raster to work, it needs to lock around each write, with something like lock.distributed.Lock(). However, that seems to directly contradict your desire here:

Our goal is to parallelize the COG chunk writing to Azure Blob Storage.

I suspect that you can't have both parallel writes and a single output file.

If you do want parallel writes, I'd suggest using Zarr.

If you want single files, then perhaps you can parallelize things at some other level of your workload?

It's fine to use the lock.distributed.Lock(), as that should be fairly fast to pass on between the different jobs in the queue. Generally, I expect the most time-consuming part is the actual compute to memory, before writing to disk (/ streaming to blob storage).
However, the distributed writing to azure is currently not supported, as rioxarray expects a str file path (here) to write to disk instead of a memory buffer.