Skip to content

xarray backend from index alone: attribute capture, serializable indexes, and S3 reads (drops the netCDF4 dependency) - #50

Open
espg wants to merge 31 commits into
gauteh:mainfrom
espg:claude/xarray-s3-engine
Open

xarray backend from index alone: attribute capture, serializable indexes, and S3 reads (drops the netCDF4 dependency)#50
espg wants to merge 31 commits into
gauteh:mainfrom
espg:claude/xarray-s3-engine

Conversation

@espg

@espg espg commented Aug 9, 2026

Copy link
Copy Markdown
Contributor

This PR makes the xarray backend self-sufficient: the hidefix Index now carries everything xarray needs (attributes, dimension names, dtypes), can be serialized to a compact portable blob, and can read its chunks over S3 via the reader from #47 — so a dataset in object storage opens through xr.open_dataset with no netCDF4/HDF5 fallback and no whole-file download. Reads stay lazy, per-variable, chunk-range.

Motivation: we aggregate MERRA-2 granules on AWS Lambda (jbbutler/antarctic_AR_dataset#2 is the consuming pipeline). An orchestrator indexes each granule once; workers receive the serialized index and fetch only the byte ranges their event footprint needs. Whole-granule fetches cost ~5× the transfer of range reads at our scale, and requiring libhdf5+netCDF4 in the worker layer is exactly what hidefix exists to avoid. Everything here is generic hidefix surface, though — nothing is specific to that pipeline.

The work is one PR in three cumulative parts.


Part 1 — metadata into the index; netCDF4 dependency dropped

Previously python/hidefix/xarray.py opened the file a second time with netCDF4 purely to read attributes and dimension names. Now the index captures them at build time:

  • src/idx/attributes.rs (new): group/dataset attribute capture and netCDF dimension-name resolution (DIMENSION_LIST dereferencing, dimension-scale self-naming, phony_dim_N fallback), plus per-dataset metadata (DatasetMeta).
  • Fixed-length string attributes are read via raw H5Aread FFI holding hdf5_sys::LOCK — libhdf5 has no fixed↔vlen conversion path (h5py converts in software), and unguarded concurrent attribute reads SIGABRT in H5SL. Trailing whitespace is preserved as stored.
  • AttributeValue carries the numeric byte-width (Int(i64, u8), Floats(Vec<f64>, u8), …) so a float32 scale_factor round-trips as float32 and xarray's CF decoding keeps the decoded dtype (widening to f64 changed decoded values by ~2e-7 and the dtype — caught in review, with a regression test asserting decoded dtype and packed scale_factor/add_offset behavior).
  • Index gains #[serde(default)] attribute/metadata fields, exposed to Python (Index.attributes, dataset_attributes, dataset_dims), returned as numpy-typed scalars/arrays matching the stored width.
  • python/hidefix/xarray.py builds the store from the index alone; the netCDF4 import is deleted and the netCDF4 dependency removed from pyproject.toml. An engine-equality test asserts hidefix == netCDF4 backends variable-by-variable (values, dtypes, attrs) on the test files.

Part 2 — serializable index for Python

  • src/idx/serialized.rs (new): a versioned wrapper — magic b"HFXI", u32 LE version, then a flexbuffers-encoded SerializedIndex { source_path, size, mtime, index }. Flexbuffers (already an optional dep) rather than bincode because the index now has #[serde(default)] fields and bincode's positional format can't tolerate them; truncated-input and bad-magic failures are distinguished.
  • Python surface: Index.save(path), hidefix.load_index(path), Index.to_bytes() / from-bytes (file IO releases the GIL).
  • xr.open_dataset(path, engine="hidefix", index=...) accepts an Index or a serialized path, with index_fingerprint='verify'|'ignore': path identity is always enforced (an index for file A never silently reads file B — review proved the 'ignore' variant of that was reachable and wrong); 'ignore' skips only size/mtime staleness (for copies/renames of identical content).

Part 3 — S3 through Python and the xarray backend

  • Builds on the Add reader for S3 (picks up #26) #47 S3Reader. s3 is added to the wheel's feature set ([tool.maturin] features) but not to the cargo python feature, so cargo consumers of python don't pull the S3 stack; the PyO3 glue is #[cfg(feature = "s3")].
  • hidefix.S3Source(bucket, key, *, region, endpoint, anonymous, access_key, secret_key, session_token, path_style) wraps a configured s3::Bucket; Dataset.with_s3(source) reads through a new DatasetD::as_s3_reader mirroring as_reader. Fetches run outside the GIL touching only Rust state. Credential modes are mutually exclusive: anonymous=True combined with explicit keys raises at construction (review found the silent anonymous-wins path). A custom endpoint (minio etc.) implies path-style addressing unless overridden. S3Source's Debug is a manual redacting impl (rust-s3's Bucket Debug prints the secret key).
  • xr.open_dataset("s3://bucket/key.nc4", engine="hidefix", index=..., region=..., endpoint=..., anonymous=...): the index is required for s3:// URIs (building one needs a local file); S3 kwargs are rejected for local paths. Reads stay lazy per-variable; opening fetches no data variables (xarray does eagerly load dimension-coordinate variables to build its indexes — documented, and the laziness test asserts exactly {coords at open} then exactly one range read on slicing).
  • S3 path identity: samefile is meaningless remotely, so 'verify' string-matches the index's source_path against the object key or full URI; 'ignore' skips it. This is deliberately asymmetric with local (where identity is unconditional): the common flow is index-locally-then-upload, in which case the embedded path is the local build path and 'ignore' is the documented escape hatch. Size/mtime staleness is not checkable remotely.

Tests

Suite Result
cargo test --lib (default / flexbuffers / s3 / python / python,s3) 48 / 53 / 52 / 58 / 62
pytest tests/python (no S3 endpoint) 27 passed, 18 skipped (re-verified 2026-08-08)
pytest tests/python with minio (HIDEFIX_S3_ENDPOINT) 35 passed, 10 skipped
cargo test --features s3 --test read_s3 vs minio 8 passed

The S3 pytest module follows tests/read_s3.rs gating (HIDEFIX_S3_ENDPOINT, skip otherwise; boto3 is importorskip'd for the upload fixture — it is not a dependency). No CI workflow changes are included; let me know if you'd like a minio setup wired in.

The clippy and Audit jobs go red on this PR, and both are pre-existing on main — nothing here introduces them. The last CI run on main at 94ae5da, the exact base of this branch, shows the same two jobs failing (run 28581563694). Checked side by side locally:

  • clippy (--workspace -- -D warnings): 30 errors on this branch and 30 on main, at the same twelve sites with the same lints — infallible TryFrom (extent.rs), manual is_multiple_of (filters/shuffle.rs, reader/dataset.rs ×6), contains() vs iter().any() (idx/dataset/dataset.rs), and hiding a lifetime (idx/index.rs ×2). Line numbers differ only where this branch inserted code above them (idx/index.rs:132→180, idx/dataset/dataset.rs:300→319). None land in the files this PR adds (attributes.rs, serialized.rs) or in python.rs. They look like newer-toolchain lints that landed after the last green run rather than anything rotten.
  • cargo audit: the same 11 advisories on both, from the shared dependency tree (bytes, openssl, pyo3, quick-xml, ring, …). main additionally warns on proc-macro-error (RUSTSEC-2024-0370), which this branch's lockfile drops — so the advisory count goes down by one, not up. The S3 stack adds no new advisory.

Happy to fix the clippy lints in a this or separate PR if you'd like them cleared — they're mechanical, your call.

Notes for review (deliberate calls & known limitations)

  • Error types are uneven: some failure paths raise RuntimeError (Rust-side errors crossing PyO3) where ValueError/KeyError might be more pythonic. Left as-is to keep the diff reviewable; happy to normalize if you have a preferred convention.
  • Encoding surface shrinks: without netCDF4 the backend no longer reports unlimited_dims, filter/compression details, or original chunk shapes in encoding. Values/dtypes/attrs are asserted equal to the netCDF4 engine; round-trip-writing workflows that relied on those encoding keys would notice.
  • Serialized-index back-compat: the HFXI wrapper is new and versioned (v1). Raw bincode blobs of the pre-existing Index serde remain possible in Rust but can't tolerate future field additions (positional format), which is why the wrapper uses flexbuffers.
  • No repath API: an index's embedded source_path can't be rewritten after build — 'ignore' is the intended mechanism for moved/uploaded files, and a repath() could be added later if wanted.
  • Dependencies: hdf5-metno 0.12.5 + netcdf 0.12.0 (the pairing is forced by the links = "hdf5" constraint), direct hdf5-sys dep for the raw attribute FFI, flexbuffers promoted into the python feature. The wheel builds with no-default-features (rustls stack only where needed).
  • Wheel portability risk (S3): rust-s3 → rustls 0.23 defaults to aws-lc-rs, so aws-lc-sys (C, needs cmake; NASM on Windows) enters the wheel build. Verified on macOS arm64; if wheel CI objects on other platforms, dropping "s3" from [tool.maturin] features is a one-line revert that keeps everything else.
  • Pre-existing, not addressed: tokio is a non-optional dependency (came in with Add reader for S3 (picks up #26) #47) even for non-S3 builds; guess_can_open won't auto-detect extensionless S3 keys (explicit engine="hidefix" works); a session-token-only S3Source (no access/secret key) falls through to ambient credentials.

espg added 25 commits July 14, 2026 18:00
@espg espg changed the title xarray backend from the index alone: attribute capture, serializable indexes, and S3 reads (drops the netCDF4 dependency) xarray backend from index alone: attribute capture, serializable indexes, and S3 reads (drops the netCDF4 dependency) Aug 9, 2026
@gauteh

gauteh commented Aug 9, 2026

Copy link
Copy Markdown
Owner

Wow... impressive! This will possibly be quite useful for some of our setups as well, maybe it's whats needed to avoid using opendap.

It's pretty immense, and I can see from the branch name that you have used claude. I am not against that. Did you go throughly through it and test it? That said, in the future, it may be a good idea to submit in smaller steps, otherwise it is difficult to review.

I have skimmed through and will try to test. Why do you need a SerializedIndex? Is it interchangeable with the regular index?

- Add #[allow(clippy::infallible_try_from)] to array tuple TryFrom impls
  in extent.rs (changing to From would conflict with the tuple macro)
- Use .is_multiple_of() in shuffle.rs and reader/dataset.rs
- Use .contains() instead of iter().any() in idx/dataset/dataset.rs
- Add explicit '_ lifetime to DatasetD in idx/index.rs

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@gauteh gauteh mentioned this pull request Aug 10, 2026
- Upgrade ndarray 0.16 -> 0.17 to unify with numpy 0.29's resolved version
- Fix dataset_attributes lifetime annotation
- Fix read_py_array borrow conflict by scoping dst before moving array
- Fix apply_fill_value_impl: add Error: Debug bound, replace
  par_mapv_inplace (unavailable on views) with Zip::par_for_each

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
gauteh and others added 4 commits August 10, 2026 15:24
- PyArray::zeros instead of PyArray::new (no uninitialized memory)
- a.readwrite().as_slice_mut() instead of unsafe as_slice_mut
- arr.readwrite().as_array_mut() instead of unsafe as_array_mut
- #[pyclass(from_py_object)] on S3Source to silence deprecation warning

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Replaces removed prepare_freethreaded_python() + Python::with_gil()
with their pyo3 0.29 equivalents.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Methods dataset(), dataset_attributes(), dataset_dims(), and datasets()
all accept an optional group parameter, but were missing the
#[pyo3(signature = (...))] attribute with default=None, causing Python
callers to get a TypeError when not passing the argument.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@gauteh
gauteh force-pushed the claude/xarray-s3-engine branch from 7e32d35 to 154adf8 Compare August 11, 2026 11:16
@espg

espg commented Aug 11, 2026

Copy link
Copy Markdown
Contributor Author

Thanks! Glad it might be useful beyond our pipeline — the opendap-replacement case (serialize an index server-side, hand it to clients that only ever do range reads) is exactly the shape this was built for. Sorry for the size on this; I generally try and keep these smaller in terms of both lines and commits, but the three parts were cumulative (part 3 needs 2's serialization, 2 needs 1's self-sufficient index), which is why they shipped together. Next time I'll split them up so that it's easier to go thru by each phase.

On SerializedIndex vs the regular Index: yes, fully interchangeable — it is not a second index type, just an at-rest envelope. Nothing reads through it: deserializing yields the ordinary in-memory Index (load_index() returns a plain Index), and the store/reader/xarray paths only ever consume Index. The envelope exists to give the persisted form three things the bare serde blob can't have:

  1. A version headerb"HFXI" + u32 version ahead of the payload, so the version is checked before the payload encoding is interpreted, and bad-magic vs truncated-input fail distinctly. Index has been serde-serializable all along (your README bincode example), but a raw blob has no format identity to evolve against.
  2. Schema tolerance — this PR adds #[serde(default)] fields to Index (attributes, dataset metadata), and bincode's positional format breaks on added fields. The envelope's payload is flexbuffers (self-describing), so old persisted indexes stay readable as Index grows. Any future field additions are covered the same way.
  3. A staleness fingerprintsource_path/size/mtime captured at build time, checked at open (index_fingerprint='verify'|'ignore'). An index at rest can outlive its file or get pointed at the wrong one; the in-memory type never has that problem, which is why these fields live on the envelope and not on Index.

For us an orchestrator indexes a granule once, and Lambda workers receive the blob and fetch only the byte ranges they need. But the index crosses process/machine/time boundaries-- it's cheaper for us to reprocess data rather than hold it longer than three or four months... which is what makes versioning and staleness detection stop being optional. But nothing in it is pipeline-specific.

On review and testing — The code was agent-written, and I put it through an adversarial review cycle before opening the PR: a fresh-context reviewer hunting for real defects, with findings folded back as fixes plus regression tests. That process caught genuine bugs — the float32 scale_factor widening that changed CF-decoded values, a reachable path where index_fingerprint='ignore' could silently read the wrong file (path identity is now unconditional locally), and a silent anonymous-wins credential precedence in S3Source. The test matrix in the PR body was run locally, not inferred: cargo across five feature combinations, pytest with and without a live minio endpoint, and the engine-equality test asserting hidefix == netCDF4 backends variable-by-variable.

What I haven't done is test it on our AWS lambda fleet deployment. That environment is a hassle to build for; it pulls in multiple libraries, including our helper library h5coro-hidefix, each with their own CI pipelines that assume things are fetch-able using cargo/pip to build, and then the lambda deployment itself is built assuming binary zips are uploaded to s3 for the lambdas to clone against. So for us, it's meant close reviews and testing locally-- but we usually have to followup with another PR and deployment cycle if there's something unexpected that shows up from the cloud deployed runs. That said, we do a decent job mocking s3 with the local testing, so the follow on 'fixup' PRs tend to be both rare and small. Let me know on the minio CI offer too, if you want S3 coverage wired into your workflows.

@gauteh

gauteh commented Aug 12, 2026

Copy link
Copy Markdown
Owner

Ok, thanks. I think having minio tests sounds like a good idea.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants