Generating a Virtual Zarr Store

Generating a Virtual Zarr Store#

The objective of this notebook is to learn how to create a virtual Zarr store for a collection of NetCDF files that together make up a complete dataset. To do this, we will use Kerchunk and VirtualiZarr. As these two packages can both create virtual Zarr stores but do it in different ways, we will utilize them both to show how they compare in combination with dask for parallel execution.

import fsspec
import xarray as xr
import ujson
import time
import kerchunk.hdf
import kerchunk.combine
from virtualizarr import open_virtual_dataset
import dask.distributed
from dask_gateway import Gateway
import logging
import os

Kerchunk vs VirtualiZarr#

To begin, let’s explain what a virtual Zarr store even is. A “virtual Zarr store” is a virtual representation of a Zarr store generated by mapping any number of real datasets in individual files (e.g., NetCDF/HDF5, GRIB2, TIFF) together into a single, sliceable dataset via an interface layer. This interface layer, which Kerchunk and VirtualiZarr generate, contains information about the original files (e.g., chunking, compression, data byte location, etc.) needed to efficiently access the data. While this could be done with xarray.open_mfdataset, we don’t want to run this command every time we open the dataset as it can be a slow and expensive process. The reason for this is that xarray.open_mfdataset performs many consistency checks as it runs, and it requires partially opening all of the datasets to get general matadata information on each of the individual files. Therefore, for numerous files, this can have significant overhead, and it would be preferable to just cache these checks and metadata for more performant future reads. This cache (specifically in Zarr format) is what a virtual Zarr store is. Once we have the virtual Zarr store, we can open the combined xarray dataset using xarray.open_dataset for an almost instantaneous read.

Now that we know what a virtual Zarr store is, let’s discuss the differences between Kerchunk and VirtualiZarr and their virtual Zarr stores. At a top level, VirtualiZarr provides almost all of the same features as Kerchunk. The primary difference is that Kerchunk supports non-Zarr-like virtual format, while VirtualiZarr is specifically focused on the Zarr format. Additionally, Kerchunk creates the virtual Zarr store and represents it in memory using json formatting (the format used for Zarr metadata). Alternatively, VirtualiZarr represents the store as array-level abstractions (which can be converted to json format). These abstractions can be cleanly wrapped by xarray for easy use of xarray.concat and xarray.merge commands to combine virtual Zarr stores. A nice table comparing the two packages can be found in the VirtualiZarr FAQs, which shows how the two packages represent virtual Zarr stores and their comparative syntax.

Spin up Dask Cluster#

To run the virtual Zarr creation in parallel, we need to spin up a Dask cluster to schedule the various workers.

os.environ['DASK_DISTRIBUTED__SCHEDULER__WORKER_SATURATION'] = "1.0"

gateway = Gateway()
_options = gateway.cluster_options()
_options.conda_environment='hytest/hytest-rechunking'  ##<< this is the conda environment we use on nebari.
_options.profile = 'Medium Worker'   
cluster = gateway.new_cluster(_options)          ##<< create cluster via the dask gateway
cluster.scale(30)                                ##<< Sets scaling parameters. 
client = cluster.get_client()
client

Client

Client-3cf91031-5684-11f0-813e-ea8215427e9e

Connection method: Cluster object	Cluster type: dask_gateway.GatewayCluster
Dashboard: https://nebari.chs.usgs.gov/gateway/clusters/nebari.672f08a8edee42099224cdeabed00d2d/status

Cluster Info

GatewayCluster

Name: nebari.672f08a8edee42099224cdeabed00d2d
Dashboard: https://nebari.chs.usgs.gov/gateway/clusters/nebari.672f08a8edee42099224cdeabed00d2d/status

Example Comparison#

With our Dask cluster ready, let’s see how Kerchunk and VirtualiZarr can be utilized to generate a vitrual Zarr store. For this example, we will use the same daily gridMET NetCDF data as used in the Writing Chunked File tutorial. Only this time we will use all of the variables not just precipitation. These include:

precipitation,
maximum relative humidity,
minimum relative humidity,
specific humidity,
downward shortwave radiation,
minimum air temperature,
maximum air temperature,
wind direction, and
wind speed.

