Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Support Zarr-Python 3
Browse files Browse the repository at this point in the history
maxrjones committed Jan 28, 2025
1 parent c7c5426 commit ba2b1b5
Showing 2 changed files with 26 additions and 70 deletions.
15 changes: 12 additions & 3 deletions src/titiler/xarray/pyproject.toml
Original file line number Diff line number Diff line change
@@ -38,15 +38,15 @@ dependencies = [

[project.optional-dependencies]
full = [
"zarr>=2,<3",
"zarr",
"h5netcdf",
"fsspec",
"s3fs",
"aiohttp",
"gcsfs",
]
minimal = [
"zarr>=2,<3",
"zarr",
"h5netcdf",
"fsspec",
]
@@ -64,7 +64,16 @@ test = [
"pytest-cov",
"pytest-asyncio",
"httpx",
"zarr>=2,<3",
"zarr",
"h5netcdf",
"fsspec",
]
test-zarr-v2 = [
"pytest",
"pytest-cov",
"pytest-asyncio",
"httpx",
"zarr<3.0",
"h5netcdf",
"fsspec",
]
81 changes: 14 additions & 67 deletions src/titiler/xarray/titiler/xarray/io.py
Original file line number Diff line number Diff line change
@@ -9,6 +9,7 @@
from morecantile import TileMatrixSet
from rio_tiler.constants import WEB_MERCATOR_TMS
from rio_tiler.io.xarray import XarrayReader
from xarray.namedarray.utils import module_available


def xarray_open_dataset( # noqa: C901
@@ -29,21 +30,6 @@ def xarray_open_dataset( # noqa: C901
"""
import fsspec # noqa

try:
import gcsfs
except ImportError: # pragma: nocover
gcsfs = None # type: ignore

try:
import s3fs
except ImportError: # pragma: nocover
s3fs = None # type: ignore

try:
import aiohttp
except ImportError: # pragma: nocover
aiohttp = None # type: ignore

try:
import h5netcdf
except ImportError: # pragma: nocover
@@ -66,55 +52,8 @@ def xarray_open_dataset( # noqa: C901

else:
assert zarr is not None, "'zarr' must be installed to read Zarr dataset"

xr_engine = "zarr"

if protocol in ["", "file"]:
filesystem = fsspec.filesystem(protocol) # type: ignore
file_handler = (
filesystem.open(src_path)
if xr_engine == "h5netcdf"
else filesystem.get_mapper(src_path)
)

elif protocol == "s3":
assert (
s3fs is not None
), "'aiohttp' must be installed to read dataset stored online"

s3_filesystem = s3fs.S3FileSystem()
file_handler = (
s3_filesystem.open(src_path)
if xr_engine == "h5netcdf"
else s3fs.S3Map(root=src_path, s3=s3_filesystem)
)

elif protocol == "gs":
assert (
gcsfs is not None
), "'gcsfs' must be installed to read dataset stored in Google Cloud Storage"

gcs_filesystem = gcsfs.GCSFileSystem()
file_handler = (
gcs_filesystem.open(src_path)
if xr_engine == "h5netcdf"
else gcs_filesystem.get_mapper(root=src_path)
)

elif protocol in ["http", "https"]:
assert (
aiohttp is not None
), "'aiohttp' must be installed to read dataset stored online"

filesystem = fsspec.filesystem(protocol) # type: ignore
file_handler = (
filesystem.open(src_path)
if xr_engine == "h5netcdf"
else filesystem.get_mapper(src_path)
)

else:
raise ValueError(f"Unsupported protocol: {protocol}, for {src_path}")
_zarr_v3 = module_available("zarr", minversion="3.0")

# Arguments for xarray.open_dataset
# Default args
@@ -135,13 +74,21 @@ def xarray_open_dataset( # noqa: C901
"lock": False,
}
)

ds = xarray.open_dataset(file_handler, **xr_open_args)
fs = fsspec.filesystem(protocol)
ds = xarray.open_dataset(fs.open(src_path), **xr_open_args)

# Fallback to Zarr
else:
ds = xarray.open_zarr(file_handler, **xr_open_args)

if _zarr_v3:
if protocol == "file":
store = zarr.storage.LocalStore(parsed.path, read_only=True)
else:
fs = fsspec.filesystem(protocol, storage_options={"asynchronous": True})
store = zarr.storage.FsspecStore(fs, path=src_path, read_only=True)
ds = xarray.open_zarr(store, **xr_open_args)
else:
fs = fsspec.filesystem(protocol)
ds = xarray.open_zarr(fs.get_mapper(src_path), **xr_open_args)
return ds


0 comments on commit ba2b1b5

Please sign in to comment.