Zarr output: write gridded and flat Zarr v3 stores - #331
Draft
robinskil wants to merge 1 commit into
Draft
Conversation
Adds a write side to beacon-arrow-zarr, mirroring the NetCDF sinks so queries can be materialized as Zarr instead of only read from it. * `ZarrStoreWriter` builds a Zarr v3 store through zarrs' filesystem store, and packs it into a zip when the caller can only take one file. * `ZarrSink` writes flat output: each column a 1-D array over a shared `obs` dimension, appended as batches arrive so memory stays bounded. * `ZarrNdSink` writes gridded output: the named dimension columns become coordinate arrays and the rest become N-D arrays indexed by them. * `ZarrFormat::create_writer_physical_plan` picks between them based on `ZarrOptions::write_dimensions`. Exposed as the `zarr` and `ndzarr` output formats, and by registering the format factory for writing so `COPY TO ... STORED AS ZARR` works. The HTTP API returns a zip (`application/zip`, `.zarr.zip`) because a response carries one file; `COPY TO` writes a plain directory store. Gridding does not reuse the NetCDF nd path's `UniqueValuesExec` + sort + per-type scatter macros. Cell positions come from a dense rank computed with stock Arrow kernels (`sort_to_indices` -> `cmp::distinct` -> accumulate) feeding `arrow::compute::take`, which is type-agnostic and needs no extra plan nodes: `DataSinkExec` already requires a single input partition. Only Zarr v3 is written. zarrs models v2 as a read-side format that is converted to v3 in memory and has no v2 writer, and Beacon's own reader discovers v3 stores only. `ZarrOptions::zarr_version` exists so adding v2 later is not a breaking change. Arrow nulls in float columns are written as NaN, the conventional CF / zarr missing-value marker, and read back as NaN rather than null. No `_FillValue` attribute is emitted: it cannot express NaN as a JSON number, and the sentinels available for the other types collide with real data.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #331 +/- ##
==========================================
+ Coverage 75.69% 76.15% +0.46%
==========================================
Files 303 308 +5
Lines 40972 42523 +1551
==========================================
+ Hits 31012 32382 +1370
- Misses 9960 10141 +181
🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds a write side to
beacon-arrow-zarr, mirroring the NetCDF sinks so queries can be materialized as Zarr instead of only read from it. Both gridded (multi-dimensional) and flat output are supported.What's here
ZarrStoreWriter(writer.rs) builds a Zarr v3 store through zarrs' filesystem store, and packs it into a zip when the caller can only take one file.ZarrSink(sink.rs) writes flat output: each column a 1-D array over a sharedobsdimension, appended as batches arrive so memory stays bounded.ZarrNdSinkwrites gridded output: the named dimension columns become coordinate arrays and the rest become N-D arrays indexed by them.ZarrFormat::create_writer_physical_planpicks between them based onZarrOptions::write_dimensions.Exposed as the
zarrandndzarroutput formats, plus a writable format-factory registration soCOPY TO ... STORED AS ZARRworks. Mirrored in the Python CLI and TS SDK.Delivery: zip vs directory
A zarr store is a directory, but an HTTP response carries one file. The query API therefore returns a zip (
application/zip,.zarr.zip), readable viazarr-python'sZipStoreor fsspec; entries areStored(uncompressed) since chunk data is already codec-compressed.COPY TOhas no such limit and writes a plain directory store.Three decisions worth reviewing
Gridding does not reuse the NetCDF nd path. That path adds
UniqueValuesExec+SortExec+SortPreservingMergeExecand scatters with ~300 lines of per-type macros. Here, cell positions come from a dense rank computed with stock Arrow kernels (sort_to_indices→cmp::distinct→ accumulate) feedingarrow::compute::take— type-agnostic, ~60 lines, and no extra plan nodes, sinceDataSinkExecalready requires a single input partition. This is a deliberate divergence from the netCDF precedent.Only Zarr v3 is written. zarrs models v2 as a read-side format converted to v3 in memory and ships no v2 writer (
MetadataConvertVersionoffers onlyDefault/V3), and Beacon's own reader discovers v3 stores only. Emitting v2 would mean hand-rolling.zgroup/.zarrayJSON and C-order chunks outside zarrs, untestable through our own reader.ZarrOptions::zarr_versionexists so adding it later is non-breaking.Float nulls become NaN, not null. NaN is the conventional CF/zarr missing-value marker and xarray reads it as missing, but a round-trip through Beacon's reader surfaces NaN rather than null. No
_FillValueattribute is written: it cannot express NaN as a JSON number, and the sentinels available for other types (0/false/"") collide with real data, so advertising them would mask real values.Testing
936 workspace tests pass, 0 failures. New coverage:
/-separated, so the archive unpacks into a usable store.zarr,ndzarr(asserts shape[2,2]anddimension_names), andCOPY TO.integration-tests/test_nd_formats.py(unzip →read_zarr) are not yet run — they need the container harness. The zip→reader path is verified only at the Rust level so far.Loose end
beacon_core::query::CopyTohas no Zarr variant. It has no references outside its own definition, so it looks like dead code — happy to add it if it's actually live.