Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Split kerchunk reader up #261

Merged
merged 29 commits into from
Oct 19, 2024
Merged
Show file tree
Hide file tree
Changes from 21 commits
Commits
Show all changes
29 commits
Select commit Hold shift + click to select a range
89ff49a
standardize zarr v3 and dmrpp readers behind dedicated open_virtual_d…
TomNicholas Oct 17, 2024
8d6c42a
refactor hdf5 reader behind open_virtual_dataset function
TomNicholas Oct 17, 2024
eb7444e
refactor netcdf3
TomNicholas Oct 17, 2024
bb39907
refactor tiff
TomNicholas Oct 17, 2024
97fc588
refactor fits
TomNicholas Oct 17, 2024
2e197e2
refactored so create VirtualBackends
TomNicholas Oct 18, 2024
f29d2ff
restore backend.py, but keep readers/common.py
TomNicholas Oct 18, 2024
5a8b18e
oops I deleted a file
TomNicholas Oct 18, 2024
84330f0
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Oct 18, 2024
0e2fa71
standardize open_virtual_dataset method signature, and raise NotImple…
TomNicholas Oct 18, 2024
996d81a
fix bug with zarr reader
TomNicholas Oct 18, 2024
bf71ae3
remove todo
TomNicholas Oct 18, 2024
79477a9
Merge branch 'split_kerchunk_reader' of https://github.com/TomNichola…
TomNicholas Oct 18, 2024
fc2f3bc
make open_virtual_dataset a staticmethod
TomNicholas Oct 18, 2024
d955e1a
try to fix mypy error about importing DataTree from versions of xarra…
TomNicholas Oct 18, 2024
4c5a2bb
mypy
TomNicholas Oct 18, 2024
e592933
sanitize drop_variables and loadable_variables
TomNicholas Oct 18, 2024
6a2179e
implement drop_variables for kerchunk reader
TomNicholas Oct 18, 2024
74a6b6c
sanitize drmpp args
TomNicholas Oct 18, 2024
6bafd5b
pass all arguments to kerchunk reader
TomNicholas Oct 18, 2024
b41e5d8
coerce kerchunk refs to our types
TomNicholas Oct 18, 2024
bf78b84
make sure all readers are passed the same set of args
TomNicholas Oct 18, 2024
0ae7437
Merge branch 'main' into split_kerchunk_reader
TomNicholas Oct 18, 2024
180a0fd
fix bad merge, and refactor determine_chunk_grid_shape a bit
TomNicholas Oct 18, 2024
8b987c6
Merge branch 'split_kerchunk_reader' of https://github.com/TomNichola…
TomNicholas Oct 18, 2024
edf0372
ensure decode_times is passed to each reader
TomNicholas Oct 19, 2024
55d152d
remove match case statement in favour of mapping
TomNicholas Oct 19, 2024
f6c75da
ensure optional dependencies aren't imported
TomNicholas Oct 19, 2024
7995b1c
release note
TomNicholas Oct 19, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
334 changes: 128 additions & 206 deletions virtualizarr/backend.py

Large diffs are not rendered by default.

13 changes: 8 additions & 5 deletions virtualizarr/manifests/array.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,13 @@

import numpy as np

from ..types.kerchunk import KerchunkArrRefs
from ..zarr import ZArray
from .array_api import MANIFESTARRAY_HANDLED_ARRAY_FUNCTIONS, _isnan
from .manifest import ChunkManifest
from virtualizarr.manifests.array_api import (
MANIFESTARRAY_HANDLED_ARRAY_FUNCTIONS,
_isnan,
)
from virtualizarr.manifests.manifest import ChunkManifest
from virtualizarr.types.kerchunk import KerchunkArrRefs
from virtualizarr.zarr import ZArray


