N-dimensional execution spine: flow zarr/netcdf through beacon.nd-encoded batches#328
Merged
Conversation
Introduce beacon-datafusion-ext::nd, a self-contained library for N-dimensional Arrow arrays and dimension-aware physical operators. - NdArrowArray: a flat Arrow array plus named C-order Dimensions; zero-copy subset and broadcast-composed materialize. - BroadcastMap: name-aligned broadcast expressed as stride arithmetic. - NdRecordBatch: columns (each on a subset of the dimensions) over a shared target grid. - NdSourceExec (+ NdBatchProvider) streams nd batches via an execute_nd side channel; NdBroadcastExec materializes them into flat RecordBatches. Filtering and projection operators will layer on top of this spine later. Depends only on arrow + datafusion; tested against an in-memory provider (31 tests).
- NdSourceExec::execute now delegates to NdBroadcastExec instead of materializing itself, so flattening lives in one place; the source is a pure nd-batch producer via execute_nd. - Add ExecutionPlanMetricsSet to both nodes (metrics() override): NdSourceExec records output grid rows, an nd_batches counter, and elapsed time; NdBroadcastExec records flattened output rows and materialization time via BaselineMetrics threaded through materialize_nd_stream.
Remove functions left dead after dropping filter/projection: - array.rs: ArraySubset + NdArrowArray::subset/contiguous_range, try_new_scalar, and the values/dims/shape/num_elements accessors; materialize loses its always-None selection parameter. - broadcast.rs: run_structure/to_run_array/RunStructure (REE), source_index_of, the target_shape accessor, and gather_indices' selection parameter; num_target_elements is now private. - dimensions.rs: Dimension::name_arc and Dimensions::is_scalar. - mod.rs: drop ArraySubset/RunStructure from the re-exports. These will return alongside filtering/projection. 23 tests pass.
DatasetNdProvider implements beacon-datafusion-ext's NdBatchProvider: it chunks a regular Dataset in C-order and yields one NdRecordBatch per chunk with each variable left un-broadcast, so a format can feed a dataset into NdSourceExec (with NdBroadcastExec materializing above). - add beacon-datafusion-ext dep to beacon-nd-array (no cycle) - expose the batch.rs chunking helpers as pub(crate) - parity test: provider -> NdSourceExec -> NdBroadcastExec matches the v1 broadcast stream at several batch sizes Build with --workspace (beacon-datafusion-ext needs workspace feature unification for arrow-schema/serde).
Add any_dataset_as_broadcast_stream in beacon-nd-array: a regular dataset flows DatasetNdProvider -> NdSourceExec -> NdBroadcastExec (the broadcast node materializes the un-broadcast nd batches); ragged datasets stay on the v1 stream. Both file-format openers now call it in place of the v1 broadcast engine, keeping DataFusion's file machinery (partitioning, projection pushdown, schema adaptation, caching, file-stats pruning) intact. Predicate chunk-pruning is dropped on the regular path for now — the nd spine has no filter node yet and the scan reports filters as inexact, so DataFusion's FilterExec preserves correctness. All zarr/netcdf tests pass.
Introduce a self-describing Arrow encoding for nd arrays: a beacon.nd
extension struct Struct{values: List<T>, dim_sizes: List<UInt32>,
dim_names: List<Utf8>} where one row is one nd array. Because it is an
ordinary Arrow array it can ride through a DataSourceExec as a normal
RecordBatch.
- nd/encoding.rs: encode/decode NdArrowArray + NdRecordBatch, the
beacon.nd extension field, and logical<->encoded schema helpers.
- NdSourceExec is now a DECODER node: it wraps a child plan producing
encoded RecordBatches and decodes them to NdRecordBatch in execute_nd
(schema = decoded/logical). Replaces the NdBatchProvider leaf model;
NdBatchProvider/MemoryNdBatchProvider removed. Tests drive it from an
in-memory DataSourceExec of encoded batches.
- re-add NdArrowArray::dims()/values() (the encoder reads them).
- beacon-nd-array: rework the format helper to chunk -> NdRecordBatch ->
materialize (broadcast) directly; still parity-tested vs v1.
Foundation for wiring formats to emit encoded batches through
DataSourceExec -> NdSourceExec -> NdBroadcastExec as real plan nodes.
The zarr and netcdf file openers now emit beacon.nd-encoded RecordBatches (one struct row per nd array) instead of flat broadcast batches. Both create_physical_plan build NdBroadcastExec(NdSourceExec(DataSourceExec)) with the DataSourceExec carrying the encoded (struct) schema, so the nd data rides up the plan as an ordinary Arrow type, is decoded by NdSourceExec, and broadcast by NdBroadcastExec — the nd nodes are now real, EXPLAIN-visible plan nodes. - encoding.rs: encoded_schema (logical->encoded field wrapping), null-struct decode (missing columns -> null scalar), zero-column count passthrough, and encode_flat_batch_as_nd (ragged: wrap on a synthetic row dim so broadcast is identity). - beacon-nd-array: any_dataset_as_encoded_stream (regular -> encoded chunks; ragged -> v1 flat then row-wrapped + encoded). - openers adapt in the encoded (struct) domain via the schema adapter; null-filled missing columns decode to all-null after broadcast. - update netcdf opener-level tests to the encoded source schema. Whole workspace: 867 passed / 0 failed.
Validates the NdBroadcastExec → NdSourceExec → DataSourceExec plan shape and that rank-0 attributes broadcast to constant columns across the grid, for both netcdf and zarr.
Adds ScaleOffsetVariableDecoder, wrapping the raw numeric decoder to decode packed variables to f64 (raw*scale+offset) with the fill value decoded by the same arithmetic so packed fill cells null out. Mirrors the zarr reader, so netcdf and zarr now decode scaled variables consistently.
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
1 participant
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
Introduces a self-contained N-dimensional execution spine in
beacon-datafusion-ext::ndand wires the netcdf and zarr readers through it. N-dimensional data now travels up the DataFusion plan as a real Arrow type (beacon.nd-encoded structs) carried by a standardDataSourceExec, with two dedicated plan nodes on top that decode and flatten it into an ordinary table.The pipeline
Each nd variable (plus its named dimensions) is packed into one
beacon.ndstruct row —Struct{ values: List<T>, dim_sizes: List<UInt32>, dim_names: List<Utf8> }— so a batch of them rides throughDataSourceExeclike any otherRecordBatch, keeping partitioning, projection pushdown, caching, and file-stats machinery.NdSourceExecdecodes them into named-dimension arrays;NdBroadcastExecbroadcasts each onto the shared grid (a single Arrowtake, or a zero-copy clone on the identity) to produce the flat table the rest of the query sees. Both nodes carry DataFusion metrics.What's included
beacon-datafusion-ext/src/nd/— the spine:Dimensions,NdArrowArray(flat values + named dims),BroadcastMap(name-aligned/xarray-style broadcast as per-axis strides),NdRecordBatch, thebeacon.ndArrow encoding (encoding.rs), and theNdSourceExec/NdBroadcastExecplan nodes.beacon-nd-arrayas an nd batch provider — chunks a regular dataset C-order intobeacon.nd-encoded batches; ragged (CF contiguous) datasets fall back to run-length expansion wrapped on a syntheticrowdim (identity broadcast).create_physical_planreturnsNdBroadcastExec(NdSourceExec(DataSourceExec)).<var>.<attr>,.<global>) flow through as rank-0beacon.ndcolumns that broadcast to constant columns across the grid.scale_factor/add_offsetin the netcdf reader — newScaleOffsetVariableDecoderdecodes packed variables tof64(raw*scale+offset), with the fill value decoded by the same arithmetic so packed fill cells null out. This aligns netcdf with the zarr reader (both now decode scaled variables toFloat64).Tests
End-to-end coverage for both formats:
NdBroadcastExec → NdSourceExec → DataSourceExecnesting.Full workspace: 819 passed / 0 failed.
Deferred (follow-ups)
NdBroadcastExecas inexactFilterExec.valid_min/valid_maxrange masking (neither reader applies it today).