Split kerchunk reader up

virtualizarr/backend.py

@@ -1,26 +1,18 @@
-import os
 import warnings
-from collections.abc import Iterable, Mapping, MutableMapping
+from collections.abc import Iterable, Mapping
 from enum import Enum, auto
-from io import BufferedIOBase
+from pathlib import Path
 from typing import (
     Any,
-    Hashable,
     Optional,
-    cast,
 )
 
-import xarray as xr
-from xarray.backends import AbstractDataStore, BackendArray
-from xarray.core.indexes import Index, PandasIndex
-from xarray.core.variable import IndexVariable
+from xarray import Dataset
+from xarray.core.indexes import Index
 
 from virtualizarr.manifests import ManifestArray
-from virtualizarr.types.kerchunk import KerchunkStoreRefs
 from virtualizarr.utils import _FsspecFSFromFilepath
 
-XArrayOpenT = str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore
-
 
 class AutoName(Enum):
     # Recommended by official Python docs for auto naming:
@@ -43,10 +35,49 @@ class FileType(AutoName):
     kerchunk = auto()
 
 
-class ManifestBackendArray(ManifestArray, BackendArray):
-    """Using this prevents xarray from wrapping the KerchunkArray in ExplicitIndexingAdapter etc."""
+def automatically_determine_filetype(
+    *,
+    filepath: str,
+    reader_options: Optional[dict[str, Any]] = {},
+) -> FileType:
+    """
+    Attempt to automatically infer the correct reader for this filetype.
 
-    ...
+    Uses magic bytes and file / directory suffixes.
+    """
+
+    # TODO this should ideally handle every filetype that we have a reader for, not just kerchunk
+
+    # TODO how do we handle kerchunk json / parquet here?
+    if Path(filepath).suffix == ".zarr":
+        # TODO we could imagine opening an existing zarr store, concatenating it, and writing a new virtual one...
+        raise NotImplementedError()
+
+    # Read magic bytes from local or remote file
+    fpath = _FsspecFSFromFilepath(
+        filepath=filepath, reader_options=reader_options
+    ).open_file()
+    magic_bytes = fpath.read(8)
+    fpath.close()
+
+    if magic_bytes.startswith(b"CDF"):
+        filetype = FileType.netcdf3
+    elif magic_bytes.startswith(b"\x0e\x03\x13\x01"):
+        raise NotImplementedError("HDF4 formatted files not supported")
+    elif magic_bytes.startswith(b"\x89HDF"):
+        filetype = FileType.hdf5
+    elif magic_bytes.startswith(b"GRIB"):
+        filetype = FileType.grib
+    elif magic_bytes.startswith(b"II*"):
+        filetype = FileType.tiff
+    elif magic_bytes.startswith(b"SIMPLE"):
+        filetype = FileType.fits
+    else:
+        raise NotImplementedError(
+            f"Unrecognised file based on header bytes: {magic_bytes}"
+        )
+
+    return filetype
 
 
 def open_virtual_dataset(
@@ -61,15 +92,14 @@ def open_virtual_dataset(
     indexes: Mapping[str, Index] | None = None,
     virtual_array_class=ManifestArray,
     reader_options: Optional[dict] = None,
-) -> xr.Dataset:
+) -> Dataset:
     """
     Open a file or store as an xarray Dataset wrapping virtualized zarr arrays.
 
     No data variables will be loaded unless specified in the ``loadable_variables`` kwarg (in which case they will be xarray lazily indexed arrays).
 
     Xarray indexes can optionally be created (the default behaviour). To avoid creating any xarray indexes pass ``indexes={}``.
 
-
     Parameters
     ----------
     filepath : str, default None
@@ -112,10 +142,6 @@ def open_virtual_dataset(
             stacklevel=2,
         )
 
-    loadable_vars: dict[str, xr.Variable]
-    virtual_vars: dict[str, xr.Variable]
-    vars: dict[str, xr.Variable]
-
     if drop_variables is None:
         drop_variables = []
     elif isinstance(drop_variables, str):
@@ -135,194 +161,95 @@ def open_virtual_dataset(
     if virtual_array_class is not ManifestArray:
         raise NotImplementedError()
 
-    # if filetype is user defined, convert to FileType
+    if reader_options is None:
+        reader_options = {}
 
     if filetype is not None:
+        # if filetype is user defined, convert to FileType
         filetype = FileType(filetype)
+    else:
+        filetype = automatically_determine_filetype(
+            filepath=filepath, reader_options=reader_options
+        )
 
-    if filetype == FileType.kerchunk:
-        from virtualizarr.readers.kerchunk import dataset_from_kerchunk_refs
-
-        fs = _FsspecFSFromFilepath(filepath=filepath, reader_options=reader_options)
-
-        # The kerchunk .parquet storage format isn't actually a parquet, but a directory that contains named parquets for each group/variable.
-        if fs.filepath.endswith("ref.parquet"):
-            from fsspec.implementations.reference import LazyReferenceMapper
-
-            lrm = LazyReferenceMapper(filepath, fs.fs)
-
-            # build reference dict from KV pairs in LazyReferenceMapper
-            # is there a better / more preformant way to extract this?
-            array_refs = {k: lrm[k] for k in lrm.keys()}
-
-            full_reference = {"refs": array_refs}
-
-            return dataset_from_kerchunk_refs(KerchunkStoreRefs(full_reference))
-
-        # JSON has no magic bytes, but the Kerchunk version 1 spec starts with 'version':
-        # https://fsspec.github.io/kerchunk/spec.html
-        elif fs.read_bytes(9).startswith(b'{"version'):
-            import ujson
+    # TODO define these through a mapping to registered pluggable entrypoints instead
+    match filetype.name.lower():
+        case "kerchunk":
+            from virtualizarr.readers.kerchunk import KerchunkVirtualBackend
 
-            with fs.open_file() as of:
-                refs = ujson.load(of)
+            return KerchunkVirtualBackend.open_virtual_dataset(filepath, reader_options)
 
-            return dataset_from_kerchunk_refs(KerchunkStoreRefs(refs))
+        case "zarr_v3":
+            from virtualizarr.readers.zarr_v3 import ZarrV3VirtualBackend
 
-        else:
-            raise ValueError(
-                "The input Kerchunk reference did not seem to be in Kerchunk's JSON or Parquet spec: https://fsspec.github.io/kerchunk/spec.html. The Kerchunk format autodetection is quite flaky, so if your reference matches the Kerchunk spec feel free to open an issue: https://github.com/zarr-developers/VirtualiZarr/issues"
+            return ZarrV3VirtualBackend.open_virtual_dataset(
+                filepath=filepath,
+                drop_variables=drop_variables,
+                indexes=indexes,
             )
 
-    if filetype == FileType.zarr_v3:
-        # TODO is there a neat way of auto-detecting this?
-        from virtualizarr.readers.zarr import open_virtual_dataset_from_v3_store
+        case "dmrpp":
+            from virtualizarr.readers.dmrpp import DMRPPVirtualBackend
 
-        return open_virtual_dataset_from_v3_store(
-            storepath=filepath, drop_variables=drop_variables, indexes=indexes
-        )
-    elif filetype == FileType.dmrpp:
-        from virtualizarr.readers.dmrpp import DMRParser
-
-        if loadable_variables != [] or indexes is None:
-            raise NotImplementedError(
-                "Specifying `loadable_variables` or auto-creating indexes with `indexes=None` is not supported for dmrpp files."
+            return DMRPPVirtualBackend.open_virtual_dataset(
+                filepath,
+                drop_variables=drop_variables,
+                loadable_variables=loadable_variables,
+                indexes=indexes,
+                reader_options=reader_options,
             )
 
-        fpath = _FsspecFSFromFilepath(
-            filepath=filepath, reader_options=reader_options
-        ).open_file()
-        parser = DMRParser(fpath.read(), data_filepath=filepath.strip(".dmrpp"))
-        vds = parser.parse_dataset()
-        vds.drop_vars(drop_variables)
-        return vds
-    else:
-        # we currently read every other filetype using kerchunks various file format backends
-        from virtualizarr.readers.kerchunk import (
-            fully_decode_arr_refs,
-            read_kerchunk_references_from_file,
-            virtual_vars_from_kerchunk_refs,
-        )
-
-        if reader_options is None:
-            reader_options = {}
+        case "netcdf3":
+            from virtualizarr.readers.netcdf3 import NetCDF3VirtualBackend
 
-        # this is the only place we actually always need to use kerchunk directly
-        # TODO avoid even reading byte ranges for variables that will be dropped later anyway?
-        vds_refs = read_kerchunk_references_from_file(
-            filepath=filepath,
-            filetype=filetype,
-            group=group,
-            reader_options=reader_options,
-        )
-        virtual_vars = virtual_vars_from_kerchunk_refs(
-            vds_refs,
-            drop_variables=drop_variables + loadable_variables,
-            virtual_array_class=virtual_array_class,
-        )
-        ds_attrs = fully_decode_arr_refs(vds_refs["refs"]).get(".zattrs", {})
-        coord_names = ds_attrs.pop("coordinates", [])
-
-        if indexes is None or len(loadable_variables) > 0:
-            # TODO we are reading a bunch of stuff we know we won't need here, e.g. all of the data variables...
-            # TODO it would also be nice if we could somehow consolidate this with the reading of the kerchunk references
-            # TODO really we probably want a dedicated xarray backend that iterates over all variables only once
-            fpath = _FsspecFSFromFilepath(
-                filepath=filepath, reader_options=reader_options
-            ).open_file()
-
-            # fpath can be `Any` thanks to fsspec.filesystem(...).open() returning Any.
-            # We'll (hopefully safely) cast it to what xarray is expecting, but this might let errors through.
-
-            ds = xr.open_dataset(
-                cast(XArrayOpenT, fpath),
-                drop_variables=drop_variables,
+            return NetCDF3VirtualBackend.open_virtual_dataset(
+                filepath,
                 group=group,
-                decode_times=decode_times,
+                drop_variables=drop_variables,
+                loadable_variables=loadable_variables,
+                indexes=indexes,
+                reader_options=reader_options,
             )
 
-            if indexes is None:
-                warnings.warn(
-                    "Specifying `indexes=None` will create in-memory pandas indexes for each 1D coordinate, but concatenation of ManifestArrays backed by pandas indexes is not yet supported (see issue #18)."
-                    "You almost certainly want to pass `indexes={}` to `open_virtual_dataset` instead."
-                )
-
-                # add default indexes by reading data from file
-                indexes = {name: index for name, index in ds.xindexes.items()}
-            elif indexes != {}:
-                # TODO allow manual specification of index objects
-                raise NotImplementedError()
-            else:
-                indexes = dict(**indexes)  # for type hinting: to allow mutation
-
-            loadable_vars = {
-                str(name): var
-                for name, var in ds.variables.items()
-                if name in loadable_variables
-            }
-
-            # if we only read the indexes we can just close the file right away as nothing is lazy
-            if loadable_vars == {}:
-                ds.close()
-        else:
-            loadable_vars = {}
-            indexes = {}
-
-        vars = {**virtual_vars, **loadable_vars}
-
-        data_vars, coords = separate_coords(vars, indexes, coord_names)
-
-        vds = xr.Dataset(
-            data_vars,
-            coords=coords,
-            # indexes={},  # TODO should be added in a later version of xarray
-            attrs=ds_attrs,
-        )
-
-        # TODO we should probably also use vds.set_close() to tell xarray how to close the file we opened
+        case "hdf5" | "netcdf4":
+            from virtualizarr.readers.hdf5 import HDF5VirtualBackend
 
-        return vds
+            return HDF5VirtualBackend.open_virtual_dataset(
+                filepath,
+                group=group,
+                drop_variables=drop_variables,
+                loadable_variables=loadable_variables,
+                indexes=indexes,
+                reader_options=reader_options,
+            )
 
+        case "grib":
+            # TODO Grib files should be handled as a DataTree object
+            # see https://github.com/TomNicholas/VirtualiZarr/issues/11
+            raise NotImplementedError(f"Unsupported file type: {filetype}")
 
-def separate_coords(
-    vars: Mapping[str, xr.Variable],
-    indexes: MutableMapping[str, Index],
-    coord_names: Iterable[str] | None = None,
-) -> tuple[dict[str, xr.Variable], xr.Coordinates]:
-    """
-    Try to generate a set of coordinates that won't cause xarray to automatically build a pandas.Index for the 1D coordinates.
+        case "tiff":
+            from virtualizarr.readers.tiff import TIFFVirtualBackend
 