class ManifestArray:
Expand Down Expand Up @@ -61,7 +64,7 @@ def __init__(

@classmethod
def _from_kerchunk_refs(cls, arr_refs: KerchunkArrRefs) -> "ManifestArray":
from virtualizarr.readers.kerchunk import (
from virtualizarr.translators.kerchunk import (
fully_decode_arr_refs,
parse_array_refs,
)
Expand Down
195 changes: 195 additions & 0 deletions virtualizarr/readers/common.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,195 @@
import os
import warnings
from abc import ABC
from collections.abc import Iterable, Mapping, MutableMapping
from io import BufferedIOBase
from typing import (
TYPE_CHECKING,
Any,
Hashable,
Optional,
cast,
)

import xarray as xr
from xarray import Dataset
from xarray.backends import AbstractDataStore, BackendArray
from xarray.core.indexes import Index, PandasIndex
from xarray.core.variable import IndexVariable, Variable

from virtualizarr.manifests import ManifestArray
from virtualizarr.utils import _FsspecFSFromFilepath

XArrayOpenT = str | os.PathLike[Any] | BufferedIOBase | AbstractDataStore

if TYPE_CHECKING:
try:
from xarray import DataTree # type: ignore[attr-defined]
except ImportError:
DataTree = Any


class ManifestBackendArray(ManifestArray, BackendArray):
"""Using this prevents xarray from wrapping the KerchunkArray in ExplicitIndexingAdapter etc."""

...


def open_loadable_vars_and_indexes(
filepath: str,
loadable_variables,
reader_options,
drop_variables,
indexes,
group,
decode_times,
) -> tuple[Mapping[str, Variable], Mapping[str, Index]]:
"""
Open selected variables and indexes using xarray.

Relies on xr.open_dataset and its auto-detection of filetypes to find the correct installed backend.
"""

# TODO get rid of this if?
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,
group=group,
decode_times=decode_times,
)

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

# TODO we should drop these earlier by using drop_variables
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 = {}

return loadable_vars, indexes


def construct_virtual_dataset(
virtual_vars,
loadable_vars,
indexes,
coord_names,
attrs,
) -> Dataset:
"""Construct a virtual Datset from consistuent parts."""

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=attrs,
)

# TODO we should probably also use vds.set_close() to tell xarray how to close the file we opened

return vds


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.

Currently requires this function as a workaround unless xarray PR #8124 is merged.

Will also preserve any loaded variables and indexes it is passed.
"""

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


class VirtualBackend(ABC):
@staticmethod
def open_virtual_dataset(
Comment on lines +172 to +174
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure that a @staticmethod is the best way to do this but for now it's just an internal implementation detail. The important thing is that the open_virtual_dataset function signature gets standardized by this approach.

filepath: str,
group: str | None = None,
drop_variables: Iterable[str] | None = None,
loadable_variables: Iterable[str] | None = None,
decode_times: bool | None = None,
indexes: Mapping[str, Index] | None = None,
reader_options: Optional[dict] = None,
) -> Dataset:
raise NotImplementedError()

@staticmethod
def open_virtual_datatree(
path: str,
group: str | None = None,
drop_variables: Iterable[str] | None = None,
loadable_variables: Iterable[str] | None = None,
decode_times: bool | None = None,
indexes: Mapping[str, Index] | None = None,
reader_options: Optional[dict] = None,
) -> "DataTree":
raise NotImplementedError()
63 changes: 49 additions & 14 deletions virtualizarr/readers/dmrpp.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,18 +2,55 @@
import warnings
from collections import defaultdict
from collections.abc import Mapping
from typing import Any, Optional
from typing import Any, Iterable, Optional
from xml.etree import ElementTree as ET

import numpy as np
import xarray as xr
from xarray import Coordinates, Dataset
from xarray.core.indexes import Index
from xarray.core.variable import Variable

from virtualizarr.manifests import ChunkManifest, ManifestArray
from virtualizarr.readers.common import VirtualBackend
from virtualizarr.types import ChunkKey
from virtualizarr.utils import _FsspecFSFromFilepath, check_for_collisions
from virtualizarr.zarr import ZArray


class DMRPPVirtualBackend(VirtualBackend):
@staticmethod
def open_virtual_dataset(
filepath: str,
group: str | None = None,
drop_variables: Iterable[str] | None = None,
loadable_variables: Iterable[str] | None = None,
decode_times: bool | None = None,
indexes: Mapping[str, Index] | None = None,
reader_options: Optional[dict] = None,
) -> Dataset:
loadable_variables, drop_variables = check_for_collisions(
drop_variables=drop_variables,
loadable_variables=loadable_variables,
)

if loadable_variables != [] or decode_times or indexes is None:
raise NotImplementedError(
"Specifying `loadable_variables` or auto-creating indexes with `indexes=None` is not supported for dmrpp files."
)

if group:
raise NotImplementedError()

fpath = _FsspecFSFromFilepath(
filepath=filepath, reader_options=reader_options
).open_file()

parser = DMRParser(fpath.read(), data_filepath=filepath.strip(".dmrpp"))
vds = parser.parse_dataset()

return vds.drop_vars(drop_variables)


class DMRParser:
"""
Parser for the OPeNDAP DMR++ XML format.
Expand Down Expand Up @@ -69,9 +106,7 @@ def __init__(self, dmr: str, data_filepath: Optional[str] = None):
data_filepath if data_filepath is not None else self.root.attrib["name"]
)

