Use kerchunk to access netcdf files

[gmaze]

One more time I came across kerchunk during ADMT25 but I finally understood what it is about !
Thanks @DocOtak !

So I gave it a try and produced the following ArgoKerchunker class, that could be used like this:

ncfile = "s3://argo-gdac-sandbox/pub/dac/coriolis/6903090/6903090_prof.nc"

# define the store where kerchunk json files/data will be stored
ak = ArgoKerchunker(store='memory')  # default
ak = ArgoKerchunker(store='local')

ak.supported(ncfile)
ak.translate(ncfile)
ak.pprint(ncfile)

ak.open_dataset(ncfile)  # Return a xarray dataset open using the kerchunk zarr version of the netcdf file

I still need to figure how to properly integrate this into argopy and if it is worth it, compared to working on more Analysis-ready and cloud-optimized (ARCO) format.

For instance, the resulting dataset has variables not casted properly.
While argopy already provides the argo engine to open properly (ie cast all variables correctly) a netcdf file, I wonder how to implement casting lazily with zarr.

Don't hesitate to shoot your opinion below !

Argo kerchunk helper class

Details

import fsspec
import xarray as xr
from typing import List, Union
from pathlib import Path
from fsspec.core import split_protocol
import json
import dask
from kerchunk.hdf import SingleHdf5ToZarr
from kerchunk.netCDF3 import NetCDF3ToZarr
from argopy.stores import memorystore, filestore
from argopy.utils import to_list


class ArgoKerchunker:
    """Argo netcdf file kerchunk helper"""

    def __init__(
        self,
        store: Union["memory", "local"] = "memory",
        inline_threshold: int = 0,
        max_chunk_size: int = 0,
    ):
        # File system to save kerchunk json files
        if store == "memory":
            self.fs = memorystore()
        elif store == "local":
            self.fs = filestore()

        # List of processed files register:
        self.kerchunk_references = {}

        self.inline_threshold = inline_threshold
        """inline_threshold: int
        Byte size below which an array will be embedded in the output. Use 0 to disable inlining.
        """

        self.max_chunk_size = max_chunk_size
        """
        max_chunk_size: int
        How big a chunk can be before triggering subchunking. If 0, there is no
        subchunking, and there is never subchunking for coordinate/dimension arrays.
        E.g., if an array contains 10,000bytes, and this value is 6000, there will
        be two output chunks, split on the biggest available dimension. [TBC]
        """

    def _tojson(self, ncfile: Union[str, Path]):
        ncfile = str(ncfile)

        magic = self._magic(ncfile)[0:3]
        if magic == b"CDF":
            chunks = NetCDF3ToZarr(
                ncfile,
                inline_threshold=self.inline_threshold,
                max_chunk_size=self.max_chunk_size,
            )
        elif magic == b"\x89HD":
            chunks = SingleHdf5ToZarr(
                ncfile,
                inline_threshold=self.inline_threshold,
                max_chunk_size=self.max_chunk_size,
            )

        zjs = chunks.translate()

        single_kerchunk_reference = Path(ncfile).name.replace(".nc", ".json")

        with self.fs.open(single_kerchunk_reference, "wb") as f:
            f.write(json.dumps(zjs).encode())

        return ncfile, single_kerchunk_reference, zjs

    def translate(self, ncfiles: List):
        """Compute kerchunks for one or a list of netcdf files"""
        ncfiles = to_list(ncfiles)
        for i, f in enumerate(ncfiles):
            ncfiles[i] = str(f)

        # results = []
        # for ncfile in ncfiles:
        #     results.append(self._tojson(ncfile))

        tasks = [dask.delayed(self._tojson)(u) for u in ncfiles]
        results = dask.compute(tasks)

        for result in results[0]:
            ncfile, single_kerchunk_reference, _ = result
            self.kerchunk_references.update({ncfile: single_kerchunk_reference})

        return results

    def pprint(self, ncfile: Union[str, Path], params: List[str] = None):
        """Pretty print kerchunk json data for a netcdf file"""
        ncfile = str(ncfile)

        if ncfile not in self.kerchunk_references:
            self.translate(ncfile)

        params = to_list(params) if params is not None else []

        # Read and load the kerchunk JSON file:
        single_kerchunk_reference = self.kerchunk_references[ncfile]
        with self.fs.open(single_kerchunk_reference, "r") as file:
            data = json.load(file)

        # Pretty print JSON data
        keys_to_select = [".zgroup", ".zattrs", ".zmetadata"]
        data_to_print = {}
        for key, value in data["refs"].items():
            if key in keys_to_select:
                if isinstance(value, str):
                    data_to_print[key] = json.loads(value)
                else:
                    data_to_print[key] = value
            for p in params:
                if p == key.split("/")[0]:
                    if isinstance(value, str):
                        data_to_print[key] = json.loads(value)
                    else:
                        data_to_print[key] = value

        print(json.dumps(data_to_print, indent=4))

    def open_dataset(self, ncfile: Union[str, Path]) -> xr.Dataset:
        """Open a netcdf Argo file with kerchunk"""
        ncfile = str(ncfile)

        if ncfile not in self.kerchunk_references:
            self.translate(ncfile)

        single_kerchunk_reference = self.kerchunk_references[ncfile]

        remote_protocol = split_protocol(ncfile)[0]

        fs_single = fsspec.filesystem(
            "reference", fo=single_kerchunk_reference, remote_protocol=remote_protocol
        )

        single_map = fs_single.get_mapper("")

        return xr.open_dataset(
            single_map, engine="zarr", backend_kwargs={"consolidated": False}
        )

    def _magic(self, ncfile: Union[str, Path]) -> str:
        fs = fsspec.filesystem(split_protocol(str(ncfile))[0])
        try:
            return fs.open(str(ncfile)).read(4)
        except:
            return None

    def supported(self, ncfile: Union[str, Path]) -> bool:
        """Check if a netcdf file can be accessed through byte ranges

        Argo GDAC supporting byte ranges:
        - s3://argo-gdac-sandbox/pub/idx
        - https://argo-gdac-sandbox.s3-eu-west-3.amazonaws.com/pub/idx

        Not supporting:
        - https://data-argo.ifremer.fr
        """
        fs = fsspec.filesystem(split_protocol(str(ncfile))[0])
        try:
            fs.open(str(ncfile)).read(4)
            return True
        except:
            return False

