Feature: Virtual datastores to reduce storage redundancies

[observingClouds]

Currently, mllam-data-prep is mostly reordering and rechunking the input datasets for better performance. Sometimes, particularly if the dataset is large or training shall be done on different subsets of the available datasets, keeping several copies of very similar datasets can be challenging and costs a lot of storage.

Creating reference datasets, i.e. datasets that only contain links/references to the original dataset, can be advantageous as no additional data is copied.

This issue is meant to document the potential options and start a discussion on potential implementations.

Option 1: Create subset of mllam-data-prep zarr
This option requires the mllam-data-prep zarr to be chunked with a chunksize of 1 along the subsetting dimension, e.g. state_features, as subsetting across chunk boundaries is not supported. This can be set in the mllam-data-prep input yaml file. After creation subsets for e.g. ablation studies can be created by:

from virtualizarr import open_virtual_dataset
from virtualizarr.parsers import ZarrParser
from obstore.store import LocalStore
from virtualizarr.registry import ObjectStoreRegistry
from pathlib import Path

zarr_store = str(Path.cwd() / 'example.danra.zarr')
store = LocalStore(prefix=zarr_store)
registry = ObjectStoreRegistry({f"file://{zarr_store}": store})
parser = ZarrParser()
vds = open_virtual_dataset( url=zarr_store,registry=registry,parser=parser)

# Select subset of mllam-data-prep dataset, e.g. reduce number of state_features
vds.isel(state_feature=slice(1,2)).vz.to_kerchunk("index.json", format='json')

# New subsetted mllam-data-prep dataset just containing references to original mllam-data-prep (e.g. no extra copy)
xr.open_zarr("reference://", storage_options={"fo":"index.json"})

Edited: this seems to require mllam-data-prep zarr files to be written in zarr format 3

Option 2: Create dataset with references to source files
Option 1 requires a full "physical" creation of the largest training dataset, while only the subsets are links. Can a reference dataset be created directly? How to handle variables that are precomputed, like statistics or time_of_day forcings?

[leifdenby]

Option 1: Create subset of mllam-data-prep zarr
This option requires the mllam-data-prep zarr to be chunked with a chunksize of 1 along the subsetting dimension, e.g. state_features, as subsetting across chunk boundaries is not supported. This can be set in the mllam-data-prep input yaml file. After creation subsets for e.g. ablation studies can be created by:

This is really cool! Are you thinking this functionality should be added to mllam-data-prep itself, or would it sit better in downstream applications (e.g. neural-lam)? If it would go in mllam-data-prep then would you change the config? It feels to me like that it would still make sense to have an unchanged config for the "full dataset", but I am not sure having config file for this virtual datasets quite makes sense, since it removes the one-to-one correspondence between the dataset config yaml and the dataset zarr that was created from it. Instead maybe we could add a convenience function, say mllam_data_prep.load_virtual_dataset_subset(source_dataset_config). Maybe we should have a chat on a whiteboard about how this could be designed?

[observingClouds]

Are you thinking this functionality should be added to mllam-data-prep itself, or would it sit better in downstream applications (e.g. neural-lam)?

The functionality of option 1 could be just documented in neural-lam docs. Before users can make use of these virtual datasets, neural-lam needs some modifications, I would particularly advocate to separate the convention that the data file has to have the same name as the config file with the only difference being the file ending (.yaml vs .zarr). The virtual datasets need for example the reference:// prefix.

[leifdenby]

The functionality of option 1 could be just documented in neural-lam docs.

I think implementing in neural-lam makes sense to me, but we could also add what you've written above in the README of mllam-data-prep and then link to the implementation in neural-lam as a working example.

Before users can make use of these virtual datasets, neural-lam needs some modifications, I would particularly advocate to separate the convention that the data file has to have the same name as the config file with the only difference being the file ending (.yaml vs .zarr). The virtual datasets need for example the reference:// prefix.

I see why you would suggest that for the added flexibility, but I think keeping with the idea that a single dataset is made from a single config (whether that be stored as a separate yaml-file or serialised into the dataset itself as well as-of v0.6.0) wrt mllam-data-prep is good idea. Otherwise we start changing the idea of what a "dataset" is, as in, for every virtual dataset you might make (using the approach you wrote above) the zarr dataset on disk as produced by mllam-data-prep wouldn't change. I think a simpler way of viewing these "virtual datastores" might be as a sort of "view" on a dataset where you only end up seeing part of the dataset, but what is on disk (and what mllam-data-prep itself outputs) doesn't change. It would be different if neural-lam used an open_dataset function implemented within mllam-data-prep (because then it was up to mllam-data-prep how to open a dataset), but I intentionally didn't design it like this because I wanted the zarr dataset on disk to be self-contained, so that anyone could call xr.open_dataset() on a dataset produced by mllam-data-prep and that would work without needing to involve mllam-data-prep itself.

We could instead introduce a VirtualDatastore in neural-lam that adds a higher-level configuration, something like this:

# neural-lam.yaml
datastore:
  kind: virtual-mdp
  config_path: danra.virtual_subset.datastore.yaml
...

# danra.virtual_subset.datastore.yaml
base_datastore_config_path: danra.datastore.yaml
subsetting:  # subsetting using coordinate values on named coordinates below
   state_feature: [u10m, v10m]