def parse_dataset(
self, group=None, indexes: Mapping[str, Index] = {}
) -> xr.Dataset:
def parse_dataset(self, group=None, indexes: Mapping[str, Index] = {}) -> Dataset:
"""
Parses the given file and creates a virtual xr.Dataset with ManifestArrays.

Expand Down Expand Up @@ -128,7 +163,7 @@ def _parse_netcdf4_dataset(
root: ET.Element,
group: Optional[str] = None,
indexes: Mapping[str, Index] = {},
) -> xr.Dataset:
) -> Dataset:
"""
Parse the dataset from the netcdf4 based dmrpp with groups, starting at the given group.
Set root to the given group.
Expand Down Expand Up @@ -201,7 +236,7 @@ def _parse_hdf5_dataset(
root: ET.Element,
group: Optional[str] = None,
indexes: Mapping[str, Index] = {},
) -> xr.Dataset:
) -> Dataset:
"""
Parse the dataset from the HDF5 based dmrpp with groups, starting at the given group.
Set root to the given group.
Expand Down Expand Up @@ -331,7 +366,7 @@ def _split_hdf5(self, root: ET.Element) -> dict[str, ET.Element]:

def _parse_dataset(
self, root: ET.Element, indexes: Mapping[str, Index] = {}
) -> xr.Dataset:
) -> Dataset:
"""
Parse the dataset using the root element of the DMR file.

Expand All @@ -353,8 +388,8 @@ def _parse_dataset(
if len(coord_names) == 0 or len(coord_names) < len(dataset_dims):
coord_names = set(dataset_dims.keys())
# Seperate and parse coords + data variables
coord_vars: dict[str, xr.Variable] = {}
data_vars: dict[str, xr.Variable] = {}
coord_vars: dict[str, Variable] = {}
data_vars: dict[str, Variable] = {}
for var_tag in self._find_var_tags(root):
variable = self._parse_variable(var_tag, dataset_dims)
if var_tag.attrib["name"] in coord_names:
Expand All @@ -365,9 +400,9 @@ def _parse_dataset(
attrs: dict[str, str] = {}
for attr_tag in self.root.iterfind("dap:Attribute", self._ns):
attrs.update(self._parse_attribute(attr_tag))
return xr.Dataset(
return Dataset(
data_vars=data_vars,
coords=xr.Coordinates(coords=coord_vars, indexes=indexes),
coords=Coordinates(coords=coord_vars, indexes=indexes),
attrs=attrs,
)

Expand Down Expand Up @@ -484,7 +519,7 @@ def _parse_multi_dims(

def _parse_variable(
self, var_tag: ET.Element, dataset_dims: dict[str, int]
) -> xr.Variable:
) -> Variable:
"""
Parse a variable from a DMR tag.

Expand Down Expand Up @@ -542,7 +577,7 @@ def _parse_variable(
)
marr = ManifestArray(zarray=zarray, chunkmanifest=chunkmanifest)
encoding = {k: attrs.get(k) for k in self._encoding_keys if k in attrs}
return xr.Variable(
return Variable(
dims=dim_shapes.keys(), data=marr, attrs=attrs, encoding=encoding
)

Expand Down
Loading
Loading