open_virtual() #1348
Conversation
… homogenization - open_virtual(): open kerchunk/parquet/json/icechunk stores from URI or DataCollection with load=True (kerchunk engine) or load=False (VirtualiZarr) - force_external: download refs, rewrite s3:// to https:// for external access - DataCollection.virtual_collection_url(), get_s3_credentials(), get_s3_filesystem() - homogenize_dataset_codec_level(): patch Zlib codec levels on ManifestArrays - virtualize() default access changed from direct to indirect - build_obstore_registry: remove bearer token from HTTPStore (proxies return 303 redirects) - Update deps: virtualizarr latest, icechunk, matplotlib, hvplot; remove pandas pin
|
I will automatically update this comment whenever this PR is modified
|
…n region access from gh actions
# Conflicts: # tests/unit/test_virtual.py
Mock get_s3_filesystem to avoid real network calls in CI. Suppress ruff SIM117 by combining nested with statements.
|
|
||
| Parameters: | ||
| uri: A ``DataCollection``, or a path/URI to the virtual store | ||
| (``.icechunk``, ``.parquet``, or ``.json``). |
There was a problem hiding this comment.
For the TEMPO month-long icechunk, it's a directory (see this example notebook for reading the current, experimental TEMPO icechunk), not a .icechunk file. Not sure if that's always the case or not for icechunks.
There was a problem hiding this comment.
yeah good catch, this was just a convention but can be changed, I think another way of making sure it's an icechunk store is the protocol prefix, icechunk:// that the icechunk store seems to know how to handle. @maxrjones, do you know the canonical way for this?
There was a problem hiding this comment.
I added a DistributionUrl to the metadata record for the TEMPO collection experimented with in that notebook (referenced in my last comment). The current code raises an error for that icechunk in S3 because it is a directory, without .icechunk suffix:
import earthaccess as ea
ea.login()
collection = ea.search_datasets(short_name="TEMPO_NO2_L3", version="V04")
vds = ea.open_virtual(collection[0]) # raises the errorerror:
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
Cell In[15], line 11
7 short_name="TEMPO_NO2_L3",
8 version="V04",
9 )
10
---> 11 vds = ea.open_virtual(collection[0])
12 vds
File [/srv/conda/envs/notebook/lib/python3.13/site-packages/earthaccess/virtual/core.py:757](https://openscapes.2i2c.cloud/srv/conda/envs/notebook/lib/python3.13/site-packages/earthaccess/virtual/core.py#line=756), in open_virtual(uri, access, storage_options, force_external, load, **kwargs)
752 if not _is_icechunk_uri(url) and not _is_kerchunk_uri(url):
753 msg = (
754 f"Unrecognised virtual store URL in collection: {url}. "
755 "Expected a .icechunk, .parquet, or .json file."
756 )
--> 757 raise ValueError(
758 msg,
759 )
760 else:
761 url = str(uri)
ValueError: Unrecognised virtual store URL in collection: s3://asdc-prod-public/virtual-reference-docs/TEMPO_L3_V04_icechunk_for_202601. Expected a .icechunk, .parquet, or .json file.
There was a problem hiding this comment.
Yeah we need to change how we validate the VDS URL. Should we come up with a convention? e.g. if any in the URL we match icechunk... and we are opening a VDS we assume icechunk. This is to avoid users having to guess the store.
There was a problem hiding this comment.
This is the leading proposal for a canonical URL representation of Zarr/Icechunk: https://github.com/jbms/url-pipeline.
It's still a draft, but it's where the ecosystem is converging. It's implemented in TensorStore and Neuroglancer, with a zarr-python PR in progress (#3369). Earthmover's Zarrs.jl already uses it. Its icechunk: scheme directly covers our case. <url>|icechunk: addresses a repo root (no node), and the spec defines directory-based format auto-detection. I'll check with Ian about the status of the zarr-python PR. I'd avoid inventing yet another URL scheme (e.g. ic+http) when this exists.
|
|
||
| Parameters: | ||
| uri: A ``DataCollection``, or a path/URI to the virtual store | ||
| (``.icechunk``, ``.parquet``, or ``.json``). |
There was a problem hiding this comment.
I added a DistributionUrl to the metadata record for the TEMPO collection experimented with in that notebook (referenced in my last comment). The current code raises an error for that icechunk in S3 because it is a directory, without .icechunk suffix:
import earthaccess as ea
ea.login()
collection = ea.search_datasets(short_name="TEMPO_NO2_L3", version="V04")
vds = ea.open_virtual(collection[0]) # raises the errorerror:
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
Cell In[15], line 11
7 short_name="TEMPO_NO2_L3",
8 version="V04",
9 )
10
---> 11 vds = ea.open_virtual(collection[0])
12 vds
File [/srv/conda/envs/notebook/lib/python3.13/site-packages/earthaccess/virtual/core.py:757](https://openscapes.2i2c.cloud/srv/conda/envs/notebook/lib/python3.13/site-packages/earthaccess/virtual/core.py#line=756), in open_virtual(uri, access, storage_options, force_external, load, **kwargs)
752 if not _is_icechunk_uri(url) and not _is_kerchunk_uri(url):
753 msg = (
754 f"Unrecognised virtual store URL in collection: {url}. "
755 "Expected a .icechunk, .parquet, or .json file."
756 )
--> 757 raise ValueError(
758 msg,
759 )
760 else:
761 url = str(uri)
ValueError: Unrecognised virtual store URL in collection: s3://asdc-prod-public/virtual-reference-docs/TEMPO_L3_V04_icechunk_for_202601. Expected a .icechunk, .parquet, or .json file.
| # --------------------------------------------------------------------------- | ||
|
|
||
|
|
||
| def _is_icechunk_uri(uri: str) -> bool: |
There was a problem hiding this comment.
This needs to be refactored, should we just match icechunk in the URI? should we use ic+http protocol notation?
There was a problem hiding this comment.
url-pipeline is the way to go eventually for both kerchunk and icechunk identification, but I also understand wanting to get this out and think this could be a TODO for later.
FWIW the definitive way to identify an icechunk store is via the magic bytes in the root file IIUC, but the url-pipeline avoids needing to send an extra request.
| repo = icechunk.Repository.open(storage=storage) | ||
| session = repo.readonly_session("main") | ||
| store = session.store | ||
| return xr.open_zarr(store, **kwargs) |
There was a problem hiding this comment.
The reason to return xarray ds instead of the store is to have a unified API with kerchunk. Also what else can we do with a store if it's not opening it with xarray?
There was a problem hiding this comment.
you could use Icechunk stores with the Zarr API instead of the Xarray API
| url: str, | ||
| access: str = "indirect", | ||
| **kwargs: Any, | ||
| ) -> xr.Dataset: |
There was a problem hiding this comment.
This superseeds the PR (draft) Julius started working #1135
|
After testing more, seems like we could go ahead with the current PR and update to more native use of virtualizarr once some of the following issues get addressed. Virtualizarr only supports fsspec for opening a kerchunk store, this is kerchunk to ManifestStore, (Kerchunk parsers) and we still have issues with decoding FillValue. One thing we need to define is how CMR will encode the store type, Max suggested following the convention here: https://github.com/jbms/url-pipeline/blob/main/schemes/icechunk.md looks good, although my personal preference would be to use the protocol prefix like git e.g. {store_type}+{net_protocol}:// so one could have What do you guys think? @maxrjones @danielfromearth |
|
Sorry - super late comments here: I was chatting with the CMR folks late last week to try and work out if this is a suitable related URL type. Apparently "VIRTUAL COLLECTION" has a different intent (grouping together data from multiple collections), and they are keen for us to adopt a different sub type. I've been tasked with getting the right vocabulary from this crowd. I was planning to use @abarciauskas-bgse's next fortnightly co-working session to get the right folks together to decide on the best terminology. One other thought - I wonder if we should actually move away from stating "VIRTUAL", too? Isn't the access pattern the same if we're pointing at a virtual store, or the entry point of an actual Zarr store? (Probably getting into the conversation we could have at Aimee's session) Once we have the type and sub type determined, the CMR folks can implement it pretty quickly. |
maxrjones
left a comment
There was a problem hiding this comment.
awesome work, @betolink!
It'd be worth trying to load a small snippet of the virtual Icechunk array in the notebook. I'm surprised you don't need any authorize_virtual_chunk_access=... logic for Icechunk stores. It's possible the demo just doesn't get to the point where virtual chunk access authentication is needed. https://www.earthmover.io/blog/secure-virtual-chunks provides helpful context for this configuration setting.
One conceptual correction about the virtual-diagram is that virtualizarr does not require any JSON for interchange. The only JSON usage is if Kerchunk is used as a starting point or output format. Otherwise, the byte ranges and uris are all handled as numpy arrays (passed directly to Rust when using Icechunk), which is more efficient than using dicts + JSON.
| "dask-labextension>=7.0.0", | ||
| "dask>=2025.7.0", | ||
| "distributed>=2025.7.0", | ||
| "holoviews>=1.22.1", | ||
| "panel>=1.8.10", |
There was a problem hiding this comment.
these are some heavy dependencies. I recommend moving them to optional dependency groups.
There was a problem hiding this comment.
Good catch, they should be optional.
| "matplotlib", | ||
| "hvplot", |
There was a problem hiding this comment.
virtualizarr capabilities don't require hvplot or matplotlib. I suggest these should be in a separate optional dependency group (e.g., viz, following xarray's naming).
| def homogenize_dataset_codec_level(ds: xr.Dataset, target_level: int = 7) -> xr.Dataset: | ||
| """Patch Zlib codec levels on all ManifestArrays in *ds* to *target_level*.""" | ||
| try: | ||
| import xarray as xr | ||
| from virtualizarr.manifests import ManifestArray | ||
| except ImportError: | ||
| msg = ( | ||
| "earthaccess.virtualize() requires `pip install earthaccess[virtualizarr]`" | ||
| ) | ||
| raise ImportError( | ||
| msg, | ||
| ) from None | ||
|
|
||
| def _patch_ma(ma: ManifestArray, level: int) -> ManifestArray: | ||
| if not isinstance(ma, ManifestArray): | ||
| return ma | ||
|
|
||
| meta = ma.metadata | ||
| codecs = list(meta.codecs) | ||
| modified = False | ||
|
|
||
| for i, codec in enumerate(codecs): | ||
| if not hasattr(codec, "to_dict"): | ||
| continue | ||
| cd = codec.to_dict() | ||
| if cd.get("name") != "numcodecs.zlib": | ||
| continue | ||
| cfg = cd.get("configuration", {}) | ||
| if isinstance(cfg, dict): | ||
| if cfg.get("level") == level: | ||
| continue | ||
| cd["configuration"] = {**cfg, "level": level} | ||
| else: | ||
| cd["configuration"] = {"level": level} | ||
| codecs[i] = type(codec).from_dict(cd) | ||
| modified = True | ||
|
|
||
| if not modified: | ||
| return ma | ||
|
|
||
| import zarr | ||
|
|
||
| new_meta = zarr.core.metadata.v3.ArrayV3Metadata( | ||
| shape=meta.shape, | ||
| data_type=meta.data_type, | ||
| chunk_grid=meta.chunk_grid, | ||
| chunk_key_encoding=meta.chunk_key_encoding, | ||
| fill_value=meta.fill_value, | ||
| codecs=tuple(codecs), | ||
| attributes=meta.attributes, | ||
| dimension_names=meta.dimension_names, | ||
| storage_transformers=meta.storage_transformers, | ||
| ) | ||
| return ManifestArray(new_meta, ma.manifest) | ||
|
|
||
| new_vars = { | ||
| name: ( | ||
| xr.DataArray( | ||
| _patch_ma(var.data, target_level), | ||
| dims=var.dims, | ||
| attrs=var.attrs, | ||
| name=name, | ||
| ) | ||
| if isinstance(var.data, ManifestArray) | ||
| else var | ||
| ) | ||
| for name, var in ds.data_vars.items() | ||
| } | ||
|
|
||
| return xr.Dataset(new_vars, coords=ds.coords, attrs=ds.attrs) |
There was a problem hiding this comment.
this seems indicative of an upstream problem. what motivated this patch? what issues were you running into?
There was a problem hiding this comment.
In some datasets like PODAAC SST, there is a change in the compression level for the codecs, Virtualizarr checks for homogeneity in the codecs along the concatenating dimension so this would fail that. However, in practice level 6 and level 7 do not alter the hashes is basically the same compression for all values in the dataset. This is an experimental feature to pass VZ checks.
| # --------------------------------------------------------------------------- | ||
|
|
||
|
|
||
| def _is_icechunk_uri(uri: str) -> bool: |
There was a problem hiding this comment.
url-pipeline is the way to go eventually for both kerchunk and icechunk identification, but I also understand wanting to get this out and think this could be a TODO for later.
FWIW the definitive way to identify an icechunk store is via the magic bytes in the root file IIUC, but the url-pipeline avoids needing to send an extra request.
| repo = icechunk.Repository.open(storage=storage) | ||
| session = repo.readonly_session("main") | ||
| store = session.store | ||
| return xr.open_zarr(store, **kwargs) |
There was a problem hiding this comment.
you could use Icechunk stores with the Zarr API instead of the Xarray API
| via the kerchunk engine. When ``False``, returns a virtual dataset | ||
| backed by ``ManifestArray`` objects via VirtualiZarr's | ||
| ``open_virtual_dataset``. | ||
| **kwargs: Additional keyword arguments forwarded to the opener. |
There was a problem hiding this comment.
since the opener isn't public, it's hard for users to now what kwargs are accepted. I personally prefer smaller, well-typed functions (in this case a separate open_kerchunk and open_icechunk function that could have kwargs aligned with the 3rd party implementation) but recognize that others prefer fewer more feature complete functions. Having separate open_icechunk and open_kerchunk functions would also resolve the need for url-pipeline resolution or other url based heuristics.
There was a problem hiding this comment.
this is a key question for this PR, I'm trying to avoid splitting this into open_icechunk() and open_kerchunk() and follow a bit what xarray does with the engines. Users won't need to be concerned about it and if CMR follows the url pipeline convention or some simpler convention to start with, we should be fine. We can also expose those at the top level and let power users decide.
There was a problem hiding this comment.
I came around to think that open_icechunk and open_kerchunk are better off as their own top level methods. I think all of them should return an xarray dataset object, when you mentioned "you could use Icechunk stores with the Zarr API instead of the Xarray API" is just the way we open it? or the instance it gets returned? @maxrjones
| def _sanitize_references_for_external(url: str) -> str: | ||
| """Download a kerchunk reference file, rewrite ``s3://`` URLs to ``https://``. | ||
|
|
||
| The HTTPS base is inferred from the reference file URL's host, which works | ||
| for all NASA Earthdata Cloud DAACs (PODAAC, NSIDC, GES DISC, LP DAAC, …). | ||
|
|
||
| Returns the local path to the sanitized file. | ||
| """ |
There was a problem hiding this comment.
it should be possible to use virtualizarr for this, but that could be a simple TODO to resolve later
| parser: KerchunkJSONParser | KerchunkParquetParser | ||
| if url.endswith(".json"): | ||
| parser = KerchunkJSONParser(skip_variables=list(kds.coords)) | ||
| elif url.endswith(".parquet"): | ||
| parser = KerchunkParquetParser(skip_variables=list(kds.coords)) | ||
| else: | ||
| msg = ( | ||
| f"Unsupported virtual store format: {url}. " | ||
| "Expected a .json or .parquet file." | ||
| ) | ||
| raise ValueError( | ||
| msg, | ||
| ) | ||
|
|
||
| vds = vz.open_virtual_dataset(ref_path, parser=parser, registry=registry, **kwargs) | ||
|
|
||
| for k in kds.coords: | ||
| vds.coords[k] = kds[k] | ||
| return vds |
There was a problem hiding this comment.
I don't understand why you need L667-685 if you already have an xarray dataset opened via L665.
This PR adds
open_virtual()a method that simplifies how we open virtual stores behind EDL authentication.We can open virtual collections listed at the collection level metadata, for now following the convention adopted by JPL using the following entry in
RelatedUrls{ "Description": "Virtual data set reference ", "URLContentType": "DistributionURL", "URL": "https://archive.podaac.earthdata.nasa.gov/../MUR-JPL-L4-GLOB-v4.1_combined-ref.json", "Type": "GET DATA", "Subtype": "VIRTUAL COLLECTION" }open_virtual()has the following kwargsuri: this can be a collection level result from collection = earthaccess.search_datasets()[0] a local path, or a publicly accessible reference URL. If we use a collection from earthaccess, we don't need to deal with the cumbersome auth scafolding for the underlaying libraries; in all cases we'll get an xarray dataset instance pointing to our virtual dataset.This can be kerchunk or Icechunk storesaccess: this one informs earthaccess if the virtual store points to S3 buckets or HTTP, without inspecting the references we need to be explicit for now. values are direct or external (same as indirect)storage_options: if we need custom authentication for the virtual store this works the same way xarray uses storage options, not needed for NASA data as it's automatic.force_external: This is a handy flag to force reference translation, this is if the only available store points to S3 but we are accessing the data from outside AWS us-west-2 setting this to True will translate the URLs on the fly for out of region access.load: if we should load the indexes in our dataset, if we do we can address the dataset as any other xarray dataset and perform operations like subsetting and aggregations, if not we get a virtualizarr dataset backed by ManifestArray instances, this is only useful if we plan to add more virtual references and update the cube.The method signature is as follows:
Another relevant aspect is that
load=Falseis slower thanload=Truefor now as an artifact of VirtualiZarr not parsing all the coordinates and variables with embedded values(inlined values). The last release of VirtualiZarr supports it but I ran into some issues. For nowload=Falserequires a round trip to a local store for the references, this worked for all stores but adds latency for example with a JSON serialized store opening the references can take up to a couple minutes depending on internet speed.load=Trueis fast as we only load coords and lazy load references (chunk manifests). This will have to change for upcoming releases to improve speed.Here is a notebook showing some the the capabilities for this new feature:
https://notebooksharing.space/view/d2eb24605b6fb6b011710d955052f67de77c38728afdbc48ccc9cc33db34ff57#displayOptions=
Another relevant topic is that Icechunk does not support out of region access for stores with authentication, if something like this PR gets mered we will have the whole set of data virtualization available and no further changes are required in earthaccess. earth-mover/icechunk#2143
Description
"Ready for review" checklist
Merge checklist
closes #1)CHANGELOG.mdupdatedREADME.mdupdatedpre-commit.ci autofixif pre-commit is failing)📚 Documentation preview 📚: https://earthaccess--1348.org.readthedocs.build/en/1348/