xarray backend from index alone: attribute capture, serializable indexes, and S3 reads (drops the netCDF4 dependency) - #50
Conversation
|
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>
- 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>
- 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>
7e32d35 to
154adf8
Compare
|
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
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 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. |
|
Ok, thanks. I think having minio tests sounds like a good idea. |
This PR makes the xarray backend self-sufficient: the hidefix
Indexnow 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 throughxr.open_datasetwith 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.pyopened 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_LISTdereferencing, dimension-scale self-naming,phony_dim_Nfallback), plus per-dataset metadata (DatasetMeta).H5AreadFFI holdinghdf5_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.AttributeValuecarries the numeric byte-width (Int(i64, u8),Floats(Vec<f64>, u8), …) so afloat32scale_factorround-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 packedscale_factor/add_offsetbehavior).Indexgains#[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.pybuilds the store from the index alone; the netCDF4 import is deleted and thenetCDF4dependency removed frompyproject.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 — magicb"HFXI", u32 LE version, then a flexbuffers-encodedSerializedIndex { 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.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 anIndexor a serialized path, withindex_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
S3Reader.s3is added to the wheel's feature set ([tool.maturin] features) but not to the cargopythonfeature, so cargo consumers ofpythondon'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 configureds3::Bucket;Dataset.with_s3(source)reads through a newDatasetD::as_s3_readermirroringas_reader. Fetches run outside the GIL touching only Rust state. Credential modes are mutually exclusive:anonymous=Truecombined with explicit keys raises at construction (review found the silent anonymous-wins path). A customendpoint(minio etc.) implies path-style addressing unless overridden.S3Source'sDebugis a manual redacting impl (rust-s3'sBucketDebug prints the secret key).xr.open_dataset("s3://bucket/key.nc4", engine="hidefix", index=..., region=..., endpoint=..., anonymous=...): the index is required fors3://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).samefileis meaningless remotely, so'verify'string-matches the index'ssource_pathagainst 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
cargo test --lib(default /flexbuffers/s3/python/python,s3)pytest tests/python(no S3 endpoint)pytest tests/pythonwith minio (HIDEFIX_S3_ENDPOINT)cargo test --features s3 --test read_s3vs minioThe S3 pytest module follows
tests/read_s3.rsgating (HIDEFIX_S3_ENDPOINT, skip otherwise; boto3 isimportorskip'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
clippyandAuditjobs go red on this PR, and both are pre-existing onmain— nothing here introduces them. The last CI run onmainat94ae5da, the exact base of this branch, shows the same two jobs failing (run 28581563694). Checked side by side locally:--workspace -- -D warnings): 30 errors on this branch and 30 onmain, 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()vsiter().any()(idx/dataset/dataset.rs), andhiding 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 inpython.rs. They look like newer-toolchain lints that landed after the last green run rather than anything rotten.bytes,openssl,pyo3,quick-xml,ring, …).mainadditionally warns onproc-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)
RuntimeError(Rust-side errors crossing PyO3) whereValueError/KeyErrormight be more pythonic. Left as-is to keep the diff reviewable; happy to normalize if you have a preferred convention.unlimited_dims, filter/compression details, or original chunk shapes inencoding. Values/dtypes/attrs are asserted equal to the netCDF4 engine; round-trip-writing workflows that relied on those encoding keys would notice.Indexserde remain possible in Rust but can't tolerate future field additions (positional format), which is why the wrapper uses flexbuffers.source_pathcan't be rewritten after build —'ignore'is the intended mechanism for moved/uploaded files, and arepath()could be added later if wanted.hdf5-metno0.12.5 +netcdf0.12.0 (the pairing is forced by thelinks = "hdf5"constraint), directhdf5-sysdep for the raw attribute FFI,flexbufferspromoted into thepythonfeature. The wheel builds withno-default-features(rustls stack only where needed).rust-s3→ rustls 0.23 defaults toaws-lc-rs, soaws-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] featuresis a one-line revert that keeps everything else.tokiois a non-optional dependency (came in with Add reader for S3 (picks up #26) #47) even for non-S3 builds;guess_can_openwon't auto-detect extensionless S3 keys (explicitengine="hidefix"works); a session-token-onlyS3Source(no access/secret key) falls through to ambient credentials.