The data is currently hosted on the HyTEST OSN as a collection NetCDF files. To access the data with both Kerchunk and VirtualiZarr, we will use fsspec to get the list of files that we are wanting to combine into a virtual Zarr store.

First we need to create the file system for accessing the files, and a second one for outputting the virtual Zarr store.

Note

We will exclude the year 2019 for now and use it later to show how to append virtual Zarr stores. Also, we will not use 2020 as it is a partial year with different chunking than the other 40 years, which is currently incompatible with Kerchunk and Virtualizarr.

# These reader options will be needed for VirtualiZarr
# We created them here to show how they fold into fsspec
reader_options = {
    'storage_options': {
        'anon': True, 
        'client_kwargs': {
            'endpoint_url': 'https://usgs.osn.mghpcc.org/'
        }
    }
}

fs = fsspec.filesystem(
    protocol='s3',
    **reader_options['storage_options']
)

fs_local = fsspec.filesystem('')
# Make directories to save the virtual zarr stores
fs_local.mkdirs('virtual_zarr/kerchunk', exist_ok=True)
fs_local.mkdirs('virtual_zarr/virtualizarr', exist_ok=True)

file_glob = fs.glob('s3://mdmf/gdp/netcdf/gridmet/gridmet/*198*.nc')
file_glob = [file for file in file_glob if (('2020' not in file) and ('2019' not in file))]

Now, we are ready to generate the virtual Zarr stores. For both Kerchunk and VirtualiZarr (for now), this consists of two steps:

Convert single original data files into individual virtual Zarr stores,
Combine the individual virtual Zarr stores into a single combined virtual Zarr store.

We will show these two steps seperately and how they are done for each package.

Opening the Virtual Zarr Stores#

To open the virtual Zarr stores, we can use the same method for both stores as we converted to Kerchunk format when saving from VirtualiZarr.

Kerchunk#

t0 = time.time()

ds = xr.open_dataset(
    'virtual_zarr/kerchunk/gridmet.json',
    chunks={},
    engine="kerchunk",
    backend_kwargs={
        "storage_options": {
            "remote_protocol": "s3",
            "remote_options": reader_options['storage_options']
        },
    }
)

kerchunk_read_time = time.time() - t0

ds

<xarray.Dataset> Size: 166GB
Dimensions:                                    (day: 3653, lat: 585, lon: 1386,
                                                crs: 1)
Coordinates:
  * crs                                        (crs) uint16 2B 3
  * day                                        (day) datetime64[ns] 29kB 1980...
  * lat                                        (lat) float64 5kB 49.4 ... 25.07
  * lon                                        (lon) float64 11kB -124.8 ... ...
Data variables:
    air_temperature                            (day, lat, lon) float64 24GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
    precipitation_amount                       (day, lat, lon) float64 24GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
    relative_humidity                          (day, lat, lon) float64 24GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
    specific_humidity                          (day, lat, lon) float64 24GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
    surface_downwelling_shortwave_flux_in_air  (day, lat, lon) float64 24GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
    wind_from_direction                        (day, lat, lon) float64 24GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
    wind_speed                                 (day, lat, lon) float64 24GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
