feat(raster-zarr): sedona-raster-zarr crate + sedona-zarr plugin#858
Merged
Conversation
Adds a sedona-raster-zarr crate that maps Zarr groups onto SedonaDB's
N-D raster Arrow schema, plus an sd_read_zarr DataFusion UDTF that
exposes it to SQL.
## sedona-raster-zarr
Two entry points emit one raster row per chunk position in the group's
chunk grid, with one band per array:
* group_to_indb_rasters(group_uri) — eager: fetches every chunk's bytes
into the Arrow data column.
* group_to_outdb_rasters(group_uri) — URI-only: emits chunk-anchor URIs
in each band's outdb_uri (<store>#array=<path>&chunk=<i0>,...).
Built on the zarrs 0.23 crate, chosen over GDAL MultiDim because zarrs
reads Zarr v3 sharding and vlen-utf8 coordinate variables.
Modules:
* source_uri — group-URI normalisation (file:// + bare path; cloud
schemes error pending the resolver work) and chunk anchor build/parse.
* geozarr — proj:wkt2 / proj:projjson / proj:epsg precedence and
spatial:transform / spatial:dims parsing from group attributes.
* dtype — zarrs DataType ↔ BandDataType via the type-erased is::<T>()
downcast (zarrs 0.23 wraps Arc<dyn DataTypeTraits>, not an enum).
* loader — opens FilesystemStore + Group, enumerates child arrays
(sorted by path for deterministic band order), validates shared chunk
grid / chunk shape / dim names, walks the chunk grid row-major,
computes per-chunk transforms by translating the group affine along
the spatial axes.
Tests: 31 in-module unit tests + 4 integration tests that build a real
Zarr group on disk via zarrs::ArrayBuilder and verify byte-exact pixel
content, per-chunk transforms, anchor URI format, and error paths.
## sd_read_zarr UDTF
DataFusion table function exposing the loader to SQL. Mirrors
sd_random_geometry's shape (Function + Provider + Exec) and registers
in SedonaContext::new_from_context.
SELECT raster FROM sd_read_zarr('file:///path/to/datacube.zarr');
SELECT count(*) FROM sd_read_zarr(
'file:///path/to/datacube.zarr',
'{\"mode\": \"outdb\", \"rows_per_batch\": 256}'
);
Returns a single-column table raster: Raster. All existing RS_* UDFs
operate on the column unchanged.
JSON options:
* mode: \"indb\" (default) or \"outdb\". Phase 1 defaults to InDb so
byte-reading kernels work end-to-end. Flips to \"outdb\" once a
format-keyed OutDb dispatcher registers a zarr loader.
* rows_per_batch: chunks per RecordBatch (default 1024).
* num_partitions: scan partitions (default 1). Anything other than 1
errors — round-robin chunk partitioning lands with the resolver work.
ZarrChunkProvider builds the StructArray once at plan time. scan()
wraps the inner exec in ProjectionExec so empty / partial projections
(e.g. SELECT count(*)) match the requested physical schema.
ZarrChunkExec slices the StructArray into rows_per_batch-sized
RecordBatches. Bounded, Incremental, single-partition.
6 in-module integration tests build a Zarr fixture and exercise the
full SQL pipeline.
## Out of scope (follow-up PR)
* OutDb byte resolver (the function the byte-loading hook calls to
fetch zarr chunks) — depends on a format-keyed multi-backend
dispatcher.
* Default mode flip from \"indb\" to \"outdb\".
* Cloud storage backends (zarrs_object_store) and the async runtime
story they require.
* Round-robin chunk partitioning past num_partitions=1.
* CREATE EXTERNAL TABLE ... STORED AS ZARR.
Adds a `con.funcs.table.sd_read_zarr(uri, *, mode, rows_per_batch,
num_partitions)` Python method that mirrors the SQL UDTF added in the
companion Rust commit. Same shape as the existing
`sd_random_geometry` wrapper: builds the JSON options object,
filters out None entries, and runs the SQL UDTF through
`self._ctx.sql`.
Python test uses `pytest.importorskip("zarr")` so it cleanly skips
when the optional `zarr` library isn't installed in the test
environment. Builds a 2×2 UInt8 Zarr group on disk via the official
`zarr` Python lib and exercises both the default-options form and
the `rows_per_batch` override.
james-willis
force-pushed
the
jw/sd-read-zarr
branch
from
e9cb19f to
3fd6420
Compare
james-willis
commented
May 19, 2026
Round of review feedback on the sd_read_zarr loader and UDTF surface: - Drop "Phase 1" references from code, error messages, and docstrings; describe behaviour in absolute terms. - Bound-check the chunk-grid product so a malformed or hostile Zarr can't drive RasterBuilder capacity to overflow / OOM. - Always populate `outdb_uri` and `outdb_format` on every band as provenance metadata; InDb vs OutDb stays discriminated by `data.is_empty()`. - Fall back to xarray's `_ARRAY_DIMENSIONS` attribute when an array lacks the first-class `dimension_names` field (covers Zarr v2 and any v3 file that uses the xarray convention). - Auto-skip 1-D arrays at selection time (typical xarray coord variables) — they can never produce a valid raster band, so reading them would always fail downstream at spatial-dim resolution. - Add an `arrays = [...]` option (Python kwarg + UDTF JSON key) to read a named subset of arrays; unknown names and 1-D arrays both error early with clear messages pointing at the option. - Error loudly when a group declares a CRS but no `spatial:transform` attribute — the high-confidence-bug case where falling back to the identity transform would silently produce wrong spatial-join results. - Warn-log when both CRS and transform are absent and we fall back to pixel-coordinate identity, so spatial-join surprises are debuggable from logs. - Flip the InDb/OutDb selector from `mode: "indb"|"outdb"` to `indb: bool` end-to-end (Python kwarg, SQL JSON key, internal options struct). The schema-level discriminator is binary, so a boolean is the more honest UI. Tests cover every new branch: auto-skip, explicit array filter, 1-D rejection, _ARRAY_DIMENSIONS fallback, CRS-without-transform error, indb-bool plumbing through SQL.
zarrs 0.23 deprecated `store_chunk_elements::<T>` in favour of `store_chunk()`. Migration is mechanical for u8 element types since the byte representation is identity. Drops the `#[allow(deprecated)]` annotations the fixtures carried.
`load_eager` more accurately describes the user-facing knob: the boolean controls whether chunk bytes are materialised eagerly or left as chunk-anchor URIs, not which storage class the resulting raster lives in. Renames the JSON key (`indb` → `load_eager`), the Python kwarg, the internal `ZarrReadOptions` field, and updates docstrings to reflect the long-term plan: once an async `RS_EnsureLoaded` UDF lands, the planner will auto-inject it over the scan output rather than the loader doing eager fetching at plan time. The user-facing flag name doesn't change when that swap happens.
paleolimbot
reviewed
May 20, 2026
Member
paleolimbot
left a comment
paleolimbot
left a comment
There was a problem hiding this comment.
This is very cool!
The main comment I have is the pattern for iterating through the chunks at the start...if we have to have it be one big forward loop (i.e., we're sure there are never enough chunks that we'll meaninfully need to inspect only a few at a time), we at least need to have the byte loading (1) be async or (2) pullable in tiny chunks.
I would also really like this to not be built in to the main sedonadb Python package. I'm certainly willing to let it go a tiny bit for progress' sake but we're hitting the point where we've got too much built in by default (pointcloud and GPU are on my radar too for how we have to separate the build). If you have it in you to work through exposing this in python/sedonadb-zarr as a simple object that implements __arrow_c_stream__ it would help me out because (1) then I don't have to do it later and (2) it will help iron out some of the missing functionality that we'll need to separate other stuff like pointclouds, which I would otherwise have to do later. An obvious hiccup in that approach is that RecordBatchReader is not async (so you'll just have to serve batches bounded by some byte transfer size to start).
- Narrow the blosc Windows cfg gate from "any Windows" to "Windows MinGW only" (`target_env = "gnu"`). MSVC Windows users (Python wheels) now get blosc-compressed Zarr support back; only the R rtools45 MinGW build path still disables it. - Drop the hardcoded `rows_per_batch = 1024` default in `ZarrChunkExec`. When the UDTF option is unset, defer at execute time to `SessionConfig::batch_size` so users tune via the standard knob instead of the loader baking in its own constant. - Python test: lift `numpy` and `pytest` imports to module level (only `zarr` remains `importorskip`-gated); materialise the resulting DataFrame via `to_arrow_table()` and assert on the Raster struct row instead of just `count()`.
The loader had two near-duplicate code paths (`group_to_indb_rasters` and `group_to_outdb_rasters`) gated by a `Mode` enum. The InDb path fetched every chunk's bytes on a single thread at plan time — exactly the antipattern that the async `RS_EnsureLoaded` resolver will replace. Rather than ship the bad path and tear it out later, collapse to a single OutDb-producing entry point now. - New single entry: `group_to_rasters(uri, arrays)`. Always emits chunk-anchor URIs in `outdb_uri`; `data` is empty. - `Mode::InDb` / `Mode::OutDb` enum and the `match mode` branches in `build_rasters` are gone. - `retrieve_chunk_bytes` is kept (the only pixel-byte read primitive in the crate; the future async resolver will want it as a starting point) but is now `#[allow(dead_code)]` and exercised by a direct unit test rather than indirectly via the deleted InDb integration test. - `ArrayInfo::array` field dropped — only the InDb path used it. - `sd_read_zarr()` default for `load_eager` flips from `true` to `false`; explicit `load_eager = true` errors with a message pointing at the future `RS_EnsureLoaded` resolver. - Tests rewritten: `indb_round_trip_emits_one_row_per_chunk_position` is folded into a single `round_trip_emits_one_row_per_chunk_position_with_outdb_anchors` test that asserts on metadata and anchor URIs. Byte-content assertions removed. New UDTF test verifies `load_eager = true` errors cleanly. Plan-time resident size for a Zarr scan drops from `N_chunks × chunk_size` (potentially GBs) to roughly `N_chunks × few-hundred-bytes`. The remaining structural concern — build-everything-upfront vs streaming emission — is unchanged in this PR and tracked as separate follow-up work.
…nalFormatSpec + UDTF migration Moves all Zarr functionality out of the `sedona` crate so the main package no longer carries a zarr dep. Three structural changes: 1. **Streaming `ZarrChunkReader: RecordBatchReader`** — replaces the eager `group_to_rasters -> StructArray` function. Opens the group + parses metadata once in `try_new`; per-batch work is just transform arithmetic and anchor URI formatting. `next()` walks `batch_size` chunk positions of the chunk grid and emits one `RecordBatch`. The chunk anchors are still OutDb-style (empty `data`, populated `outdb_uri`); pixel-byte resolution is the future async resolver's job. 2. **`ZarrFormatSpec: ExternalFormatSpec`** — wraps the reader with the standard SedonaDB datasource API so `con.read_format(spec, uri)` works the same shape as PyogrioFormatSpec. Lives in `sedona-raster-zarr` gated under a `zarr` feature (default on). 3. **UDTF migrates from `sedona` to `sedona-raster-zarr`.** The `ZarrReadFunction` / `ZarrChunkProvider` / `ZarrChunkExec` code moves to `sedona-raster-zarr/src/udtf.rs` (also `zarr`-feature-gated). The exec is now genuinely streaming — it constructs a fresh `ZarrChunkReader` in `execute()` and adapts it via `RecordBatchStreamAdapter`. The new `register(ctx: &SessionContext)` helper is what the Python plugin package will call. The `sedona` crate's `zarr_read` module is deleted; the `sedona_raster_zarr` workspace dep is dropped; the `SedonaContext::new_from_context` UDTF registration is removed. The Python `sd_read_zarr` wrapper method in `functions/table.py` and the `test_read_zarr` test go too — they get re-introduced in the new `sedonadb-zarr` plugin package (next commit). Plan-time validation: `ZarrChunkProvider::try_new` opens the reader once for validation and drops it, so URI / metadata / arrays-filter errors still surface at `ctx.sql(...).await` rather than at collect time. Cost is one extra group-open per scan. Result: `sedona` crate has zero zarr code; `sedona-raster-zarr` is fully self-contained with both the format-spec and UDTF integration surfaces under `feature = "zarr"`. The plugin Python package wires both via `sedonadb_zarr.register(con)`.
Introduces `python/sedonadb-zarr` as a separate maturin-built Python package. After `sedonadb_zarr.register(con)`, the `sd_read_zarr` SQL UDTF is attached to the session; before that call (or without importing the package at all) the main `sedonadb` package has zero zarr awareness. Package contents: - `python/sedonadb-zarr/src/lib.rs` — PyO3 binding exposing a single `register_udtf(internal_ctx)` function that calls `sedona_raster_zarr::register_udtf` on the wrapped DataFusion `SessionContext`. The plugin imports `python/sedonadb`'s PyO3 classes via the `::_lib` extern crate (sedonadb's rustc crate name per its `lib.name = "_lib"` override that maturin requires). - `python/sedonadb-zarr/python/sedonadb_zarr/__init__.py` — Python `register(con)` helper. Locates the `_impl: InternalContext` on a sedonadb `Context` and forwards to the Rust extension. - `python/sedonadb-zarr/pyproject.toml`, `Cargo.toml`, `README.md` — build scaffolding mirroring `python/sedonadb`'s layout. Required upstream change: - `python/sedonadb/Cargo.toml` gains `rlib` alongside `cdylib` in `crate-type` so plugin crates can link against it as a Rust dependency. - `python/sedonadb/src/lib.rs` makes the `context` module `pub` so plugins can extract `InternalContext` from Python objects via PyO3's downcast machinery. The other modules stay private. - Workspace root `Cargo.toml` adds `python/sedonadb-zarr` to `[workspace] members` and exposes `sedonadb` as a workspace path-dep. Tests: - `python/sedonadb-zarr/tests/test_zarr.py` covers the smoke path (register + SQL works), the negative case (SQL fails before register), and idempotence (double-register is fine). - `python/sedonadb/tests/test_no_zarr_by_default.py` is the architectural regression test: it lives in the main package and asserts `sd_read_zarr` is unknown without the plugin imported. Catches accidental re-bundling. ExternalFormatSpec is not yet exposed through the Python plugin — the Rust impl exists in `sedona-raster-zarr` and is callable from Rust code; surfacing it via `con.read_format(ZarrFormatSpec(), uri)` requires wiring `ZarrChunkReader` through the Arrow C Stream FFI and is left as a follow-up.
Completes the second user-facing surface promised in the plugin
design. After `sedonadb_zarr.register(con)`, both work:
con.sql("SELECT * FROM sd_read_zarr(...)")
con.read_format(sedonadb_zarr.ZarrFormatSpec(), uri)
Implementation:
- New PyO3 class `PyZarrChunkReader` in the plugin's Rust extension
wraps `sedona_raster_zarr::ZarrChunkReader` and exposes
`__arrow_c_stream__` via `arrow_array::ffi_stream::FFI_ArrowArrayStream`.
The reader is consumed on first `__arrow_c_stream__` call (matches
pyarrow's `RecordBatchReader` convention).
- Python `ZarrFormatSpec(ExternalFormatSpec)` subclass in the plugin
package wraps `PyZarrChunkReader`. Supports `with_options` for
`load_eager` and `arrays`; `open_reader` returns the Rust-backed
reader; default `infer_schema` works because the reader exposes
schema via the Arrow C stream protocol.
Tests cover the new surface end-to-end: read via `con.read_format`,
options threading, and the `load_eager=True` rejection.
Per Dewey's review on test_funcs.py:58 — strengthen the format-spec roundtrip test to actually materialise and inspect the raster row. `arrow_tab["raster"][0].as_py()` returns a dict with the raster struct's fields (transform, bands, etc.). Asserts on the top-level shape (transform + bands present), the band-level shape (data is empty for OutDb scans), and that the chunk anchor URI threads through to the band's outdb_uri field. We don't have a great way to inspect Raster structs end-to-end in Python yet, but this is the minimum that proves the FFI bridge delivers a coherent struct value to the user.
james-willis
force-pushed
the
jw/sd-read-zarr
branch
from
6a8276b to
e7638eb
Compare
- pre-commit (ruff format): test_zarr.py had a multi-line read_format call that ruff wanted as one line. - docs (`cargo doc --workspace`): both `sedonadb` and `sedonadb-zarr` originally declared `[lib] name = "_lib"` (maturin requires that for the Python `<pkg>._lib` import path). rustdoc errors with "document output filename collision" because both produce a `_lib` rustdoc output. Rename the plugin's rustc crate to `sedonadb_zarr_lib`; maturin still renames the compiled cdylib to `_lib.<ext>` for the wheel via `module-name`, so the Python import path is unaffected. - python ubuntu-latest: pytest collects from `python/` which now includes `python/sedonadb-zarr/tests/`. The plugin wasn't installed in the test env, so its `import sedonadb_zarr` errored at collection time. Install the plugin alongside sedonadb in the Install step.
james-willis
force-pushed
the
jw/sd-read-zarr
branch
from
e7638eb to
06e479b
Compare
The plugin refactor exposed `python/sedonadb` as an rlib so plugin
crates (e.g. `sedonadb-zarr`) could link it and extract
`InternalContext`. That worked on macOS but broke Linux CI with:
rust-lld: error: duplicate symbol: PyInit__lib
Both `sedonadb` and `sedonadb-zarr` declare `#[pymodule] fn _lib`
(each needs the symbol so its respective Python wheel imports as
`<pkg>._lib`). When the plugin's cdylib statically links sedonadb's
rlib, both `PyInit__lib` symbols collide. macOS's two-level namespace
tolerated it; GNU ld doesn't.
Fix: put sedonadb's `#[pymodule] fn _lib` behind a new
`extension-module` Cargo feature, default-on. Maturin's wheel build
inherits the default and gets the symbol; plugin Rust deps set
`default-features = false` (already done at the workspace path-dep
level) and the pymodule's symbol vanishes from the rlib they link.
Adds a blanket `#![cfg_attr(not(feature = "extension-module"),
allow(dead_code, unused_imports))]` so the helper `#[pyfunction]`s
that are only reachable via the pymodule don't generate warnings in
plugin-rlib builds.
The Status section claimed ExternalFormatSpec was a Rust-only follow-up, but `con.read_format(ZarrFormatSpec(), uri)` was actually shipped in the same PR (PyZarrChunkReader + Python wrapper + tests). Tighten the README — show both surfaces, drop the stale "follow-up" claim, and unwrap paragraphs per GFM's hard-line-break rendering.
The feature was a defensive response to Dewey's "at the very very least feature-flag it" comment, but that was the fallback assuming we *didn't* move the UDTF out of the `sedona` crate. We did — the UDTF and ExternalFormatSpec now live in `sedona-raster-zarr` itself — so the load-bearing concern is addressed. The only consumer of this crate is `sedonadb-zarr`, which uses the full surface. No hypothetical "slim build" wants just the streaming reader without DataFusion. The feature gate added cfg-clutter without real value. Drop the gate, drop the `[features]` table, make the previously- optional deps required. `lib.rs` declares the modules unconditionally.
MATURIN_PEP517_ARGS="--features s2geography" set for sedonadb's install was still in scope for the subsequent pip install of sedonadb-zarr, which doesn't have that feature. Unset before the plugin install.
The previous fix gated the pymodule behind `extension-module` so plugin rlib consumers wouldn't duplicate `PyInit__lib`. The feature was default-on; Cargo's workspace-wide feature unification meant a workspace build (e.g. `cargo test --workspace`, `cargo build --workspace`) forced the feature on for everyone — including the plugin's rlib link of sedonadb — and the duplicate-symbol collision came back. Move `extension-module` out of `default`. Maturin enables it explicitly via `[tool.maturin] features` in `pyproject.toml` for the wheel build. Workspace cargo runs don't request it; sedonadb's rlib in those builds has no pymodule symbol; plugin crates link clean.
After moving `extension-module` out of sedonadb's default features (to
avoid Cargo workspace-feature-unification dragging the pymodule symbol
into plugin rlib links), the wheel builds need to enable it
explicitly. Five workflows / scripts set MATURIN_PEP517_ARGS and need
the feature added:
* `.github/workflows/python.yml` — the test job's editable install
* `.github/workflows/packaging.yml` — docs-and-deploy
* `ci/scripts/wheels-build-{linux,macos,windows}.{sh,ps1}` — wheels
Without it, the produced cdylib doesn't define `PyInit__lib` and any
attempt to `import sedonadb` fails.
cibuildwheel's wheel build (macOS-arm64, linux-arm64) failed cargo
metadata with:
error: failed to select a version for the requirement
`sedonadb = "^0.4.0"`
The workspace dep table had `sedonadb` as a self-reference (the
sedonadb crate exposed via its own path). Only sedonadb-zarr consumed
it; inlining the path-dep on sedonadb-zarr's side and dropping the
workspace entry removes the self-reference. cibuildwheel's isolated
build env (where the workspace root may not be in scope the same way)
no longer chokes.
`cargo build --workspace --all-features` (the build CI job) ignores the workspace dep's `default-features = false` because feature unification merges feature requests across the workspace. With sedonadb's `extension-module` forced on, its rlib's `PyInit__lib` symbol gets linked into sedonadb-zarr's cdylib, which also defines `PyInit__lib` — and the linker errors with duplicate symbol. Rename sedonadb-zarr's pymodule from `_lib` to `_zarr_lib`. The generated PyInit symbol is now `PyInit__zarr_lib`, distinct from sedonadb's `PyInit__lib` regardless of feature unification. This means the Python wrapper imports from `sedonadb_zarr._zarr_lib` instead of `sedonadb_zarr._lib`; the user-facing `register(con)` / `ZarrFormatSpec` API is unchanged.
PyO3's `#[pyclass]` extraction uses a per-cdylib type-id static, so `sedonadb_zarr` couldn't extract `sedonadb`'s `InternalContext` — the two extensions saw distinct types even though they linked against the same Rust definition. The plugin now hands its UDTF across the extension boundary in a `PyCapsule` carrying `Arc<dyn TableFunctionImpl>`, which sedonadb clones into its own `SessionContext`. The test helper `_read_format` bridges to `_impl.read_external_format` because `SedonaContext.read_format` doesn't exist yet; once a public helper lands the body collapses to a one-liner.
…th; drop external_listing_table Two of the three review follow-ups on the directory-format provider: - `SingleObjectExternalTable::try_new` now resolves the `Arc<dyn ObjectStore>` from the session's runtime registry and threads it through to `spec.infer_schema`. The listing path has always done this; the single-object path previously passed `store: None`, which was a silent asymmetry for any future spec whose `infer_schema` needs to peek at the store. Zarr's schema is constant so this didn't bite in practice. - `external_listing_table` had no internal callers after the `external_table` dispatcher landed. Inlined as private `listing_table_provider`. One public entry point, one return type (`Arc<dyn TableProvider>`).
paleolimbot
reviewed
May 21, 2026
Member
paleolimbot
left a comment
paleolimbot
left a comment
There was a problem hiding this comment.
Thank you for working on this, and especially thank you for going with me on the plugin concept.
I think we can solve the linking issues with by dropping the table function. This is a trivial table function (in the sense that it is not slow in the outdb case and never particularly large) and we can make a system like the ExternalFormatSpec (i.e., just routing a RecordBatchReader, which already can be FFIed) for this if the table function is a truly important way to interact with this. I also left some comments that will let the table function hang around without causing a linking issue (I'm pretty sure).
Three independent simplifications from Dewey's review: - `sedona-raster-zarr` no longer depends on the umbrella `datafusion` crate in `[dependencies]`. Replaced with the narrow sub-crates it actually uses: `datafusion-execution`, `datafusion-physical-expr`, `datafusion-physical-plan`. The umbrella moves to `[dev-dependencies]` for tests that build a real `SessionContext`. The `register(&SessionContext)` helper is gone — it had no callers (the plugin constructs `Arc::new(ZarrReadFunction::default())` directly), and removing it was what unblocked dropping the umbrella. - Drop the `load_eager` option from `ZarrFormatSpec` and `sd_read_zarr`. It currently errors anyway; pixel-byte materialisation lands with the async `RS_EnsureLoaded` resolver and we can decide the user- facing shape (option, separate UDF, planner injection) at that point. Removes the only surface defending a not-yet-supported flag. - `retrieve_chunk_bytes` is now `#[cfg(test)]` rather than `#[allow(dead_code)]`. The unit test for it lives in the same file so the implementation doesn't bit-rot before the resolver crate consumes it. The blanket lint-silencer is gone. Also drops `test_no_zarr_by_default.py` — the architectural regression test for zarr-not-in-default sedonadb. The build itself catches reintroduction at this point, so the test is overkill.
Rewires the cross-cdylib UDTF registration path onto datafusion-ffi's `FFI_TableFunction`, matching what `datafusion-python` does for its own plugins. Three big wins: 1. `python/sedonadb-zarr` no longer depends on `python/sedonadb` as a path rlib. The plugin links only `datafusion-catalog`, `datafusion-ffi`, `sedona-raster-zarr` and `pyo3` for its host- facing surface — no transitive sedonadb-* / pyo3 dupe of host internals. Resolves the cibuildwheel "another Python package's private crate" objection from review. 2. mimalloc returns to defaults. All cross-cdylib traffic now flows through Arrow C Data/Stream (release callbacks owned by the producing cdylib), so allocations are always freed in the cdylib that made them. The previous segfault required passing native `Arc<dyn TableProvider>` across cdylibs, which we no longer do. 3. Host-side codec + capsule machinery lives in `python/sedonadb/src/plugin.rs`: a `SessionTaskContextProvider` adapts our `SessionContext` to datafusion-ffi's `TaskContextProvider`, and `create_session_capsule` / `ffi_table_function_from_capsule` pack and unpack the `c"datafusion_logical_extension_codec"` and `c"datafusion_table_function"` capsules the way datafusion-python does. Plugins built against either ecosystem can register on a sedonadb context. Plugin side: `ZarrTableFunction` is the Python-visible spec class with `__datafusion_table_function__(session)`. The host calls it from `InternalContext.register_udtf_capsule(name, spec)`.
… lifetime The `FFI_LogicalExtensionCodec` produced by `create_session_capsule` stored a `Weak<dyn TaskContextProvider>`. The strong `Arc` was a local in the function, so by the time the plugin invoked `FFI_TableFunction::call` the `Weak` couldn't upgrade — DataFusion failed with `FFI error: TaskContextProvider went out of scope over FFI boundary`. Caught by running the plugin smoke tests locally against the FFI rework. Fix: stash one `Arc<dyn TaskContextProvider + Send + Sync>` on `InternalContext` at construction time. Every `register_udtf_capsule` call clones it into the session capsule. The strong ref lives as long as the session, so the FFI codec's `Weak` can always upgrade. All 7 sedonadb-zarr Python tests pass locally (including the mimalloc-on smoke test); 1699 sedonadb tests pass unchanged.
`register_udtf_capsule` and the `__datafusion_table_function__` contract it consumes are 0.4.0 features. Pinning prevents an install that picks up an older sedonadb at resolution time.
The target-gated `[target.'cfg(...)'.dependencies]` block existed for rtools45's pthread conflict when this crate got pulled into the R sedonadb.dll build. R sedonadb has no path into `sedona-raster-zarr` in this PR — the crate is only consumed by the Python sedonadb-zarr plugin — so the gate is dead and the `blosc` feature moves to the unconditional zarrs feature list. R parity is tracked as a follow-up issue.
paleolimbot
reviewed
May 21, 2026
Adds `test_dtype_mapping_roundtrips` covering every arm of
`zarr_to_band_data_type` (bool / int{8,16,32,64} / uint{8,16,32,64} /
float{32,64}). Each variant writes a tiny Zarr group and reads back
two chunks via `sd_read_zarr` — fast sanity check that the dtype
table doesn't silently drift.
Per Dewey's review note on `rust/sedona-raster-zarr/src/dtype.rs:40`.
Both existed solely to support `sedonadb-zarr` linking `sedonadb` as an rlib, which the datafusion-ffi rework eliminated. With no rlib consumer: - `extension-module` Cargo feature on `python/sedonadb` is removed. The `#[pymodule] fn _lib` is unconditional. `[lib].crate-type` is back to just `cdylib`. CI workflows and wheel-build scripts drop the `extension-module` MATURIN_PEP517_ARGS entry. - `python/sedonadb-zarr` renames the rustc crate name and pymodule function from `_zarr_lib` back to the conventional `_lib`. No `PyInit_*` symbol collision is possible now. Python import path becomes `sedonadb_zarr._lib`, matching `sedonadb._lib`. Also adds a `.gitignore` for `python/sedonadb-zarr` mirroring the sedonadb one so editable-install `.so` files don't get staged.
Post-FFI cleanup pass on the plugin handoff and adjacent code: - `mod context;` and `mod plugin;` are now private. They were `pub` only when `sedonadb-zarr` linked the host crate as an rlib; no external Rust consumer remains. - `python/sedonadb/src/plugin.rs`: capsule-name constants and module docstring trimmed; consts are no longer `pub`. - `python/sedonadb-zarr/src/lib.rs`: module/struct/fn docstrings collapsed to one line each. Behaviour unchanged. - `python/sedonadb-zarr/python/sedonadb_zarr/__init__.py`: docstrings converted from ReST (`:func:`, `:meth:`, NumPy `Parameters\n----`) to Markdown (mkdocs reads these as Markdown). - `python/sedonadb-zarr/README.md`: architecture line referenced a non-existent `register_udtf` PyO3 function; updated to the actual `ZarrTableFunction` / `datafusion-ffi` capsule story. - `python/sedonadb-zarr/Cargo.toml`: dropped the "why no sedonadb dep" essay; the absence speaks for itself. - `rust/sedona-datasource/src/provider.rs`: `SingleObjectExternalTable` docs and helper-fn docs collapsed. - `python/sedonadb/src/context.rs::register_udtf_capsule` docstring trimmed; mechanism docs live in `crate::plugin`. No public Rust API surface changes (the constants that lost `pub` had no Rust consumers — capsule names are part of the Python-visible ABI).
Ships the Zarr plugin with the `ExternalFormatSpec` surface only — `con.read_format(ZarrFormatSpec(), uri)`. The SQL UDTF form (`SELECT * FROM sd_read_zarr(...)`) and the `datafusion-ffi` plugin handoff machinery it requires are deferred so the plugin pattern can be designed without the time pressure of the zarr feature. Removed in this commit: - `rust/sedona-raster-zarr/src/udtf.rs` (ZarrReadFunction + ZarrChunkProvider + ZarrChunkExec) - `python/sedonadb-zarr/src/lib.rs` lost ZarrTableFunction and the `ffi_logical_codec_from_pycapsule` helper; the plugin's PyO3 module exports only `PyZarrChunkReader`. - `python/sedonadb/src/plugin.rs` (host-side session/codec capsule helpers). - `InternalContext::register_udtf_capsule` + `udtf_task_provider` field on the host side. - `python/sedonadb-zarr/python/sedonadb_zarr/__init__.py` lost `register(con)`; `con.read_format` is the only surface. - UDTF-only datafusion sub-crates (`datafusion-catalog`, `datafusion-execution`, `datafusion-physical-expr`, `datafusion-physical-plan`) and `datafusion-ffi` are dropped from `sedona-raster-zarr` and `sedonadb-zarr`. - Workspace `datafusion-proto` dep dropped (no remaining consumer). - UDTF-specific tests in `python/sedonadb-zarr/tests/test_zarr.py`; `test_dtype_mapping_roundtrips` rewired to use `read_format`. UDTF work preserved on `jw/sd-read-zarr-udtf` and tracked as DB-34.
paleolimbot
approved these changes
May 22, 2026
Member
paleolimbot
left a comment
paleolimbot
left a comment
There was a problem hiding this comment.
Thank you! Looking forward to this!
Member
|
Ah, I also see in CI the note about duplicate docs for |
Contributor
Author
|
sorry dewey, I havent finished dealing with your last round of comments yet! |
Two small fixes:
- `cargo doc --workspace` was failing with an output-path collision
("The lib `_lib` in sedonadb-zarr ... same name as the lib `_lib`
in sedonadb. Only one may be documented at once.") introduced when
the plugin renamed `_zarr_lib` → `_lib` per review. The plugin
cdylib has no Rust API surface (pure `#[pyclass]`), so `doc = false`
on its `[lib]` is the targeted fix — keeps the conventional `_lib`
name without re-introducing the collision dodge in the lib name.
- `rust/sedona`'s `[dev-dependencies]` carried a `sedona-raster` entry
added earlier in this PR that has no consumer in the crate (per
Dewey's nudge — `rg use sedona_raster\b` in `rust/sedona/` returns
nothing).
Per Dewey's top-level note: both `python/sedonadb` and `python/sedonadb-zarr` are pure Python extension modules with no Rust API surface for rustdoc, so neither should participate in `cargo doc --workspace`. Marking only one (the previous commit) left the other doc'ing into `target/doc/_lib/` solo — no collision today, but a foot-gun for any future plugin that also lands as `_lib`. Both crates already had `publish = false`; this just adds the matching `doc = false`.
The old comment said "so sedonadb's features don't leak" without naming what the actual leak is. After the `extension-module` cleanup the only feature in `MATURIN_PEP517_ARGS` is `s2geography`, which is sedonadb-only — the plugin crate has no `[features]` section, so a leaked `--features s2geography` would error. The trimmed comment says that directly.
`con.read_format(spec, paths)` is the Python plugin entry point — called out in README + docstrings since apache#858 opened, but the wrapper itself was never wired up. Tests had been ducking into the private `_impl.read_external_format`. Now there's a real method on `SedonaContext` matching `read_parquet` / `read_pyogrio` shape, and the plugin tests use it directly. Also rewords the loader-module docstring in `rust/sedona-raster-zarr/src/loader.rs` to describe the actual streaming-reader behaviour instead of the long-removed `group_to_indb_rasters` / `group_to_outdb_rasters` entry points.
Sweeping audit follow-ups now that the UDTF is split out: - Delete `rust/sedona-raster-zarr/src/format_spec.rs` and its re-export. The Rust `ZarrFormatSpec` had no consumer on either this branch or PR apache#866 (the UDTF uses `ZarrChunkReader` via `ZarrChunkProvider` directly). The R wrapper (DB-32) and URL-as- table support (DB-33) will revive the impl with a real caller; ~80 lines of trait delegation is cheap to reconstruct from git history. - Drop `sedona-datasource`, `async-trait`, and `arrow-array` deps from `sedona-raster-zarr`'s `[dependencies]`: format_spec.rs was their only consumer. (arrow-array stays as a transitive dep through `arrow-schema` / `RecordBatch` in loader.rs.) - Wait, `arrow-array` is still used by `loader.rs` for `RecordBatch` / `RecordBatchReader` — kept. - Drop `zarrs_object_store` workspace dep — no crate references it. - Rewrite `rust/sedona-raster-zarr/src/lib.rs` module docstring to describe the actual streaming reader, not the removed "async OutDb resolver registered separately" path. - Drop the unreachable `json.loads(arrays)` branch in `ZarrFormatSpec.open_reader`. It was for the SQL UDTF's JSON-string option form; `with_options` callers always pass real lists. Removes the only `import json` from the module too. - Trim plugin's `src/lib.rs` module docstring — claimed "via `ExternalFormatSpec` + Arrow C Stream" but the Rust shim doesn't touch `ExternalFormatSpec` (Python side handles that). Reworded to name `PyZarrChunkReader` and Arrow C Stream only.
`read_list_single_object` previously fell back to `False` when the attribute was absent — meant to support "older spec implementations that predate the directory-format path". No such impl exists: every `ExternalFormatSpec` shipped from this repo inherits the property from the base class (defaults to `False`). A duck-typed spec without the attribute now raises `AttributeError` directly, which is the clearer failure mode.
Ruff F821 on the new `SedonaContext.read_format()` signature — the forward-ref string annotation referenced a class that was never imported in this module. Pulling the import under `TYPE_CHECKING` keeps the annotation resolvable for static analysers without introducing a runtime cycle (`datasource` imports from `_lib`, which itself is built into the same package).
… polish docs)
Six small follow-ups from the final review pass:
- `ZarrFormatSpec.with_options` validates option keys against the
supported set (`{"arrays"}`) and raises on unknown keys — typos
like `"array"` vs `"arrays"` no longer silently read the whole
group. Docstring drops the stale "JSON-string" claim.
- `python/sedonadb/src/datasource.rs` doc comments updated to the
current `PyExternalFormat` name (three spots referenced the long-
renamed `PyExternalFormatSpec`).
- `rust/sedona-raster-zarr/src/lib.rs` shrinks the crate's public
surface: `dtype`, `geozarr`, `loader`, `source_uri` are now
private modules; only `ZarrChunkReader` is re-exported. Stops
internal types like `ChunkAnchor`, `GroupGeoMetadata`,
`group_uri_to_filesystem_path` from being part of the crate's
contract before they have any external consumer.
- `ChunkAnchor` + `parse_chunk_anchor` in `source_uri.rs` are now
`#[cfg(test)]`. They have no production consumer today (the async
byte resolver will resurrect them); the in-file tests still cover
the parse direction.
- `SingleObjectExternalTable::scan` documents why `_filters` is
discarded (no `supports_filters_pushdown` override → spatial
bbox pushdown doesn't fire for single-object formats).
- Mixed-store error message drops the `external_table:` prefix that
leaked the function name into user-facing text.
- `SedonaContext.read_format` docstring now points users at
`spec.with_options({...})` as the way to pass format-specific
options (no `options=` kwarg like `read_pyogrio`).
Contributor
Author
|
@paleolimbot now I am happy with it. review/approve/merge at your discretion. ETA: I think the test failure is a flakey test in GDAL |
james-willis
changed the title
paleolimbot
approved these changes
May 23, 2026
Comment on lines
+349
to
+355
| def read_format( | ||
| self, | ||
| spec: "ExternalFormatSpec", | ||
| table_paths: Union[str, Path, Iterable[str]], | ||
| check_extension: bool = False, | ||
| ) -> DataFrame: | ||
| """Read one or more paths using a Python-defined `ExternalFormatSpec`. |
Contributor
Author
There was a problem hiding this comment.
I almost forgot to create it! Would have made this feature challenging to use 😆
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
2 participants
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.
Summary
Adds Zarr support to SedonaDB as an opt-in Python plugin (
sedonadb-zarr). Each chunk position in the Zarr group becomes one row of a single-column raster table, similar to how each TIFF becomes one row.The main
sedonadbpackage has zero Zarr code or dependencies — applications that don'timport sedonadb_zarrpay no build or runtime cost. This is the plugin pattern future format crates (sedonadb-cog,sedonadb-icechunk, …) will follow.We built on the
zarrs0.23 crate rather than GDAL MultiDim because two GDAL gaps are independently disqualifying:What's in this PR
New
sedona-raster-zarrcrate. A streamingZarrChunkReader: RecordBatchReaderthat walks the group's chunk grid lazily, emitting one raster row per chunk position with one band per array. Per-batch memory isO(rows_per_batch × per-row-metadata), notO(total_chunks × chunk_size). Each row carries anoutdb_urichunk anchor;datastays empty until the async byte resolver lands.New
sedonadb-zarrPython plugin package. A PyO3 shim aroundsedona-raster-zarr:PyZarrChunkReaderimplements__arrow_c_stream__;ZarrFormatSpec(ExternalFormatSpec)consumes it from Python. Plugin Cargo.toml has no host-crate dependency — the plugin and host talk through the PythonExternalFormatSpeccontract and Arrow C Stream FFI.New host machinery in
sedonadb/sedona-datasource.SedonaContext.read_format(spec, paths)Python API for plugin-defined formats. Parallelsread_parquet/read_pyogrioin shape.ExternalFormatSpec.list_single_objectproperty (defaultFalse) on the Python base class and the matching Rust trait method.SingleObjectExternalTableRustTableProviderthat bypasses DataFusion'sListingTablefor directory-shape formats — a Zarr group is a directory, so the listing layer would return zero matching objects at a.zarrprefix.external_tabledispatcher that picks betweenSingleObjectExternalTableandListingTablebased on the spec'slist_single_object.Loader feature surface.
proj:wkt2/proj:projjson/proj:epsgfor CRS;spatial:transform/spatial:dimsfor affine + spatial axes.dimension_namesfield or the xarray-on-Zarr_ARRAY_DIMENSIONSattribute (supports xarray-produced v2 Zarrs).arrays = [...]option for an explicit subset; default auto-skips 1-D coord variables. Explicit 1-D names error at parse time.spatial:transformerrors; both-absent falls back to identity +log::warn!.bool,u8/i8,u16/i16,u32/i32,u64/i64,f32/f64— each exercised by a parameterized smoke test.file://URIs or bare paths).Out of scope (separate follow-ups)
sd_read_zarr). Split out to draft PR feat(sedonadb):sd_read_zarrUDTF via datafusion-ffi plugin handoff #866 to keep this PR focused on the loader + plugin pattern.SELECT * FROM 'foo.zarr'). Needs filesystem-catalog customisation in DataFusion's resolver.sedonadbzarrpackage. Mirrors the Python plugin pattern for R sedonadb users.RS_EnsureLoadedUDF. Materialises the chunk-anchor URIs into actual pixel bytes; until then every metadata-only kernel (count(*),RS_Envelope,RS_Width, …) works against the anchor-only rows.s3://,gs://,az://,https://).GeoTransformfromc/sedona-gdalintosedona-rasterso all raster backends share the same affine-transform type.