[ocefpaf]

BTW, this is built on top of kerchunk and may be useful too:

https://virtualizarr.readthedocs.io/en/latest/

[gmaze]

thanks @ocefpaf for the ref
this looks great very for nicely formatted data
But I'm afraid that for current Argo netcdf files, where N-D arrays never have the same size, it may be difficult to use such a high level library
We don't "have a bunch of legacy files which together tile along a dimension to form a large dataset"

[gmaze]

I added a updated version of the above ArgoKerchunker into #385
Lastest version can thus be seen at:

argopy/argopy/stores/kerchunker.py

Line 41 in 0f01c3c

class ArgoKerchunker:

        >>> ncfile = "s3://argo-gdac-sandbox/pub/dac/coriolis/6903090/6903090_prof.nc"

        >>> ak = ArgoKerchunker(store='memory')  # default
        >>> ak = ArgoKerchunker(store='local', root='.')
        >>> ak = ArgoKerchunker(store='local', root='/kerchunk_data_folder')  # Custom local storage folder
        >>> ak = ArgoKerchunker(store=fsspec.filesystem('dir', path='s3://.../kerchunk_data_folder/', target_protocol='s3'))

        >>> ak.supported(ncfile)
        >>> ak.translate(ncfile)  # takes 1 or a list of uris, requires the kerchunk library
        >>> ak.to_kerchunk(ncfile)  # Take 1 netcdf file uri, return kerchunk json data (translate or load from store)
        >>> ak.pprint(ncfile)

        >>> ak.open_dataset(ncfile)  #  Return xarray dataset of the netcdf file using zarr engine from kerchunk data

[gmaze]

The ArgoKerchunker class has evolved a lot: most notably I removed the open_dataset that was in fact a nightmare to handle because it was placing the new class both downstream and upstream of an argopy store, which was not satisfactory.

So I adopted an upstream positioning whereby to open a dataset lazily, users would necessarily do it through one of the argopy stores open_dataset method. In doing so the ArgoKerchunker class could easily inherit of the most appropriate file system.

The remaining downstream use case, is when translating files without the need to up; in the scenario of producing/sharing references for less experience users.

[gmaze]

but this remains highly experimental and could change anytime

[gmaze]

In fact, the performance overhead of translating a remote netcdf file to a zarr reference is not worth it for mono/multi profile files that are only a few MB.

On the other hand, this could come in handy with larger netcdf files from the geo folder, gathering daily data from all floats.

[gmaze]

Implemented and release in v1.1.0