Attributes: (12/19)
    Conventions:                CF-1.6
    author:                     John Abatzoglou - University of Idaho, jabatz...
    coordinate_system:          EPSG:4326
    date:                       02 July 2019
    geospatial_bounds:          POLYGON((-124.7666666333333 49.40000000000000...
    geospatial_bounds_crs:      EPSG:4326
    ...                         ...
    geospatial_lon_units:       decimal_degrees east
    note1:                      The projection information for this file is: ...
    note2:                      Citation: Abatzoglou, J.T., 2013, Development...
    note3:                      Data in slices after last_permanent_slice (1-...
    note4:                      Data in slices after last_provisional_slice (...
    note5:                      Days correspond approximately to calendar day...

VirtualiZarr#

t0 = time.time()

ds = xr.open_dataset(
    'virtual_zarr/virtualizarr/gridmet.json',
    chunks={},
    engine="kerchunk",
    backend_kwargs={
        "storage_options": {
            "remote_protocol": "s3",
            "remote_options": reader_options['storage_options']
        },
    }
)

virtualizarr_read_time = time.time() - t0

ds

<xarray.Dataset> Size: 166GB
Dimensions:                                    (day: 3653, lat: 585, lon: 1386,
                                                crs: 1)
Coordinates:
  * crs                                        (crs) float32 4B 3.0
  * day                                        (day) datetime64[ns] 29kB 1980...
  * lat                                        (lat) float64 5kB 49.4 ... 25.07
  * lon                                        (lon) float64 11kB -124.8 ... ...
Data variables:
    air_temperature                            (day, lat, lon) float64 24GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
    precipitation_amount                       (day, lat, lon) float64 24GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
    relative_humidity                          (day, lat, lon) float64 24GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
    specific_humidity                          (day, lat, lon) float64 24GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
    surface_downwelling_shortwave_flux_in_air  (day, lat, lon) float64 24GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
    wind_from_direction                        (day, lat, lon) float64 24GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
    wind_speed                                 (day, lat, lon) float64 24GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
Attributes: (12/19)
    Conventions:                CF-1.6
    author:                     John Abatzoglou - University of Idaho, jabatz...
    coordinate_system:          EPSG:4326
    date:                       02 July 2019
    geospatial_bounds:          POLYGON((-124.7666666333333 49.40000000000000...
    geospatial_bounds_crs:      EPSG:4326
    ...                         ...
    geospatial_lon_units:       decimal_degrees east
    note1:                      The projection information for this file is: ...
    note2:                      Citation: Abatzoglou, J.T., 2013, Development...
    note3:                      Data in slices after last_permanent_slice (1-...
    note4:                      Data in slices after last_provisional_slice (...
    note5:                      Days correspond approximately to calendar day...

Reading with `xarray.open_mfdataset`#

As a comparison of read times, let’s also compile the dataset using xarray.open_mfdataset in parallel with Dask. This way we can see if we will be saving time in the future by having the compiled virtual Zarr for faster reads.

%%time 
t0 = time.time()

ds = xr.open_mfdataset(
    [fs.open(file) for file in file_glob],
    chunks={},
    parallel=True,
    engine='h5netcdf'
)

open_mfdataset_time = time.time() - t0

ds

CPU times: user 966 ms, sys: 65.1 ms, total: 1.03 s
Wall time: 1min 5s

<xarray.Dataset> Size: 166GB
Dimensions:                                    (day: 3653, lat: 585, lon: 1386,
                                                crs: 1)
Coordinates:
  * lon                                        (lon) float64 11kB -124.8 ... ...
  * lat                                        (lat) float64 5kB 49.4 ... 25.07
  * day                                        (day) datetime64[ns] 29kB 1980...
  * crs                                        (crs) uint16 2B 3
Data variables:
    precipitation_amount                       (day, lat, lon) float64 24GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
    relative_humidity                          (day, lat, lon) float64 24GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
    specific_humidity                          (day, lat, lon) float64 24GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
    surface_downwelling_shortwave_flux_in_air  (day, lat, lon) float64 24GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
    wind_from_direction                        (day, lat, lon) float64 24GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
    air_temperature                            (day, lat, lon) float64 24GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
    wind_speed                                 (day, lat, lon) float64 24GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
Attributes: (12/19)
    geospatial_bounds_crs:      EPSG:4326
    Conventions:                CF-1.6
    geospatial_bounds:          POLYGON((-124.7666666333333 49.40000000000000...
    geospatial_lat_min:         25.066666666666666
    geospatial_lat_max:         49.40000000000000
    geospatial_lon_min:         -124.7666666333333
    ...                         ...
    date:                       02 July 2019
    note1:                      The projection information for this file is: ...
    note2:                      Citation: Abatzoglou, J.T., 2013, Development...
    note3:                      Data in slices after last_permanent_slice (1-...
    note4:                      Data in slices after last_provisional_slice (...
    note5:                      Days correspond approximately to calendar day...

Now, let’s compare the computational times!

print("Kerchunk virtual Zarr creation time: "
      f"{kerchunk_time:.0f}s ({kerchunk_time/60:.1f} min)")
print("VirtualiZarr virtual Zarr creation time: "
      f"{virtualizarr_time:.0f}s ({virtualizarr_time/60:.1f} min)")
print("open_mfdataset dataset creation time: "
      f"{open_mfdataset_time:.0f}s ({open_mfdataset_time/60:.1f} min)")
print(f"Time ratio: Kerchunk to open_mfdataset = {kerchunk_time/open_mfdataset_time:.2f}\n"
      f"            VirtualiZarr to open_mfdataset = {virtualizarr_time/open_mfdataset_time:.2f}\n"
      f"            Kerchunk to VirtualiZarr = {kerchunk_time/virtualizarr_time:.2f}")

Kerchunk virtual Zarr creation time: 225s (3.8 min)
VirtualiZarr virtual Zarr creation time: 133s (2.2 min)
open_mfdataset dataset creation time: 65s (1.1 min)
Time ratio: Kerchunk to open_mfdataset = 3.44
            VirtualiZarr to open_mfdataset = 2.03
            Kerchunk to VirtualiZarr = 1.70

As we can see the Kerchunk and VirtualiZarr take about the same amount of time to create a virtual zarr store. open_mfdataset can create the Dataset on the fly in slightly less time than it takes to create the virtual zarr stores (about 30 seconds less, in this case). However, if you use open_mfdataset, you will have to create this dataset each time you run your code, whereas Kerchunk and Virtualizarr can create the dataset once and then read it much faster on future uses. Looking at read speed after the virtual Zarr store creation:

print("Kerchunk virtual Zarr read time: "
      f"{kerchunk_read_time:.2f}s")
print("VirtualiZarr virtual Zarr read time: "
      f"{virtualizarr_read_time:.2f}s")
print("open_mfdataset dataset read/creation time: "
      f"{open_mfdataset_time:.0f}s ({open_mfdataset_time/60:.1f} min)")
print(f"Time ratio: Kerchunk to open_mfdataset = {kerchunk_read_time/open_mfdataset_time:.3f}\n"
      f"            VirtualiZarr to open_mfdataset = {virtualizarr_read_time/open_mfdataset_time:.3f}\n"
      f"            Kerchunk to VirtualiZarr = {kerchunk_read_time/virtualizarr_read_time:.3f}")

Kerchunk virtual Zarr read time: 0.55s
VirtualiZarr virtual Zarr read time: 0.06s
open_mfdataset dataset read/creation time: 65s (1.1 min)
Time ratio: Kerchunk to open_mfdataset = 0.008
            VirtualiZarr to open_mfdataset = 0.001
            Kerchunk to VirtualiZarr = 9.096

From this, it is very clear that performing more than one read using either the Kerchunk or VirtualiZarr virtual Zarr store is more efficient that reading with open_mfdataset. Additionally, the differences in read times between Kerchunk and Virtualizarr, while appearing drastic, is likely not going to be significant in any workflow.

Appending to Existing Virtual Zarr Store#

As noted when introducing the gridMET data, we did not utilize the 2019 data in order to show how to append it to a virtual Zarr store. The ability to append more data to the virtual Zarr store is highly convienient, as plenty of datasets are continuously updated as new data becomes available. So, let’s appends some data to our virtual Zarr stores we just made.

First, we create the 2019 file glob.

file_glob_2019 = fs.glob('s3://mdmf/gdp/netcdf/gridmet/gridmet/*_2019.nc')

Create New Virtual Zarr for New File#

Next, we need to get our 2019 NetCDFs into a virtual Zarr store.

Kerchunk#

We will do this for Kerchunk the same way we did before, by using kerchunk.hdf.SingleHdf5ToZarr, which translates the content of one HDF5 (NetCDF4) file into Zarr metadata.

tasks = [generate_single_virtual_zarr(file) for file in file_glob_2019]
single_virtual_zarrs_2019 = dask.compute(*tasks)

VirtualiZarr#

And for VirtualiZarr, we will use virtualizarr.open_virtual_dataset.

tasks = [
    dask.delayed(open_virtual_dataset)(
        f's3://{file}',
        indexes={},
        loadable_variables=['day', 'lat', 'lon', 'crs'],
        decode_times=True,
        reader_options=reader_options
    )
    for file in file_glob_2019
]

virtual_datasets_2019 = dask.compute(*tasks)

Double Check New Stores#

Finally, let’s read in the appended stores to make sure that we correctly appended the 2019 data.

ds = xr.open_dataset(
    'virtual_zarr/kerchunk/gridmet_appended.json',
    engine="kerchunk",
    chunks={},
    backend_kwargs={
        "storage_options": {
            "remote_protocol": "s3",
            "remote_options": reader_options['storage_options']
        },
    }
)
ds

<xarray.Dataset> Size: 182GB
Dimensions:                                    (day: 4018, lat: 585, lon: 1386,
                                                crs: 1)
Coordinates:
  * crs                                        (crs) uint16 2B 3
  * day                                        (day) datetime64[ns] 32kB 1980...
  * lat                                        (lat) float64 5kB 49.4 ... 25.07
  * lon                                        (lon) float64 11kB -124.8 ... ...
Data variables:
    air_temperature                            (day, lat, lon) float64 26GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
    precipitation_amount                       (day, lat, lon) float64 26GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
    relative_humidity                          (day, lat, lon) float64 26GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
    specific_humidity                          (day, lat, lon) float64 26GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
    surface_downwelling_shortwave_flux_in_air  (day, lat, lon) float64 26GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
    wind_from_direction                        (day, lat, lon) float64 26GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
    wind_speed                                 (day, lat, lon) float64 26GB dask.array<chunksize=(61, 98, 231), meta=np.ndarray>
Attributes: (12/19)
    Conventions:                CF-1.6
    author:                     John Abatzoglou - University of Idaho, jabatz...
    coordinate_system:          EPSG:4326
    date:                       02 March 2020
    geospatial_bounds:          POLYGON((-124.7666666333333 49.40000000000000...
    geospatial_bounds_crs:      EPSG:4326
    ...                         ...
    geospatial_lon_units:       decimal_degrees east
    note1:                      The projection information for this file is: ...
    note2:                      Citation: Abatzoglou, J.T., 2013, Development...
    note3:                      Data in slices after last_permanent_slice (1-...
    note4:                      Data in slices after last_provisional_slice (...
    note5:                      Days correspond approximately to calendar day...

Nice! The 2019 data is now appended and showing on the day coordinate.

Clean Up#

Rather than deleting the virtual Zarr stores that we created, we will actually keep them for use in future tutorials. However, we will do want to conform with best practices and close our Dask client and cluster.

client.close()
cluster.close()

Generating a Virtual Zarr Store

Contents

Generating a Virtual Zarr Store#

Kerchunk vs VirtualiZarr#

Spin up Dask Cluster#

Client

Cluster Info

GatewayCluster

Example Comparison#

Generate Individual Virtual Zarr Stores#

Kerchunk#

VirtualiZarr#

Combine Individual Virtual Zarr Stores#

Kerchunk#

VirtualiZarr#

Opening the Virtual Zarr Stores#

Kerchunk#

VirtualiZarr#

Reading with `xarray.open_mfdataset`#

Appending to Existing Virtual Zarr Store#

Create New Virtual Zarr for New File#

Kerchunk#

VirtualiZarr#

Append to Existing Store#

Kerchunk#

VirtualiZarr#

Double Check New Stores#

Clean Up#

Generating a Virtual Zarr Store

Contents

Generating a Virtual Zarr Store#

Kerchunk vs VirtualiZarr#

Spin up Dask Cluster#

Client

Cluster Info

GatewayCluster

Example Comparison#

Generate Individual Virtual Zarr Stores#

Kerchunk#

VirtualiZarr#

Combine Individual Virtual Zarr Stores#

Kerchunk#

VirtualiZarr#

Opening the Virtual Zarr Stores#

Kerchunk#

VirtualiZarr#

Reading with xarray.open_mfdataset#

Appending to Existing Virtual Zarr Store#

Create New Virtual Zarr for New File#

Kerchunk#

VirtualiZarr#

Append to Existing Store#

Kerchunk#

VirtualiZarr#

Double Check New Stores#

Clean Up#

Reading with `xarray.open_mfdataset`#