-    Currently requires this function as a workaround unless xarray PR #8124 is merged.
+            return TIFFVirtualBackend.open_virtual_dataset(
+                filepath,
+                group=group,
+                drop_variables=drop_variables,
+                loadable_variables=loadable_variables,
+                indexes=indexes,
+                reader_options=reader_options,
+            )
 
-    Will also preserve any loaded variables and indexes it is passed.
-    """
+        case "fits":
+            from virtualizarr.readers.fits import FITSVirtualBackend
 
-    if coord_names is None:
-        coord_names = []
-
-    # split data and coordinate variables (promote dimension coordinates)
-    data_vars = {}
-    coord_vars: dict[
-        str, tuple[Hashable, Any, dict[Any, Any], dict[Any, Any]] | xr.Variable
-    ] = {}
-    for name, var in vars.items():
-        if name in coord_names or var.dims == (name,):
-            # use workaround to avoid creating IndexVariables described here https://github.com/pydata/xarray/pull/8107#discussion_r1311214263
-            if len(var.dims) == 1:
-                dim1d, *_ = var.dims
-                coord_vars[name] = (dim1d, var.data, var.attrs, var.encoding)
-
-                if isinstance(var, IndexVariable):
-                    # unless variable actually already is a loaded IndexVariable,
-                    # in which case we need to keep it and add the corresponding indexes explicitly
-                    coord_vars[str(name)] = var
-                    # TODO this seems suspect - will it handle datetimes?
-                    indexes[name] = PandasIndex(var, dim1d)
-            else:
-                coord_vars[name] = var
-        else:
-            data_vars[name] = var
-
-    coords = xr.Coordinates(coord_vars, indexes=indexes)
-
-    return data_vars, coords
+            return FITSVirtualBackend.open_virtual_dataset(
+                filepath,
+                group=group,
+                drop_variables=drop_variables,
+                loadable_variables=loadable_variables,
+                indexes=indexes,
+                reader_options=reader_options,
+            )
+        case _:
+            raise NotImplementedError(f"Unsupported file type: {filetype.name}")