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

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