If we then refactor the current MDPDatastore to move the actual opening of the zarr dataset to a private method called _open_dataset() -> xr.Dataset we could then override that in a new subclass called VirtualMDPDatastore with the approach you wrote above. How does that sound? Of course it might be nicer not to introduce a third configuration file, in that case we could modify the neural-lam config to look like this:

# neural-lam.yaml
datastore:
  kind: virtual-mdp
  config_path: danra.datastore.yaml
  subsetting:
     state_feature: [u10m, v10m]
...

and only allow the subsetting config parameter to be set of virtual-mdp is used for the datastore kind.

[Jayant-kernel]

I've been following the discussion on storage efficiency and the consensus on moving the virtual datastore logic into neural-lam. I'd love to take this on!

My plan :

Subclassing: Creating a VirtualMDPDatastore in neural-lam that inherits from the current datastore class.

Refactoring: Moving the xr.open_zarr logic into a private _open_dataset() method to allow for easy overriding.

Virtualization: Integrating virtualizarr within that override to generate Kerchunk-style references on the fly based on the subsetting parameters in the config.

Config Update: Extending the neural-lam.yaml schema to support the kind: virtual-mdp and subsetting fields.

I'm ready to start drafting a PR for this—could I be assigned to this issue?

[observingClouds]

@leifdenby I like your concept. The question of integrating the subsetting setting into an independent datastore, or into neural-lam.yaml basically decides where we would like to have this capability implemented.

It probably feels more natural to implement the subsetting into neural-lam, because mllam-data-prep needs to be used already to create the original merged dataset. If used iteratively a second time to create a virtual sub-set, we will quickly loose track of the initial data source. However, statistics and forcings might need to be recomputed and we might also want to do subsetting e.g. in time or directly from data-sources. These virtualizations would probably be better included in a datastore.yaml.

@Jayant-kernel I think we still need to discuss how to best implement this. You propose to include the sub-setting in neural-lam. Feel free to share your views on what you prefer and why.

[Jayant-kernel]

@observingClouds . This is a great architectural trade-off.

I actually lean towards implementing the primary subsetting in datastore.yaml, with the option for neural-lam.yaml to override it. Here’s why:

Robustness: Keeping subsetting in datastore.yaml ensures statistics and coordinate metadata stay 'attached' to the data definition, preventing the 'loss of track' you mentioned.

Reusability: A virtual datastore.yaml remains valid and portable even if we switch training frameworks or models in the future.

Separation of Concerns: We can treat datastore.yaml as the Physical/Virtual Schema and neural-lam.yaml as the Client Query.

Revised Plan:
Draft VirtualMDPDatastore to prioritize reading from these virtual configs.
Include a small utility to 'export' a virtual datastore.yaml from an existing Zarr to simplify user setup.

[Jayant-kernel]

@observingClouds
Thanks for the discussion — to move forward, should we first align on:

Whether subsetting lives primarily in datastore.yaml (with optional neural-lam override), and

Whether we proceed with introducing VirtualMDPDatastore as a separate class?

If this direction sounds good, I can draft a minimal design proposal (class diagram + config schema) before starting implementation.

[leifdenby]

However, statistics and forcings might need to be recomputed and we might also want to do subsetting e.g. in time or directly from data-sources

Could you give me an example where statistics and forcings would need to be recomputed? Recomputing statistics could be quite costly, so you wouldn't want to do that every time you create a virtual datastore, but then where would you save the statistics? And wrt to the statistics then the train/val/test split might invalid too, right?

If we want to do quite costly operations and persist the results of these to disk maybe what we are looking at is a new config-file for mllam-data-prep where want can derive one dataset from another, i.e. the config file for this virtual datastore would point to (path to) a different actual datastore, and then have specific config arguments for the extra operations one might want to do. The config would look quite different though since you are operating in the transformed space of the output of the original mllam-data-prep dataset create operation. Or maybe we just see if we can add a virtual argument to the inputs in the existing mllam-data-prep config an implement all the other logic on top of that...

[Jayant-kernel]

@leifdenby Thanks — that makes sense. I agree we shouldn’t be recomputing stats/forcings every time someone makes a “virtual” datastore.

A few concrete cases where recompute really would be needed:

• Time subsetting (e.g. 2010–2018 vs 2019–2021): mean/std changes + the train/val/test split needs to be regenerated for the new timeline.
• Feature/channel subsetting (dropping/adding variables): normalization stats need to match the final feature set.
• (Later) spatial subsetting/masking: distribution shifts, so stats change.

To keep this sane, maybe we split it into two modes:

• “Virtual view” (cheap): only supports ops that don’t create new derived arrays; stats/forcings are reused from the parent datastore.
• “Derived datastore” (explicit + persisted): if user asks for recompute (stats/forcings/splits), run it as a data-prep job and write outputs to disk.

If you’re aligned, I can start with an MVP PR for the “virtual view” only (no recompute):

• allow simple feature subsetting (and maybe time indexing as a view)
• add a tiny helper to export a virtual_datastore.yaml pointing at an existing zarr

Does that direction sound reasonable, or would you prefer we keep everything inside an mllam-data-prep config from day one?