xarray like lazy loading support

[rajeeja]

Proposed new feature or change:

Leverage Xarray's existing mechanisms for lazy I/O, which operate even without Dask.

[philipc2]

This is something we've actively been considering, but might be challenging considering our explicit requirement to convert grids into the UGRID conventions.

At the most basic form, this involves

• Renaming dimensions
• Converting to zero index
• Replacing fill values

For grids like SCRIP, we generate the connectivity information directly, which can't be done lazily unless we delay the execution until it is required. See for example what we do with HEALPix

uxarray/uxarray/grid/grid.py

Lines 1151 to 1155 in 30a534b

	if (
	"face_node_connectivity" not in self._ds
	and self.source_grid_spec == "HEALPix"
	):
	_populate_healpix_boundaries(self._ds)

See some discussion in #1092

[philipc2]

We'd need to significantly re-write all of our grid readers to support this. The following operations can be done lazily using Xarray's existing mechanisms

• Inspecting Metadata: .dims, .sizes, etc.
• Metadata Operations: .rename(), .assign_attrs(), etc.
• Index Selections: .isel()

However, the following cause the data to be eagerly loaded into memory

• Arithmatic Operations
• Type Casting
• Masking & Conditional Logic

Those operations above are crucial for our conversion to the UGRID conventions.

Dask allows us to delay the execution of those eager operations. For grid formats that don't require any conversions (i.e. some UGRID files), the data can be already loaded in without doing any operation.

[rajeeja]

Good thought and thanks for the detailed planning. I'm thinking as a start we can simply consider lazy loading Data Variables: Variables containing physical quantities defined on the grid nodes, faces, or edges.

[philipc2]

Good thought and thanks for the detailed planning. I'm thinking as a start we can simply consider lazy loading Data Variables: Variables containing physical quantities defined on the grid nodes, faces, or edges.

That is already the default behavior. We only change the dimension name of data variables.

[rajeeja]

Xarray's doesn't load and instead, represents objects with

xarray.core.indexing.LazilyIndexedArray

that wrap the underlying access, what do we do currently?

[rajeeja]

Xarray's doesn't load and instead, represents objects with

xarray.core.indexing.LazilyIndexedArray

that wrap the underlying access, what do we do currently?

Gotcha, browsing through the code, I find that we don't use the xarray.core, rather we use of xr.open_dataset() that inherently supports lazy loading by default, as xarray loads data into memory only when it is explicitly accessed