Zarr output: write gridded and flat Zarr v3 stores

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
         }

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"] }

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![] },

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<String>, // 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<String>, // 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(),
         }
     }

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<crate::query::output::QueryOutputFile> 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),
         }
     }

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<String> {
+        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<u8> {
+        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).

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()),
     ];