diff --git a/Cargo.lock b/Cargo.lock index 54bd5613..3361a149 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1525,19 +1525,23 @@ dependencies = [ "beacon-common", "beacon-datafusion-ext", "beacon-nd-array", + "beacon-object-storage", "datafusion", "futures", "hifitime", "indexmap 2.14.0", "ndarray 0.17.2", "object_store 0.13.2", + "serde", "serde_json", "tempfile", + "thiserror 2.0.18", "tokio", "tracing", "zarrs", "zarrs_object_store", "zarrs_storage", + "zip 2.4.2", ] [[package]] @@ -1681,6 +1685,7 @@ dependencies = [ "url", "utoipa", "uuid", + "zip 2.4.2", ] [[package]] diff --git a/beacon-api/src/axum/client/query.rs b/beacon-api/src/axum/client/query.rs index 401398d0..cf9952e3 100644 --- a/beacon-api/src/axum/client/query.rs +++ b/beacon-api/src/axum/client/query.rs @@ -152,6 +152,16 @@ async fn handle_query_output_file( QueryOutputFile::Odv(named_temp_file) => { file_stream_response(named_temp_file.path(), "application/zip", "zip", query_id).await } + QueryOutputFile::Zarr(named_temp_file) => { + // A zarr store is a directory tree, delivered as a zip archive. + file_stream_response( + named_temp_file.path(), + "application/zip", + "zarr.zip", + query_id, + ) + .await + } QueryOutputFile::NetCDF(named_temp_file) => { file_stream_response(named_temp_file.path(), "application/netcdf", "nc", query_id).await } diff --git a/beacon-core/Cargo.toml b/beacon-core/Cargo.toml index 3e3f16c5..6686be3f 100644 --- a/beacon-core/Cargo.toml +++ b/beacon-core/Cargo.toml @@ -58,6 +58,8 @@ test-util = [] [dev-dependencies] deltalake = { workspace = true } url = {workspace = true} +# Inspecting the zipped zarr stores produced by the zarr output formats. +zip = { workspace = true } # Self-reference enabling `test-util` for this crate's own integration tests # (tests/), so they can build an ephemeral in-memory-auth runtime. beacon-core = { path = ".", features = ["test-util"] } diff --git a/beacon-core/src/query/from.rs b/beacon-core/src/query/from.rs index 6efaa2ff..eaa19c39 100644 --- a/beacon-core/src/query/from.rs +++ b/beacon-core/src/query/from.rs @@ -237,7 +237,7 @@ mod tests { ), ( FromFormat::Zarr { paths: vec![] }, - Arc::new(beacon_arrow_zarr::datafusion::ZarrFormatFactory), + Arc::new(beacon_arrow_zarr::datafusion::ZarrFormatFactory::new()), ), ( FromFormat::Bbf { paths: vec![] }, diff --git a/beacon-core/src/query/output.rs b/beacon-core/src/query/output.rs index 04d39389..9ac86ee0 100644 --- a/beacon-core/src/query/output.rs +++ b/beacon-core/src/query/output.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use beacon_arrow_netcdf::datafusion::{options::NetcdfOptions, NetCDFFormatFactory, NetcdfConfig}; use beacon_arrow_odv::datafusion::OdvFileFormatFactory; +use beacon_arrow_zarr::datafusion::{ZarrFormatFactory, ZarrOptions}; use beacon_arrow_odv::writer::OdvOptions; use beacon_arrow_csv::datafusion::CsvFormatFactory; use beacon_arrow_geoparquet::datafusion::{GeoParquetFormatFactory, GeoParquetOptions}; @@ -92,6 +93,15 @@ pub enum OutputFormat { /// Columns to use as dimensions for the ND NetCDF output. dimension_columns: Vec, // Cannot be empty }, + /// Flat (record-oriented) Zarr v3 store, returned as a zip archive. + Zarr, + /// Multi-dimensional (gridded) Zarr v3 store, returned as a zip archive. + /// The named columns become the output dimensions; must not be empty. + #[serde(alias = "nd_zarr")] + NdZarr { + /// Columns to use as dimensions for the ND Zarr output. + dimension_columns: Vec, // Cannot be empty + }, /// GeoParquet format with optional longitude/latitude columns. GeoParquet { /// Name of the longitude column, if any. @@ -113,6 +123,8 @@ impl OutputFormat { OutputFormat::GeoParquet { .. } => QueryOutputFile::GeoParquet(temp_file), OutputFormat::NetCDF => QueryOutputFile::NetCDF(temp_file), OutputFormat::NdNetCDF { .. } => QueryOutputFile::NetCDF(temp_file), + OutputFormat::Zarr => QueryOutputFile::Zarr(temp_file), + OutputFormat::NdZarr { .. } => QueryOutputFile::Zarr(temp_file), OutputFormat::Odv(_) => QueryOutputFile::Odv(temp_file), } } @@ -158,6 +170,28 @@ impl OutputFormat { NetcdfConfig::default(), ))) } + OutputFormat::Zarr => { + // A zarr store is a directory; the query API can only return one + // file, so it is packed into a zip archive. + format_as_file_type(Arc::new(ZarrFormatFactory::new_for_write( + datasets_store.clone(), + ZarrOptions { + zip_output: true, + ..Default::default() + }, + ))) + } + OutputFormat::NdZarr { dimension_columns } => { + format_as_file_type(Arc::new(ZarrFormatFactory::new_for_write( + datasets_store.clone(), + ZarrOptions { + zip_output: true, + write_dimensions: Some(dimension_columns.clone()), + unique_value_columns: dimension_columns.clone(), + ..Default::default() + }, + ))) + } OutputFormat::Odv(options) => { format_as_file_type(Arc::new(OdvFileFormatFactory::new(Some(options.clone())))) } @@ -180,6 +214,8 @@ pub enum QueryOutputFile { NetCDF(NamedTempFile), /// ODV output file. Odv(NamedTempFile), + /// Zarr store, packed as a zip archive. + Zarr(NamedTempFile), /// GeoParquet output file. GeoParquet(NamedTempFile), } @@ -194,6 +230,7 @@ impl QueryOutputFile { QueryOutputFile::Parquet(file) => Ok(file.path().metadata()?.len()), QueryOutputFile::NetCDF(file) => Ok(file.path().metadata()?.len()), QueryOutputFile::Odv(file) => Ok(file.path().metadata()?.len()), + QueryOutputFile::Zarr(file) => Ok(file.path().metadata()?.len()), QueryOutputFile::GeoParquet(file) => Ok(file.path().metadata()?.len()), } } @@ -206,6 +243,7 @@ impl QueryOutputFile { QueryOutputFile::Parquet(file) => file.path(), QueryOutputFile::NetCDF(file) => file.path(), QueryOutputFile::Odv(file) => file.path(), + QueryOutputFile::Zarr(file) => file.path(), QueryOutputFile::GeoParquet(file) => file.path(), } } diff --git a/beacon-core/src/query_result.rs b/beacon-core/src/query_result.rs index da0e2c8a..107c02de 100644 --- a/beacon-core/src/query_result.rs +++ b/beacon-core/src/query_result.rs @@ -38,6 +38,8 @@ pub enum QueryOutputFile { Parquet(NamedTempFile), NetCDF(NamedTempFile), Odv(NamedTempFile), + /// Zarr store, packed as a zip archive. + Zarr(NamedTempFile), GeoParquet(NamedTempFile), } @@ -50,6 +52,7 @@ impl QueryOutputFile { QueryOutputFile::Parquet(file) => Ok(file.path().metadata()?.len()), QueryOutputFile::NetCDF(file) => Ok(file.path().metadata()?.len()), QueryOutputFile::Odv(file) => Ok(file.path().metadata()?.len()), + QueryOutputFile::Zarr(file) => Ok(file.path().metadata()?.len()), QueryOutputFile::GeoParquet(file) => Ok(file.path().metadata()?.len()), } } @@ -62,6 +65,7 @@ impl QueryOutputFile { QueryOutputFile::Parquet(file) => file.path(), QueryOutputFile::NetCDF(file) => file.path(), QueryOutputFile::Odv(file) => file.path(), + QueryOutputFile::Zarr(file) => file.path(), QueryOutputFile::GeoParquet(file) => file.path(), } } @@ -76,6 +80,7 @@ impl From for QueryOutputFile { crate::query::output::QueryOutputFile::Parquet(file) => Self::Parquet(file), crate::query::output::QueryOutputFile::NetCDF(file) => Self::NetCDF(file), crate::query::output::QueryOutputFile::Odv(file) => Self::Odv(file), + crate::query::output::QueryOutputFile::Zarr(file) => Self::Zarr(file), crate::query::output::QueryOutputFile::GeoParquet(file) => Self::GeoParquet(file), } } diff --git a/beacon-core/src/runtime.rs b/beacon-core/src/runtime.rs index dda25679..2d5c4928 100644 --- a/beacon-core/src/runtime.rs +++ b/beacon-core/src/runtime.rs @@ -1641,6 +1641,188 @@ mod client_query_tests { } } + /// Zarr output is a *directory* of chunk files, so the query API packs it + /// into a zip archive — the only shape a single-file HTTP response can + /// carry. This asserts the archive lands under the configured tmp dir and + /// holds the group metadata plus a `zarr.json` for the selected column. + #[tokio::test(flavor = "multi_thread")] + async fn query_with_zarr_output_returns_a_zipped_store() { + let config = std::sync::Arc::new(beacon_config::Config::load().unwrap()); + let tmp_dir = config.storage.tmp_dir.clone(); + let runtime = Runtime::new_with_in_memory_auth(config) + .await + .expect("runtime should start"); + let suffix = uuid::Uuid::new_v4().simple(); + let table = format!("zarrout_{suffix}"); + + run_sql(&runtime, &format!("CREATE TABLE {table} (a BIGINT)")).await; + run_sql(&runtime, &format!("INSERT INTO {table} VALUES (1), (2)")).await; + + let result = runtime + .run_query( + query(serde_json::json!({ + "from": table, + "select": ["a"], + "output": { "format": "zarr" }, + })), + beacon_auth::AuthIdentity::empty(), + ) + .await + .expect("zarr query with output should run"); + + match result.query_output { + QueryOutput::File(file) => { + let got = std::fs::canonicalize(file.path().parent().unwrap()) + .expect("canonicalize output parent"); + let want = std::fs::canonicalize(&tmp_dir).expect("canonicalize tmp dir"); + assert_eq!( + got, want, + "zarr output should be written under the configured tmp dir" + ); + + let entries = zip_entry_names(file.path()); + assert!( + entries.iter().any(|name| name == "zarr.json"), + "archive should hold the root group metadata, got {entries:?}" + ); + assert!( + entries.iter().any(|name| name == "a/zarr.json"), + "archive should hold the 'a' array metadata, got {entries:?}" + ); + // The scratch build directory must not survive alongside the zip. + let scratch = file.path().with_extension("zarr-build"); + assert!(!scratch.exists(), "scratch store dir should be cleaned up"); + } + QueryOutput::Stream(_) => panic!("expected a file output"), + } + } + + /// Gridded zarr output turns the named dimension columns into coordinate + /// arrays and every other column into an array indexed by them. Two `time` + /// values and two `depth` values must yield a 2×2 grid, not four rows. + #[tokio::test(flavor = "multi_thread")] + async fn query_with_nd_zarr_output_writes_gridded_arrays() { + let config = std::sync::Arc::new(beacon_config::Config::load().unwrap()); + let runtime = Runtime::new_with_in_memory_auth(config) + .await + .expect("runtime should start"); + let suffix = uuid::Uuid::new_v4().simple(); + let table = format!("ndzarrout_{suffix}"); + + run_sql( + &runtime, + &format!("CREATE TABLE {table} (time BIGINT, depth BIGINT, temp DOUBLE)"), + ) + .await; + run_sql( + &runtime, + &format!( + "INSERT INTO {table} VALUES \ + (1, 10, 1.5), (1, 20, 2.5), (2, 10, 3.5), (2, 20, 4.5)" + ), + ) + .await; + + let result = runtime + .run_query( + query(serde_json::json!({ + "from": table, + "select": ["time", "depth", "temp"], + "output": { + "format": { "ndzarr": { "dimension_columns": ["time", "depth"] } }, + }, + })), + beacon_auth::AuthIdentity::empty(), + ) + .await + .expect("nd zarr query with output should run"); + + match result.query_output { + QueryOutput::File(file) => { + let entries = zip_entry_names(file.path()); + for expected in ["zarr.json", "time/zarr.json", "depth/zarr.json", "temp/zarr.json"] + { + assert!( + entries.iter().any(|name| name == expected), + "archive should hold {expected}, got {entries:?}" + ); + } + + let temp_meta: serde_json::Value = + serde_json::from_slice(&read_zip_entry(file.path(), "temp/zarr.json")) + .expect("temp metadata should be valid JSON"); + assert_eq!( + temp_meta["shape"], + serde_json::json!([2, 2]), + "temp should be gridded as time x depth" + ); + assert_eq!( + temp_meta["dimension_names"], + serde_json::json!(["time", "depth"]) + ); + + let time_meta: serde_json::Value = + serde_json::from_slice(&read_zip_entry(file.path(), "time/zarr.json")) + .expect("time metadata should be valid JSON"); + assert_eq!( + time_meta["shape"], + serde_json::json!([2]), + "time should be a 1-D coordinate array of its distinct values" + ); + } + QueryOutput::Stream(_) => panic!("expected a file output"), + } + } + + /// `COPY TO ... STORED AS ZARR` writes a plain directory store rather than a + /// zip: unlike the HTTP API, a copy target is not limited to a single file. + #[tokio::test(flavor = "multi_thread")] + async fn copy_to_zarr_writes_a_directory_store() { + let config = std::sync::Arc::new(beacon_config::Config::load().unwrap()); + let tmp_dir = config.storage.tmp_dir.clone(); + let runtime = Runtime::new_with_in_memory_auth(config) + .await + .expect("runtime should start"); + let suffix = uuid::Uuid::new_v4().simple(); + let table = format!("zarrcopy_{suffix}"); + let target = format!("copy_{suffix}.zarr"); + + run_sql(&runtime, &format!("CREATE TABLE {table} (a BIGINT)")).await; + run_sql(&runtime, &format!("INSERT INTO {table} VALUES (1), (2)")).await; + run_sql( + &runtime, + &format!("COPY (SELECT a FROM {table}) TO '{target}' STORED AS ZARR"), + ) + .await; + + let store = tmp_dir.join(&target); + assert!(store.is_dir(), "expected a directory store at {store:?}"); + assert!(store.join("zarr.json").exists(), "root group metadata"); + assert!(store.join("a/zarr.json").exists(), "array metadata"); + } + + /// Entry names inside a zip archive, for asserting on zarr store layout. + fn zip_entry_names(path: &std::path::Path) -> Vec { + let file = std::fs::File::open(path).expect("open zip output"); + let mut archive = zip::ZipArchive::new(file).expect("output should be a zip archive"); + (0..archive.len()) + .map(|i| archive.by_index(i).expect("zip entry").name().to_string()) + .collect() + } + + /// Read one entry out of a zip archive. + fn read_zip_entry(path: &std::path::Path, name: &str) -> Vec { + use std::io::Read as _; + let file = std::fs::File::open(path).expect("open zip output"); + let mut archive = zip::ZipArchive::new(file).expect("output should be a zip archive"); + let mut entry = archive + .by_name(name) + .unwrap_or_else(|_| panic!("archive should contain {name}")); + let mut bytes = Vec::new(); + entry.read_to_end(&mut bytes).expect("read zip entry"); + bytes + } + /// `validate_query_plan` is the single permission gate: non-super-users may run /// read-only SELECTs but not DDL/DML (standard nodes) nor any beacon extension /// operation (super-user-only). diff --git a/beacon-data-lake/src/file_formats.rs b/beacon-data-lake/src/file_formats.rs index 1f63f1ae..9f5e443c 100644 --- a/beacon-data-lake/src/file_formats.rs +++ b/beacon-data-lake/src/file_formats.rs @@ -13,7 +13,7 @@ use beacon_arrow_ipc::datafusion::ArrowFormatFactory; use beacon_arrow_netcdf::datafusion::{NetCDFFormatFactory, options::NetcdfOptions}; use beacon_arrow_parquet::datafusion::ParquetFormatFactory; use beacon_arrow_tiff::datafusion::TiffFormatFactory; -use beacon_arrow_zarr::datafusion::ZarrFormatFactory; +use beacon_arrow_zarr::datafusion::{ZarrFormatFactory, ZarrOptions}; use beacon_datafusion_ext::format_ext::FileFormatFactoryExt; use beacon_object_storage::DatasetsStore; use datafusion::prelude::SessionContext; @@ -44,7 +44,12 @@ pub fn file_formats( config.atlas.clone(), )), Arc::new(TiffFormatFactory::new(Default::default())), - Arc::new(ZarrFormatFactory), + // Writable: `COPY TO ... STORED AS ZARR` builds a directory store under + // the datasets store's tmp root, the same way NetCDF does. + Arc::new(ZarrFormatFactory::new_for_write( + datasets_object_store.clone(), + ZarrOptions::default(), + )), Arc::new(BBFFormatFactory::new(config.bbf.clone())), Arc::new(GeoParquetFormatFactory::default()), ]; diff --git a/beacon-file-formats/beacon-arrow-zarr/Cargo.toml b/beacon-file-formats/beacon-arrow-zarr/Cargo.toml index 532266ba..320b1e45 100644 --- a/beacon-file-formats/beacon-arrow-zarr/Cargo.toml +++ b/beacon-file-formats/beacon-arrow-zarr/Cargo.toml @@ -14,15 +14,19 @@ datafusion = { workspace = true } ndarray = {workspace = true} anyhow = { workspace = true } async-trait = { workspace = true } +serde = { workspace = true } serde_json = { workspace = true } indexmap = { workspace = true } tracing = { workspace = true } hifitime = {workspace = true} futures = { workspace = true } +thiserror = { workspace = true } +zip = { workspace = true } beacon-nd-array = { path = "../beacon-nd-array" } beacon-datafusion-ext = { path = "../../beacon-datafusion-ext" } beacon-common = { path = "../../beacon-common" } +beacon-object-storage = { path = "../../beacon-object-storage" } [dev-dependencies] tempfile = { workspace = true } \ No newline at end of file diff --git a/beacon-file-formats/beacon-arrow-zarr/src/datafusion/mod.rs b/beacon-file-formats/beacon-arrow-zarr/src/datafusion/mod.rs index 6dbbfaef..c441c936 100644 --- a/beacon-file-formats/beacon-arrow-zarr/src/datafusion/mod.rs +++ b/beacon-file-formats/beacon-arrow-zarr/src/datafusion/mod.rs @@ -36,24 +36,63 @@ use crate::{ }, }; +pub mod options; +pub mod sink; pub mod source; +pub use options::ZarrOptions; +pub use sink::{ZarrNdSink, ZarrSink}; pub use source::ZarrSource; // ─── Factory ───────────────────────────────────────────────────────────────── #[derive(Default)] -pub struct ZarrFormatFactory; +pub struct ZarrFormatFactory { + /// Datasets store, used only for writing: a zarr store is built on the + /// local filesystem under its configured tmp root. `None` for read-only + /// registrations, where `create_writer_physical_plan` is never reached. + datasets_store: Option>, + /// Write-side defaults applied to formats this factory creates. + options: ZarrOptions, +} impl std::fmt::Debug for ZarrFormatFactory { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("ZarrFormatFactory").finish() + f.debug_struct("ZarrFormatFactory") + .field("options", &self.options) + .finish() } } impl ZarrFormatFactory { + /// A read-only factory. Writes through it fail with a clear error. pub fn new() -> Self { - Self + Self { + datasets_store: None, + options: ZarrOptions::default(), + } + } + + /// A factory that can also write, using `datasets_store`'s tmp root as the + /// destination for the store directory. + pub fn new_for_write( + datasets_store: Arc, + options: ZarrOptions, + ) -> Self { + Self { + datasets_store: Some(datasets_store), + options, + } + } + + fn format(&self, options: ZarrOptions) -> ZarrFormat { + ZarrFormat { + options, + output_dir: self + .datasets_store + .as_ref() + .map(|store| store.storage().tmp_dir.clone()), + } } } @@ -69,19 +108,27 @@ impl FileFormatFactory for ZarrFormatFactory { _state: &dyn Session, format_options: &std::collections::HashMap, ) -> datafusion::error::Result> { - // Per-table override from `CREATE EXTERNAL TABLE ... OPTIONS (...)`. - let read_dimensions = format_options.get("read_dimensions").map(|value| { + // Per-table overrides from `CREATE EXTERNAL TABLE ... OPTIONS (...)`. + let comma_separated = |value: &String| -> Vec { value .split(',') .map(|s| s.trim().to_string()) .filter(|s| !s.is_empty()) .collect() - }); - Ok(Arc::new(ZarrFormat::new(read_dimensions))) + }; + + let mut options = self.options.clone(); + if let Some(value) = format_options.get("read_dimensions") { + options.read_dimensions = Some(comma_separated(value)); + } + if let Some(value) = format_options.get("write_dimensions") { + options.write_dimensions = Some(comma_separated(value)); + } + Ok(Arc::new(self.format(options))) } fn default(&self) -> Arc { - Arc::new(ZarrFormat::default()) + Arc::new(self.format(self.options.clone())) } fn as_any(&self) -> &dyn Any { @@ -128,18 +175,30 @@ impl FileFormatFactoryExt for ZarrFormatFactory { #[derive(Debug, Clone, Default)] pub struct ZarrFormat { - /// Explicit dimensions requested via `read_zarr(paths, ['dims'])` or a - /// `CREATE EXTERNAL TABLE ... OPTIONS (read_dimensions '...')`. When set, - /// only variables whose dimensions are a subset of these are read; when - /// `None`, a broadcast-compatible default is auto-selected. - pub read_dimensions: Option>, + /// Read and write behaviour for this format. + pub options: ZarrOptions, + /// Local directory zarr stores are written into — the datasets store's tmp + /// root. `None` on read-only formats, which reject writes. + pub output_dir: Option, } impl ZarrFormat { - /// Build a format that reads only the variables belonging to + /// Build a read-only format that reads only the variables belonging to /// `read_dimensions` (or auto-selects a default when `None`). pub fn new(read_dimensions: Option>) -> Self { - Self { read_dimensions } + Self { + options: ZarrOptions { + read_dimensions, + ..Default::default() + }, + output_dir: None, + } + } + + /// Explicit dimensions requested via `read_zarr(paths, ['dims'])` or a + /// `CREATE EXTERNAL TABLE ... OPTIONS (read_dimensions '...')`. + fn read_dimensions(&self) -> Option> { + self.options.read_dimensions.clone() } } @@ -216,7 +275,7 @@ impl FileFormat for ZarrFormat { // default so the inferred schema matches what the scan returns. let any = match resolve_read_dimensions( &any, - self.read_dimensions.clone(), + self.read_dimensions(), Some("read_zarr"), ) { Some(dims) => any @@ -300,7 +359,7 @@ impl FileFormat for ZarrFormat { // source — rebuilding the source below would otherwise drop it. let projection = conf.file_source().projection().cloned(); let source = ZarrSource::new(table_schema) - .with_read_dimensions(self.read_dimensions.clone()) + .with_read_dimensions(self.read_dimensions()) .with_projection(projection); let conf = FileScanConfigBuilder::from(conf) .with_file_groups(file_groups) @@ -314,11 +373,50 @@ impl FileFormat for ZarrFormat { Ok(Arc::new(broadcast)) } + async fn create_writer_physical_plan( + &self, + input: Arc, + _state: &dyn Session, + conf: datafusion::datasource::physical_plan::FileSinkConfig, + order_requirements: Option, + ) -> datafusion::error::Result> { + // A zarr store is a directory tree of chunk files, which `zarrs` writes + // through a local filesystem store. Like NetCDF, we build it under the + // configured tmp root rather than streaming to an object store. + let output_dir = self.output_dir.clone().ok_or_else(|| { + datafusion::error::DataFusionError::Execution( + "Zarr writing requires a datasets store; this format was registered read-only" + .to_string(), + ) + })?; + + // `DataSinkExec` already requires a single input partition, so the + // gridded sink sees every row without extra plan nodes. + let sink: Arc = + match &self.options.write_dimensions { + Some(dimension_columns) if !dimension_columns.is_empty() => Arc::new( + sink::ZarrNdSink::new( + conf, + self.options.clone(), + dimension_columns.clone(), + output_dir, + )?, + ), + _ => Arc::new(sink::ZarrSink::new(conf, self.options.clone(), output_dir)), + }; + + Ok(Arc::new(datafusion::datasource::sink::DataSinkExec::new( + input, + sink, + order_requirements, + ))) + } + fn file_source( &self, table_schema: datafusion::datasource::table_schema::TableSchema, ) -> Arc { - Arc::new(ZarrSource::new(table_schema).with_read_dimensions(self.read_dimensions.clone())) + Arc::new(ZarrSource::new(table_schema).with_read_dimensions(self.read_dimensions())) } } diff --git a/beacon-file-formats/beacon-arrow-zarr/src/datafusion/options.rs b/beacon-file-formats/beacon-arrow-zarr/src/datafusion/options.rs new file mode 100644 index 00000000..a9a6b36c --- /dev/null +++ b/beacon-file-formats/beacon-arrow-zarr/src/datafusion/options.rs @@ -0,0 +1,90 @@ +//! Configuration for the Zarr file format in DataFusion queries. +//! +//! Controls both reading (which dimensions to unpack) and writing (flat vs. +//! gridded output, chunking, compression, and whether the resulting store is +//! packed into a single zip file). + +/// Compression applied to each chunk of a written zarr array. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum ZarrCompression { + /// No `bytes -> bytes` codec; chunks are stored raw. + None, + /// Zstandard, the zarr v3 default in most toolchains. + #[default] + Zstd, + /// Gzip, for readers without zstd support. + Gzip, +} + +/// Zarr specification version to emit. +/// +/// Only v3 can be written today: `zarrs` models v2 purely as a read-side format +/// that is converted to v3 in memory, and Beacon's own zarr reader discovers v3 +/// stores only. The option exists so adding a v2 writer later is non-breaking. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub enum ZarrVersion { + #[default] + #[serde(rename = "v3", alias = "3")] + V3, + #[serde(rename = "v2", alias = "2")] + V2, +} + +/// Configuration options for the Zarr file format. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct ZarrOptions { + /// Explicit dimensions requested when **reading**, via `read_zarr(paths, + /// ['dims'])` or `CREATE EXTERNAL TABLE ... OPTIONS (read_dimensions '...')`. + /// When `None`, a broadcast-compatible default set is auto-selected. + #[serde(default)] + pub read_dimensions: Option>, + /// Columns to use as zarr dimensions when **writing**. + /// + /// * `None` / empty → flat output via [`ZarrSink`](super::sink::ZarrSink): + /// every column becomes a 1-D array over a shared `obs` dimension. + /// * `Some(["lat", "lon", …])` → gridded output via + /// [`ZarrNdSink`](super::sink::ZarrNdSink): the listed columns become + /// coordinate arrays and the remaining columns become data arrays indexed + /// by them. + #[serde(default)] + pub write_dimensions: Option>, + /// Column names whose unique values are collected during execution to form + /// the dimension axes of a gridded write. + #[serde(default)] + pub unique_value_columns: Vec, + /// Pack the written store into a single zip file at the output path. + /// + /// The HTTP query API sets this because it can only hand back one file; a + /// `COPY TO` leaves it off and writes a plain directory store. + #[serde(default)] + pub zip_output: bool, + /// Chunk length along each dimension. A zarr store is a directory of chunk + /// files, so this trades file count against per-read granularity. + #[serde(default = "default_chunk_size")] + pub chunk_size: usize, + /// Chunk compression codec. + #[serde(default)] + pub compression: ZarrCompression, + /// Zarr specification version to emit. + #[serde(default)] + pub zarr_version: ZarrVersion, +} + +impl Default for ZarrOptions { + fn default() -> Self { + Self { + read_dimensions: None, + write_dimensions: None, + unique_value_columns: Vec::new(), + zip_output: false, + chunk_size: default_chunk_size(), + compression: ZarrCompression::default(), + zarr_version: ZarrVersion::default(), + } + } +} + +fn default_chunk_size() -> usize { + 64 * 1024 +} diff --git a/beacon-file-formats/beacon-arrow-zarr/src/datafusion/sink.rs b/beacon-file-formats/beacon-arrow-zarr/src/datafusion/sink.rs new file mode 100644 index 00000000..05b5ccc0 --- /dev/null +++ b/beacon-file-formats/beacon-arrow-zarr/src/datafusion/sink.rs @@ -0,0 +1,608 @@ +//! DataFusion sinks that materialize Arrow [`RecordBatch`]es into Zarr v3 stores. +//! +//! Two [`DataSink`] implementations mirror the NetCDF pair: +//! +//! * [`ZarrSink`] — writes each column as a 1-D array over a shared `obs` +//! dimension, producing a *flat* store for table-shaped data. Batches stream +//! straight through, so memory stays bounded. +//! +//! * [`ZarrNdSink`] — reshapes tabular rows into an N-dimensional grid: the +//! named dimension columns become coordinate arrays and the remaining columns +//! become arrays indexed by them (e.g. `lat × lon × time`). +//! +//! The write mode is chosen by [`ZarrOptions::write_dimensions`]; see +//! [`ZarrFormat::create_writer_physical_plan`](super::ZarrFormat). +//! +//! # Output layout +//! +//! A zarr store is a directory, but callers such as the HTTP query API can only +//! return one file. With [`ZarrOptions::zip_output`] set, the store is built in +//! a scratch directory and then packed into a zip archive at the output path. + +use std::any::Any; +use std::fmt::Formatter; +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, BooleanArray, RecordBatch, UInt32Array}; +use arrow::datatypes::SchemaRef; +use datafusion::{ + datasource::{physical_plan::FileSinkConfig, sink::DataSink}, + error::DataFusionError, + execution::{SendableRecordBatchStream, TaskContext}, + physical_plan::{DisplayAs, DisplayFormatType}, +}; +use futures::StreamExt; + +use crate::datafusion::options::ZarrOptions; +use crate::writer::{OBS_DIMENSION, ZarrStoreWriter, archive_store, is_supported_field}; + +fn exec_err(e: impl std::fmt::Display) -> DataFusionError { + DataFusionError::Execution(e.to_string()) +} + +/// Where a sink builds its store and where the finished artifact lands. +struct OutputPaths { + /// Directory the zarr store is written into. + store_dir: PathBuf, + /// Zip archive to pack `store_dir` into, when zipping is enabled. + zip_target: Option, +} + +impl OutputPaths { + fn resolve(output_dir: &Path, prefix: &str, zip_output: bool) -> Self { + let target = output_dir.join(prefix); + if zip_output { + // Build beside the target so the rename/zip stays on one filesystem. + let mut scratch = target.clone().into_os_string(); + scratch.push(".zarr-build"); + Self { + store_dir: PathBuf::from(scratch), + zip_target: Some(target), + } + } else { + Self { + store_dir: target, + zip_target: None, + } + } + } + + /// Pack the store into its zip target, or leave the directory in place. + fn finalize(self) -> datafusion::error::Result<()> { + let Some(zip_target) = self.zip_target else { + return Ok(()); + }; + // The caller pre-creates an empty temp file at the target path. + let file = std::fs::File::create(&zip_target).map_err(exec_err)?; + archive_store(&self.store_dir, std::io::BufWriter::new(file)).map_err(exec_err)?; + if let Err(e) = std::fs::remove_dir_all(&self.store_dir) { + tracing::warn!(error = %e, path = %self.store_dir.display(), "failed to clean up zarr scratch directory"); + } + Ok(()) + } + + /// Clear anything already at the store path — `COPY TO` may have created an + /// empty placeholder file where we need a directory. + fn prepare(&self) -> datafusion::error::Result<()> { + if self.store_dir.is_file() { + std::fs::remove_file(&self.store_dir).map_err(exec_err)?; + } else if self.store_dir.is_dir() { + std::fs::remove_dir_all(&self.store_dir).map_err(exec_err)?; + } + Ok(()) + } +} + +/// Rejects columns whose Arrow type has no zarr equivalent, naming all of them +/// at once rather than failing on the first. +fn check_supported(schema: &SchemaRef) -> datafusion::error::Result<()> { + let unsupported: Vec = schema + .fields() + .iter() + .filter(|f| !is_supported_field(f)) + .map(|f| format!("{} ({})", f.name(), f.data_type())) + .collect(); + if unsupported.is_empty() { + Ok(()) + } else { + Err(exec_err(format!( + "Zarr output does not support these columns: {}", + unsupported.join(", ") + ))) + } +} + +// ─── Flat sink ───────────────────────────────────────────────────────────── + +/// [`DataSink`] that writes batches to a flat zarr store. +/// +/// Every column becomes a 1-D array over the `obs` dimension. Batches are +/// appended as they arrive: the arrays grow with the stream, so nothing is +/// buffered beyond the current batch. +#[derive(Debug, Clone)] +pub struct ZarrSink { + sink_config: FileSinkConfig, + options: ZarrOptions, + /// Directory the store is written to — the configured tmp store root. + output_dir: PathBuf, +} + +impl ZarrSink { + pub fn new(sink_config: FileSinkConfig, options: ZarrOptions, output_dir: PathBuf) -> Self { + Self { + sink_config, + options, + output_dir, + } + } +} + +impl DisplayAs for ZarrSink { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "ZarrSink") + } +} + +#[async_trait::async_trait] +impl DataSink for ZarrSink { + fn as_any(&self) -> &dyn Any { + self + } + + fn schema(&self) -> &SchemaRef { + self.sink_config.output_schema() + } + + fn metrics(&self) -> Option { + None + } + + async fn write_all( + &self, + data: SendableRecordBatchStream, + _context: &Arc, + ) -> datafusion::error::Result { + let schema = self.sink_config.output_schema().clone(); + check_supported(&schema)?; + + let paths = OutputPaths::resolve( + &self.output_dir, + self.sink_config.table_paths[0].prefix().as_ref(), + self.options.zip_output, + ); + paths.prepare()?; + tracing::info!("Writing flat Zarr store to path: {:?}", paths.store_dir); + + let mut writer = ZarrStoreWriter::new(paths.store_dir.clone(), self.options.clone()) + .map_err(exec_err)?; + + let dimensions = vec![OBS_DIMENSION.to_string()]; + for field in schema.fields() { + // Shape 0 for now: `append` grows the `obs` extent as batches land. + writer + .create_array(field, vec![0], &dimensions) + .map_err(exec_err)?; + } + + let mut rows_written: u64 = 0; + let mut stream = std::pin::pin!(data); + while let Some(batch) = stream.next().await { + let batch = batch?; + for (index, field) in schema.fields().iter().enumerate() { + writer + .append(field, rows_written, batch.column(index).as_ref()) + .map_err(exec_err)?; + } + rows_written += batch.num_rows() as u64; + } + + writer.finish().map_err(exec_err)?; + paths.finalize()?; + + Ok(rows_written) + } +} + +// ─── N-dimensional sink ──────────────────────────────────────────────────── + +/// [`DataSink`] that reshapes tabular rows into a gridded zarr store. +/// +/// The columns named in `dimension_columns` become the axes, in the order +/// given; each is written as a 1-D coordinate array of its sorted distinct +/// values. Every other column becomes an N-D array indexed by those axes, with +/// cells that no row covers left as the array's fill value. +/// +/// Rows whose dimension values are null are dropped — they have no position in +/// the grid. +#[derive(Debug, Clone)] +pub struct ZarrNdSink { + sink_config: FileSinkConfig, + options: ZarrOptions, + dimension_columns: Vec, + /// Directory the store is written to — the configured tmp store root. + output_dir: PathBuf, +} + +impl ZarrNdSink { + /// Create a gridded sink. + /// + /// Fails if a named dimension column is missing from the output schema. + pub fn new( + sink_config: FileSinkConfig, + options: ZarrOptions, + dimension_columns: Vec, + output_dir: PathBuf, + ) -> datafusion::error::Result { + let schema = sink_config.output_schema(); + for column in &dimension_columns { + if schema.index_of(column).is_err() { + return Err(exec_err(format!( + "Zarr dimension column '{column}' is not in the query output" + ))); + } + } + Ok(Self { + sink_config, + options, + dimension_columns, + output_dir, + }) + } +} + +impl DisplayAs for ZarrNdSink { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "ZarrNdSink") + } +} + +#[async_trait::async_trait] +impl DataSink for ZarrNdSink { + fn as_any(&self) -> &dyn Any { + self + } + + fn schema(&self) -> &SchemaRef { + self.sink_config.output_schema() + } + + fn metrics(&self) -> Option { + None + } + + async fn write_all( + &self, + data: SendableRecordBatchStream, + _context: &Arc, + ) -> datafusion::error::Result { + let schema = self.sink_config.output_schema().clone(); + check_supported(&schema)?; + + let paths = OutputPaths::resolve( + &self.output_dir, + self.sink_config.table_paths[0].prefix().as_ref(), + self.options.zip_output, + ); + paths.prepare()?; + tracing::info!("Writing gridded Zarr store to path: {:?}", paths.store_dir); + + // Gridding needs every row at once: a cell's position depends on the + // full set of distinct dimension values. + let mut rows_written: u64 = 0; + let mut batches = Vec::new(); + let mut stream = std::pin::pin!(data); + while let Some(batch) = stream.next().await { + let batch = batch?; + rows_written += batch.num_rows() as u64; + batches.push(batch); + } + let batch = arrow::compute::concat_batches(&schema, batches.iter())?; + + let grid = Grid::build(&batch, &self.dimension_columns)?; + + let mut writer = ZarrStoreWriter::new(paths.store_dir.clone(), self.options.clone()) + .map_err(exec_err)?; + + // Coordinate arrays: 1-D, each indexed by its own dimension. + for axis in &grid.axes { + let field = schema.field_with_name(&axis.name)?; + writer + .create_array( + field, + vec![axis.values.len() as u64], + std::slice::from_ref(&axis.name), + ) + .map_err(exec_err)?; + writer + .write_full(field, axis.values.as_ref()) + .map_err(exec_err)?; + } + + // Data arrays: N-D, indexed by every axis. + let shape: Vec = grid.axes.iter().map(|a| a.values.len() as u64).collect(); + let dimension_names: Vec = grid.axes.iter().map(|a| a.name.clone()).collect(); + for field in schema.fields() { + if self.dimension_columns.contains(field.name()) { + continue; + } + let column = batch.column(schema.index_of(field.name())?); + let slab = arrow::compute::take(column.as_ref(), &grid.cell_to_row, None)?; + writer + .create_array(field, shape.clone(), &dimension_names) + .map_err(exec_err)?; + writer.write_full(field, slab.as_ref()).map_err(exec_err)?; + } + + writer.finish().map_err(exec_err)?; + paths.finalize()?; + + Ok(rows_written) + } +} + +// ─── Gridding ────────────────────────────────────────────────────────────── + +/// One dimension of the output grid. +struct Axis { + name: String, + /// The dimension's sorted distinct values — the coordinate array. + values: ArrayRef, + /// Position along this axis for each input row. + row_positions: Vec, +} + +/// The mapping from tabular rows onto grid cells. +struct Grid { + axes: Vec, + /// For each cell of the flattened grid, the row that fills it (null where no + /// row covers the cell). Feeds `arrow::compute::take`. + cell_to_row: UInt32Array, +} + +impl Grid { + fn build(batch: &RecordBatch, dimension_columns: &[String]) -> datafusion::error::Result { + let schema = batch.schema(); + + // A row only has a grid position if all of its dimension values are + // non-null; drop the rest. + let mut usable = vec![true; batch.num_rows()]; + for column in dimension_columns { + let array = batch.column(schema.index_of(column)?); + if array.null_count() > 0 { + for (row, keep) in usable.iter_mut().enumerate() { + *keep &= array.is_valid(row); + } + } + } + + let mut axes = Vec::with_capacity(dimension_columns.len()); + for column in dimension_columns { + let array = batch.column(schema.index_of(column)?); + axes.push(dense_rank_axis(column.clone(), array, &usable)?); + } + + let total_cells: usize = axes.iter().map(|a| a.values.len()).product(); + let mut cell_to_row: Vec> = vec![None; total_cells]; + for row in 0..batch.num_rows() { + if !usable[row] { + continue; + } + // Row-major (C-order) flattening, matching the array shape. + let mut cell = 0usize; + for axis in &axes { + cell = cell * axis.values.len() + axis.row_positions[row] as usize; + } + // Later rows win, mirroring the NetCDF sink's last-write-wins scatter. + cell_to_row[cell] = Some(row as u32); + } + + Ok(Self { + axes, + cell_to_row: UInt32Array::from(cell_to_row), + }) + } +} + +/// Build one grid axis from a dimension column. +/// +/// Positions come from a dense rank over the column's sorted distinct values, +/// which works for any sortable Arrow type without per-type dispatch: sort the +/// column, mark where adjacent values differ, and accumulate. Unusable rows are +/// dropped up front so they cannot introduce a coordinate of their own. +fn dense_rank_axis( + name: String, + array: &ArrayRef, + usable: &[bool], +) -> datafusion::error::Result { + let kept: UInt32Array = usable + .iter() + .enumerate() + .filter(|(_, keep)| **keep) + .map(|(row, _)| row as u32) + .collect::>() + .into(); + + let kept_values = arrow::compute::take(array.as_ref(), &kept, None)?; + let num_kept = kept_values.len(); + + let sorted_indices = arrow::compute::sort_to_indices(kept_values.as_ref(), None, None)?; + let sorted = arrow::compute::take(kept_values.as_ref(), &sorted_indices, None)?; + + // `distinct` is null-aware equality: true wherever a sorted neighbour + // differs, which is exactly where the dense rank advances. + let changed: BooleanArray = if num_kept > 1 { + arrow::compute::kernels::cmp::distinct( + &sorted.slice(1, num_kept - 1), + &sorted.slice(0, num_kept - 1), + )? + } else { + BooleanArray::from(Vec::::new()) + }; + + let mut row_positions = vec![0u32; array.len()]; + // First occurrence (in sorted order) of each distinct value, used to + // materialize the coordinate array. + let mut first_of_rank: Vec = Vec::new(); + + let mut rank = 0u32; + for position in 0..num_kept { + if position > 0 && changed.value(position - 1) { + rank += 1; + } + let kept_index = sorted_indices.value(position); + if first_of_rank.len() as u32 == rank { + first_of_rank.push(kept_index); + } + row_positions[kept.value(kept_index as usize) as usize] = rank; + } + + // Coordinates are indices into the kept subset, so read them back from it. + let values = arrow::compute::take( + kept_values.as_ref(), + &UInt32Array::from(first_of_rank), + None, + )?; + Ok(Axis { + name, + values, + row_positions, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{AsArray, Float64Array, Int64Array}; + use arrow::datatypes::{DataType, Field, Schema}; + + /// Rows keyed by (time, depth), deliberately unsorted and with a gap so the + /// scatter has to place values rather than copy them through. + fn ungridded_batch() -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("time", DataType::Int64, true), + Field::new("depth", DataType::Int64, true), + Field::new("temp", DataType::Float64, true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from(vec![2, 1, 1])), + Arc::new(Int64Array::from(vec![20, 10, 20])), + Arc::new(Float64Array::from(vec![4.0, 1.0, 2.0])), + ], + ) + .unwrap() + } + + #[test] + fn grid_axes_are_sorted_distinct_values() { + let batch = ungridded_batch(); + let grid = Grid::build(&batch, &["time".to_string(), "depth".to_string()]).unwrap(); + + assert_eq!(grid.axes.len(), 2); + assert_eq!( + grid.axes[0].values.as_ref().as_primitive::().values(), + &[1, 2], + "time axis should be the sorted distinct values" + ); + assert_eq!( + grid.axes[1].values.as_ref().as_primitive::().values(), + &[10, 20], + "depth axis should be the sorted distinct values" + ); + } + + /// The cell mapping must place each row at `time × depth` in row-major + /// order, and leave uncovered cells empty. + #[test] + fn grid_scatters_rows_into_cells() { + let batch = ungridded_batch(); + let grid = Grid::build(&batch, &["time".to_string(), "depth".to_string()]).unwrap(); + + let temp = batch.column(2); + let slab = arrow::compute::take(temp.as_ref(), &grid.cell_to_row, None).unwrap(); + let slab = slab.as_ref().as_primitive::(); + + assert_eq!(slab.len(), 4, "2 times x 2 depths"); + // (time=1, depth=10) -> 1.0, (time=1, depth=20) -> 2.0 + assert_eq!(slab.value(0), 1.0); + assert_eq!(slab.value(1), 2.0); + // (time=2, depth=10) has no row. + assert!(slab.is_null(2), "uncovered cell should stay empty"); + // (time=2, depth=20) -> 4.0 + assert_eq!(slab.value(3), 4.0); + } + + /// A row with a null dimension value has no grid position and must not + /// introduce a coordinate of its own. + #[test] + fn grid_drops_rows_with_null_dimensions() { + let schema = Arc::new(Schema::new(vec![ + Field::new("time", DataType::Int64, true), + Field::new("temp", DataType::Float64, true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from(vec![Some(1), None, Some(2)])), + Arc::new(Float64Array::from(vec![1.0, 9.0, 2.0])), + ], + ) + .unwrap(); + + let grid = Grid::build(&batch, &["time".to_string()]).unwrap(); + assert_eq!( + grid.axes[0].values.as_ref().as_primitive::().values(), + &[1, 2], + "the null time must not become a coordinate" + ); + + let slab = arrow::compute::take(batch.column(1).as_ref(), &grid.cell_to_row, None).unwrap(); + let slab = slab.as_ref().as_primitive::(); + assert_eq!(slab.len(), 2); + assert_eq!(slab.value(0), 1.0); + assert_eq!(slab.value(1), 2.0); + } + + /// Zipping puts the store under a scratch path so the archive can be + /// written at the requested output path. + #[test] + fn zip_output_builds_in_a_scratch_directory() { + let zipped = OutputPaths::resolve(Path::new("/tmp"), "out.tmp", true); + assert_eq!(zipped.zip_target.as_deref(), Some(Path::new("/tmp/out.tmp"))); + assert_ne!(zipped.store_dir, PathBuf::from("/tmp/out.tmp")); + + let plain = OutputPaths::resolve(Path::new("/tmp"), "out.zarr", false); + assert_eq!(plain.store_dir, PathBuf::from("/tmp/out.zarr")); + assert!(plain.zip_target.is_none()); + } + + #[test] + fn nd_sink_rejects_unknown_dimension_columns() { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, true)])); + let config = FileSinkConfig { + original_url: String::new(), + object_store_url: datafusion::execution::object_store::ObjectStoreUrl::local_filesystem(), + file_group: datafusion::datasource::physical_plan::FileGroup::default(), + table_paths: vec![ + datafusion::datasource::listing::ListingTableUrl::parse("file:///tmp/out").unwrap(), + ], + output_schema: schema, + table_partition_cols: vec![], + insert_op: datafusion::logical_expr::dml::InsertOp::Overwrite, + keep_partition_by_columns: false, + file_extension: "zarr".to_string(), + file_output_mode: datafusion::datasource::physical_plan::FileOutputMode::SingleFile, + }; + + let err = ZarrNdSink::new( + config, + ZarrOptions::default(), + vec!["missing".to_string()], + PathBuf::from("/tmp"), + ) + .unwrap_err(); + assert!(err.to_string().contains("missing"), "{err}"); + } +} diff --git a/beacon-file-formats/beacon-arrow-zarr/src/lib.rs b/beacon-file-formats/beacon-arrow-zarr/src/lib.rs index 1ae8300d..2aa524b0 100644 --- a/beacon-file-formats/beacon-arrow-zarr/src/lib.rs +++ b/beacon-file-formats/beacon-arrow-zarr/src/lib.rs @@ -16,3 +16,4 @@ pub mod data_types; pub mod datafusion; pub mod reader; pub mod util; +pub mod writer; diff --git a/beacon-file-formats/beacon-arrow-zarr/src/writer.rs b/beacon-file-formats/beacon-arrow-zarr/src/writer.rs new file mode 100644 index 00000000..cf6f7995 --- /dev/null +++ b/beacon-file-formats/beacon-arrow-zarr/src/writer.rs @@ -0,0 +1,616 @@ +//! Writing Arrow data out as a Zarr v3 store. +//! +//! A zarr store is a *directory* of metadata and chunk files, not a single +//! file. [`ZarrStoreWriter`] builds that directory on the local filesystem, and +//! [`archive_store`] optionally packs it into one zip file for callers (such as +//! the HTTP query API) that can only hand back a single artifact. +//! +//! Two shapes are supported, mirroring the NetCDF sinks: +//! +//! * **Flat** — every column becomes a 1-D array over a shared `obs` dimension. +//! Rows stream in and are appended, so memory stays bounded. +//! * **Gridded** — named dimension columns become coordinate arrays and the +//! remaining columns become N-D arrays indexed by them. See +//! [`crate::datafusion::sink::ZarrNdSink`]. +//! +//! Arrays carry CF attributes (`units`/`calendar` for temporal columns) so that +//! Beacon's own zarr reader — and xarray — decode what we wrote back to the +//! original Arrow types. See [`fill_value_for`] for how nulls are represented. + +use std::collections::HashMap; +use std::io::{Seek, Write}; +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +use arrow::array::{Array, AsArray}; +use arrow::datatypes::{DataType, Field, TimeUnit}; +use zarrs::array::{ArrayBuilder, DimensionName}; +use zarrs::filesystem::FilesystemStore; +use zarrs::group::GroupBuilder; + +use crate::datafusion::options::{ZarrCompression, ZarrOptions, ZarrVersion}; + +/// Dimension name used by flat (record-oriented) output. +pub const OBS_DIMENSION: &str = "obs"; + +/// Errors raised while writing a zarr store. +#[derive(Debug, thiserror::Error)] +pub enum ZarrWriteError { + #[error("unsupported Arrow data type {data_type} in column '{column}'")] + UnsupportedDataType { column: String, data_type: DataType }, + #[error("zarr v2 output is not supported yet; only v3 stores can be written")] + UnsupportedVersion, + #[error("zarr error: {0}")] + Zarr(String), + #[error(transparent)] + Arrow(#[from] arrow::error::ArrowError), + #[error(transparent)] + Io(#[from] std::io::Error), +} + +type Result = std::result::Result; + +fn zarr_err(e: impl std::fmt::Display) -> ZarrWriteError { + ZarrWriteError::Zarr(e.to_string()) +} + +// ─── Arrow → zarr type mapping ─────────────────────────────────────────────── + +/// The zarr v3 data type name a column is stored as, plus the CF attributes +/// needed to decode it back into the original Arrow type. +struct ZarrColumnType { + /// Zarr v3 data type name (e.g. `"float64"`). + name: &'static str, + /// Attributes written alongside the array (CF `units`, `calendar`, …). + attributes: Vec<(&'static str, serde_json::Value)>, +} + +/// Maps an Arrow field onto a zarr v3 data type. +/// +/// Temporal types are stored as their underlying integer with CF `units`, which +/// is what [`crate::compat::array_to_nd_array`]'s `CfTimeBackend` reads back. +fn zarr_type_for(field: &Field) -> Result { + let plain = |name| ZarrColumnType { + name, + attributes: Vec::new(), + }; + + Ok(match field.data_type() { + DataType::Boolean => plain("bool"), + DataType::Int8 => plain("int8"), + DataType::Int16 => plain("int16"), + DataType::Int32 => plain("int32"), + DataType::Int64 => plain("int64"), + DataType::UInt8 => plain("uint8"), + DataType::UInt16 => plain("uint16"), + DataType::UInt32 => plain("uint32"), + DataType::UInt64 => plain("uint64"), + DataType::Float32 => plain("float32"), + DataType::Float64 => plain("float64"), + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => plain("string"), + DataType::Timestamp(unit, _) => { + let unit = match unit { + TimeUnit::Second => "seconds", + TimeUnit::Millisecond => "milliseconds", + TimeUnit::Microsecond => "microseconds", + TimeUnit::Nanosecond => "nanoseconds", + }; + ZarrColumnType { + name: "int64", + attributes: vec![ + ( + "units", + format!("{unit} since 1970-01-01T00:00:00Z").into(), + ), + ("calendar", "gregorian".into()), + ("standard_name", "time".into()), + ], + } + } + DataType::Date32 => ZarrColumnType { + name: "int32", + attributes: vec![ + ("units", "days since 1970-01-01T00:00:00Z".into()), + ("calendar", "gregorian".into()), + ], + }, + DataType::Date64 => ZarrColumnType { + name: "int64", + attributes: vec![ + ("units", "milliseconds since 1970-01-01T00:00:00Z".into()), + ("calendar", "gregorian".into()), + ], + }, + other => { + return Err(ZarrWriteError::UnsupportedDataType { + column: field.name().clone(), + data_type: other.clone(), + }); + } + }) +} + +/// Whether a field can be written at all. +pub fn is_supported_field(field: &Field) -> bool { + zarr_type_for(field).is_ok() +} + +// ─── Store writer ──────────────────────────────────────────────────────────── + +/// Builds a Zarr v3 store on the local filesystem from Arrow data. +/// +/// Arrays are created lazily on first write and their metadata is only flushed +/// in [`finish`](Self::finish), which lets the flat path grow the `obs` +/// dimension as batches arrive without knowing the row count up front. +pub struct ZarrStoreWriter { + store: Arc, + root: PathBuf, + options: ZarrOptions, + arrays: HashMap>, + /// Insertion order of `arrays`, so metadata is written deterministically. + order: Vec, +} + +impl ZarrStoreWriter { + /// Create a store rooted at `root`, which must not already exist as a file. + pub fn new(root: PathBuf, options: ZarrOptions) -> Result { + if options.zarr_version != ZarrVersion::V3 { + return Err(ZarrWriteError::UnsupportedVersion); + } + std::fs::create_dir_all(&root)?; + let store = Arc::new(FilesystemStore::new(&root).map_err(zarr_err)?); + + GroupBuilder::new() + .build(store.clone(), "/") + .map_err(zarr_err)? + .store_metadata() + .map_err(zarr_err)?; + + Ok(Self { + store, + root, + options, + arrays: HashMap::new(), + order: Vec::new(), + }) + } + + /// The directory the store is being written into. + pub fn root(&self) -> &Path { + &self.root + } + + /// Declare an array of `shape` indexed by `dimensions`. + /// + /// The array is created with its full shape; the flat path passes a shape of + /// zero and grows it via [`append`](Self::append). + pub fn create_array( + &mut self, + field: &Field, + shape: Vec, + dimensions: &[String], + ) -> Result<()> { + let column_type = zarr_type_for(field)?; + let chunk_shape = self.chunk_shape(&shape); + + let mut builder = ArrayBuilder::new( + shape, + chunk_shape, + column_type.name, + fill_value_for(column_type.name), + ); + builder.dimension_names(Some( + dimensions + .iter() + .map(|d| DimensionName::Some(d.clone())) + .collect::>(), + )); + + for (key, value) in column_type.attributes { + builder.attributes_mut().insert(key.to_string(), value); + } + match self.options.compression { + ZarrCompression::None => {} + ZarrCompression::Zstd => { + builder.bytes_to_bytes_codecs(vec![Arc::new( + zarrs::array::codec::bytes_to_bytes::zstd::ZstdCodec::new(3, false), + )]); + } + ZarrCompression::Gzip => { + builder.bytes_to_bytes_codecs(vec![Arc::new( + zarrs::array::codec::bytes_to_bytes::gzip::GzipCodec::new(5).map_err(zarr_err)?, + )]); + } + } + + let array = builder + .build(self.store.clone(), &format!("/{}", field.name())) + .map_err(zarr_err)?; + + self.order.push(field.name().clone()); + self.arrays.insert(field.name().clone(), array); + Ok(()) + } + + /// Chunk shape for an array: the configured chunk size along the first + /// (record) dimension, whole extents along the rest. + /// + /// Gridded arrays are typically small along the trailing coordinate axes, so + /// chunking those too would produce a large number of tiny files. + fn chunk_shape(&self, shape: &[u64]) -> Vec { + let chunk = self.options.chunk_size.max(1) as u64; + shape + .iter() + .enumerate() + .map(|(i, &extent)| { + if i == 0 { + // A zero extent means the array grows as rows stream in, so + // the chunk size cannot be clamped to it yet. + if extent == 0 { + chunk + } else { + chunk.min(extent) + } + } else { + extent.max(1) + } + }) + .collect() + } + + /// Write `values` into `array_name` over `offset..offset + len` of the first + /// dimension, growing the array if needed. + /// + /// Used by the flat path, where the total row count is unknown until the + /// stream ends. + pub fn append(&mut self, field: &Field, offset: u64, values: &dyn Array) -> Result<()> { + let array = self + .arrays + .get_mut(field.name()) + .ok_or_else(|| zarr_err(format!("array '{}' was not created", field.name())))?; + + let len = values.len() as u64; + if len == 0 { + return Ok(()); + } + + let end = offset + len; + if array.shape()[0] < end { + let mut shape = array.shape().to_vec(); + shape[0] = end; + array.set_shape(shape).map_err(zarr_err)?; + } + + let subset = zarrs::array::ArraySubset::new_with_ranges(&[offset..end]); + store_subset(array, &subset, field, values) + } + + /// Write a fully-formed slab covering the whole extent of `field`'s array. + /// + /// Used by the gridded path, which knows the complete shape up front. + pub fn write_full(&mut self, field: &Field, values: &dyn Array) -> Result<()> { + let array = self + .arrays + .get_mut(field.name()) + .ok_or_else(|| zarr_err(format!("array '{}' was not created", field.name())))?; + + let subset = zarrs::array::ArraySubset::new_with_shape(array.shape().to_vec()); + if subset.num_elements() != values.len() as u64 { + return Err(zarr_err(format!( + "column '{}' has {} values but its zarr array holds {}", + field.name(), + values.len(), + subset.num_elements() + ))); + } + store_subset(array, &subset, field, values) + } + + /// Flush every array's metadata, completing the store. + pub fn finish(self) -> Result<()> { + for name in &self.order { + self.arrays[name].store_metadata().map_err(zarr_err)?; + } + Ok(()) + } +} + +/// Encodes an Arrow array into the zarr subset, substituting the array's fill +/// value for nulls. +fn store_subset( + array: &zarrs::array::Array, + subset: &zarrs::array::ArraySubset, + field: &Field, + values: &dyn Array, +) -> Result<()> { + /// Downcast a primitive Arrow column and store it, mapping nulls to `$fill`. + macro_rules! store_primitive { + ($arrow_type:ty, $rust_type:ty, $fill:expr) => {{ + let typed = values.as_primitive::<$arrow_type>(); + let elements: Vec<$rust_type> = (0..typed.len()) + .map(|i| { + if typed.is_null(i) { + $fill + } else { + typed.value(i) as $rust_type + } + }) + .collect(); + array + .store_array_subset(subset, elements) + .map_err(zarr_err) + }}; + } + + use arrow::datatypes::*; + + match field.data_type() { + DataType::Boolean => { + let typed = values.as_boolean(); + let elements: Vec = (0..typed.len()) + .map(|i| !typed.is_null(i) && typed.value(i)) + .collect(); + array + .store_array_subset(subset, elements) + .map_err(zarr_err) + } + DataType::Int8 => store_primitive!(Int8Type, i8, 0), + DataType::Int16 => store_primitive!(Int16Type, i16, 0), + DataType::Int32 => store_primitive!(Int32Type, i32, 0), + DataType::Int64 => store_primitive!(Int64Type, i64, 0), + DataType::UInt8 => store_primitive!(UInt8Type, u8, 0), + DataType::UInt16 => store_primitive!(UInt16Type, u16, 0), + DataType::UInt32 => store_primitive!(UInt32Type, u32, 0), + DataType::UInt64 => store_primitive!(UInt64Type, u64, 0), + DataType::Float32 => store_primitive!(Float32Type, f32, f32::NAN), + DataType::Float64 => store_primitive!(Float64Type, f64, f64::NAN), + DataType::Date32 => store_primitive!(Date32Type, i32, 0), + DataType::Date64 => store_primitive!(Date64Type, i64, 0), + DataType::Timestamp(unit, _) => match unit { + TimeUnit::Second => store_primitive!(TimestampSecondType, i64, 0), + TimeUnit::Millisecond => store_primitive!(TimestampMillisecondType, i64, 0), + TimeUnit::Microsecond => store_primitive!(TimestampMicrosecondType, i64, 0), + TimeUnit::Nanosecond => store_primitive!(TimestampNanosecondType, i64, 0), + }, + DataType::Utf8 => store_strings(array, subset, values.as_string::().iter()), + DataType::LargeUtf8 => store_strings(array, subset, values.as_string::().iter()), + DataType::Utf8View => store_strings(array, subset, values.as_string_view().iter()), + other => Err(ZarrWriteError::UnsupportedDataType { + column: field.name().clone(), + data_type: other.clone(), + }), + } +} + +fn store_strings<'a>( + array: &zarrs::array::Array, + subset: &zarrs::array::ArraySubset, + values: impl Iterator>, +) -> Result<()> { + let elements: Vec = values.map(|v| v.unwrap_or_default().to_string()).collect(); + array + .store_array_subset(subset, elements) + .map_err(zarr_err) +} + +/// The zarr fill value for a data type — what Arrow nulls are encoded as. +/// +/// Floats use NaN, the conventional missing-value marker that xarray and CF +/// tooling already treat as missing. Other types have no sentinel that cannot +/// collide with real data, so their nulls become zero / `false` / `""` and are +/// indistinguishable from those values on read. No `_FillValue` attribute is +/// written: it cannot express NaN as a JSON number, and advertising a colliding +/// sentinel for the other types would mask real values. +fn fill_value_for(zarr_type: &str) -> zarrs::metadata::FillValueMetadata { + use zarrs::metadata::FillValueMetadata; + match zarr_type { + "float32" | "float64" => FillValueMetadata::String("NaN".to_string()), + "bool" => FillValueMetadata::Bool(false), + "string" => FillValueMetadata::String(String::new()), + _ => FillValueMetadata::Number(serde_json::Number::from(0)), + } +} + +// ─── Zip packaging ─────────────────────────────────────────────────────────── + +/// Pack a zarr store directory into a zip archive. +/// +/// The entries are `Stored` (uncompressed) because chunk data is already +/// compressed by its codec, and stored entries let `zarr-python`'s `ZipStore` +/// and fsspec read chunks without inflating the whole archive. +pub fn archive_store(store_root: &Path, writer: W) -> Result<()> { + let mut zip = zip::ZipWriter::new(writer); + let options: zip::write::FileOptions<'_, ()> = + zip::write::FileOptions::default().compression_method(zip::CompressionMethod::Stored); + + let mut stack = vec![store_root.to_path_buf()]; + while let Some(dir) = stack.pop() { + for entry in std::fs::read_dir(&dir)? { + let path = entry?.path(); + // Zip entries are `/`-separated paths relative to the store root, so + // the archive unpacks straight into a usable store. + let name = path + .strip_prefix(store_root) + .map_err(|e| zarr_err(e))? + .components() + .map(|c| c.as_os_str().to_string_lossy()) + .collect::>() + .join("/"); + + if path.is_dir() { + stack.push(path); + } else { + zip.start_file(name, options).map_err(zarr_err)?; + let mut file = std::fs::File::open(&path)?; + std::io::copy(&mut file, &mut zip)?; + } + } + } + + zip.finish().map_err(zarr_err)?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Float64Array, Int64Array, RecordBatch, StringArray}; + use arrow::datatypes::Schema; + use datafusion::datasource::listing::{ + ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, + }; + use datafusion::prelude::SessionContext; + + use crate::datafusion::ZarrFormat; + + fn sample_batch() -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("time", DataType::Int64, false), + Field::new("temp", DataType::Float64, true), + Field::new("label", DataType::Utf8, true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from(vec![10, 20, 30])), + Arc::new(Float64Array::from(vec![Some(1.5), None, Some(3.5)])), + Arc::new(StringArray::from(vec![Some("a"), Some("b"), None])), + ], + ) + .unwrap() + } + + /// Write a flat store, then read it back through the crate's own reader. + /// This is the round-trip that matters: what we emit must be a store Beacon + /// (and any other zarr v3 reader) can open. + #[tokio::test] + async fn flat_store_round_trips_through_the_reader() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join("flat.zarr"); + let batch = sample_batch(); + let schema = batch.schema(); + + let mut writer = ZarrStoreWriter::new(root.clone(), ZarrOptions::default()).unwrap(); + let dimensions = vec![OBS_DIMENSION.to_string()]; + for field in schema.fields() { + writer.create_array(field, vec![0], &dimensions).unwrap(); + } + for (index, field) in schema.fields().iter().enumerate() { + writer.append(field, 0, batch.column(index).as_ref()).unwrap(); + } + writer.finish().unwrap(); + + assert!(root.join("zarr.json").exists(), "root group metadata"); + assert!(root.join("temp/zarr.json").exists(), "array metadata"); + + let ctx = SessionContext::new(); + let table_path = + ListingTableUrl::parse(format!("file://{}/", root.display())).unwrap(); + let listing_options = ListingOptions::new(Arc::new(ZarrFormat::default())) + .with_file_extension("zarr.json"); + let config = ListingTableConfig::new(table_path) + .with_listing_options(listing_options) + .infer_schema(&ctx.state()) + .await + .unwrap(); + ctx.register_table("flat", Arc::new(ListingTable::try_new(config).unwrap())) + .unwrap(); + + let batches = ctx + .sql("SELECT time, temp, label FROM flat ORDER BY time") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let read_back = arrow::compute::concat_batches(&batches[0].schema(), batches.iter()).unwrap(); + assert_eq!(read_back.num_rows(), 3); + assert_eq!( + read_back + .column(0) + .as_primitive::() + .values(), + &[10, 20, 30] + ); + // Arrow nulls in a float column are written as NaN — the conventional + // zarr/CF missing-value marker — so they read back as NaN, not null. + let temp = read_back + .column(1) + .as_primitive::(); + assert!(temp.value(1).is_nan(), "null should be encoded as NaN"); + assert_eq!(temp.value(0), 1.5); + } + + /// The flat path grows the `obs` dimension as batches arrive, so a store + /// written across several batches must hold every row in order — including + /// across a chunk boundary, where a partial chunk is re-read and extended. + #[test] + fn appending_multiple_batches_grows_the_obs_dimension() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join("multi.zarr"); + let field = Field::new("a", DataType::Int64, false); + + // A chunk size of 2 forces the second and third batches to land in + // chunks that are already partially written. + let options = ZarrOptions { + chunk_size: 2, + ..Default::default() + }; + let mut writer = ZarrStoreWriter::new(root.clone(), options).unwrap(); + writer + .create_array(&field, vec![0], &[OBS_DIMENSION.to_string()]) + .unwrap(); + + let mut offset = 0u64; + for values in [vec![1i64, 2, 3], vec![4, 5], vec![6]] { + let array = Int64Array::from(values); + writer.append(&field, offset, &array).unwrap(); + offset += array.len() as u64; + } + writer.finish().unwrap(); + + let store = Arc::new(FilesystemStore::new(&root).unwrap()); + let array = zarrs::array::Array::open(store, "/a").unwrap(); + assert_eq!(array.shape(), &[6]); + let elements = array + .retrieve_array_subset::>(&zarrs::array::ArraySubset::new_with_shape(vec![6])) + .unwrap(); + assert_eq!(elements, vec![1, 2, 3, 4, 5, 6]); + } + + /// Unsupported Arrow types are rejected up front rather than producing a + /// half-written store. + #[test] + fn unsupported_types_are_rejected() { + let field = Field::new( + "nested", + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))), + true, + ); + assert!(!is_supported_field(&field)); + assert!(is_supported_field(&Field::new("x", DataType::Float32, true))); + } + + /// The zip archive must contain store-relative, `/`-separated paths so it + /// unpacks straight into a usable store. + #[test] + fn archive_uses_store_relative_paths() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join("z.zarr"); + std::fs::create_dir_all(root.join("temp/c/0")).unwrap(); + std::fs::write(root.join("zarr.json"), b"{}").unwrap(); + std::fs::write(root.join("temp/zarr.json"), b"{}").unwrap(); + std::fs::write(root.join("temp/c/0/0"), b"chunk").unwrap(); + + let mut buffer = std::io::Cursor::new(Vec::new()); + archive_store(&root, &mut buffer).unwrap(); + + let mut archive = zip::ZipArchive::new(buffer).unwrap(); + let mut names: Vec = (0..archive.len()) + .map(|i| archive.by_index(i).unwrap().name().to_string()) + .collect(); + names.sort(); + assert_eq!(names, vec!["temp/c/0/0", "temp/zarr.json", "zarr.json"]); + } +} diff --git a/clients/beacon-datalake-cli/beacon_datalake_cli/formats.py b/clients/beacon-datalake-cli/beacon_datalake_cli/formats.py index c9a2de37..e31520ac 100644 --- a/clients/beacon-datalake-cli/beacon_datalake_cli/formats.py +++ b/clients/beacon-datalake-cli/beacon_datalake_cli/formats.py @@ -11,7 +11,7 @@ from typing import Any # Simple (string) formats keyed by canonical name. -SIMPLE_FORMATS = {"csv", "ipc", "parquet", "netcdf", "odv"} +SIMPLE_FORMATS = {"csv", "ipc", "parquet", "netcdf", "odv", "zarr"} # File-extension -> canonical format name. EXT_TO_FORMAT = { @@ -22,6 +22,8 @@ ".nc": "netcdf", ".geoparquet": "geoparquet", ".odv": "odv", + # A zarr store is a directory, so it comes back packed as a zip. + ".zarr": "zarr", } # Aliases accepted on the --format flag. @@ -42,7 +44,7 @@ def infer_format_name(path: str, override: str | None) -> str: return EXT_TO_FORMAT[ext] raise ExportFormatError( f"cannot infer output format from '{path}'; pass --format " - f"(csv, parquet, ipc/arrow, netcdf, nd_netcdf, geoparquet, odv)" + f"(csv, parquet, ipc/arrow, netcdf, nd_netcdf, zarr, nd_zarr, geoparquet, odv)" ) @@ -64,7 +66,11 @@ def build_output_format( if not dimension_columns: raise ExportFormatError("nd_netcdf requires --dimension-columns") return {"nd_netcdf": {"dimension_columns": dimension_columns}} + if name == "nd_zarr": + if not dimension_columns: + raise ExportFormatError("nd_zarr requires --dimension-columns") + return {"nd_zarr": {"dimension_columns": dimension_columns}} raise ExportFormatError( f"unknown output format '{name}'; expected one of: " - f"csv, parquet, ipc/arrow, netcdf, nd_netcdf, geoparquet, odv" + f"csv, parquet, ipc/arrow, netcdf, nd_netcdf, zarr, nd_zarr, geoparquet, odv" ) diff --git a/clients/beacon-ts/src/types.ts b/clients/beacon-ts/src/types.ts index 71934d88..4e61ff07 100644 --- a/clients/beacon-ts/src/types.ts +++ b/clients/beacon-ts/src/types.ts @@ -84,11 +84,24 @@ export type OutputFormat = | SimpleOutputFormat | { ndnetcdf: { dimension_columns: string[] } } | { nd_netcdf: { dimension_columns: string[] } } + | { ndzarr: { dimension_columns: string[] } } + | { nd_zarr: { dimension_columns: string[] } } | { geoparquet: GeoParquetOptions } | { odv: Record }; -/** Output formats expressed as a bare string. `arrow` is an alias for `ipc`. */ -export type SimpleOutputFormat = "csv" | "ipc" | "arrow" | "parquet" | "netcdf"; +/** + * Output formats expressed as a bare string. `arrow` is an alias for `ipc`. + * + * `zarr` returns a zip archive: a zarr store is a directory, and a single HTTP + * response can only carry one file. + */ +export type SimpleOutputFormat = + | "csv" + | "ipc" + | "arrow" + | "parquet" + | "netcdf" + | "zarr"; export interface GeoParquetOptions { longitude_column?: string | null; diff --git a/integration-tests/test_nd_formats.py b/integration-tests/test_nd_formats.py index e3b140fa..ef621bcf 100644 --- a/integration-tests/test_nd_formats.py +++ b/integration-tests/test_nd_formats.py @@ -127,3 +127,51 @@ def test_nd_netcdf_output_roundtrip(client, datasets_dir): # 50 unique time values -> 50 cells along the single dimension. assert n == 50 assert "temperature" in client.sql_rows(f"SELECT * FROM read_netcdf(['{rel}']) LIMIT 1")[0] + + +# --- Zarr output round-trips -------------------------------------------------- +def _unzip_to_datasets(datasets_dir, rel_path: str, data: bytes) -> str: + """Unpack a zipped zarr store into the datasets dir and return its path. + + Zarr output arrives as a zip archive because an HTTP response can only carry + one file; ``read_zarr`` wants the directory store, so unpack it first. + """ + import io + import zipfile + + dst = datasets_dir / rel_path + if dst.exists(): + shutil.rmtree(dst) + dst.mkdir(parents=True) + with zipfile.ZipFile(io.BytesIO(data)) as archive: + archive.extractall(dst) + return rel_path + + +def test_flat_zarr_output_roundtrip(client, datasets_dir): + """Flat (record-oriented) Zarr output is readable again by read_zarr.""" + src = "SELECT temperature, salinity, depth FROM read_parquet(['obs/*.parquet']) LIMIT 100" + data = client.query_bytes(src, "zarr") + # Zip local file header magic. + assert data[:2] == b"PK", data[:8] + + rel = _unzip_to_datasets(datasets_dir, "roundtrip/flat.zarr", data) + back = client.sql_rows(f"SELECT * FROM read_zarr(['{rel}'])") + assert len(back) - 1 == 100 + assert "temperature" in back[0] + + +def test_nd_zarr_output_roundtrip(client, datasets_dir): + """ND-Zarr output keyed on the unique `time` column round-trips.""" + src = ( + "SELECT time, temperature, salinity FROM read_parquet(['obs/*.parquet']) " + "ORDER BY time LIMIT 50" + ) + data = client.query_bytes(src, {"ndzarr": {"dimension_columns": ["time"]}}) + assert data[:2] == b"PK", data[:8] + + rel = _unzip_to_datasets(datasets_dir, "roundtrip/nd.zarr", data) + n = client.count(f"SELECT * FROM read_zarr(['{rel}'])") + # 50 unique time values -> 50 cells along the single dimension. + assert n == 50 + assert "temperature" in client.sql_rows(f"SELECT * FROM read_zarr(['{rel}']) LIMIT 1")[0]