Add support for zarr

Cargo.toml

@@ -1,5 +1,5 @@
 [workspace]
-members = ["beacon-api", "beacon-arrow-netcdf", "beacon-arrow-odv", "beacon-common", "beacon-config", "beacon-core", "beacon-functions", "beacon-query", "beacon-planner", "beacon-data-lake", "beacon-formats"]
+members = ["beacon-api", "beacon-arrow-netcdf", "beacon-arrow-odv", "beacon-common", "beacon-config", "beacon-core", "beacon-functions", "beacon-query", "beacon-planner", "beacon-data-lake", "beacon-formats", "beacon-arrow-zarr"]
 
 [workspace.dependencies]
 tokio = { version = "1.47.1", features = ["full"] }
@@ -16,19 +16,19 @@ utoipa-axum = "0.2.0"
 utoipa-scalar = { version = "0.3.0", features = ["axum"] }
 utoipa-swagger-ui = { version = "9.0.0", features = ["axum"] }
 
-serde = { version = "=1.0.200", features = ["rc", "derive"] }
+serde = { version = "^1.0.2", features = ["rc", "derive"] }
 serde_json = "=1.0.120"
 anyhow = "1.0.95"
 thiserror = "2.0.12"
 
 tracing = "0.1.41"
 tracing-subscriber = { version = "0.3.19", features = ["env-filter"] }
-
 glob = "0.3.2"
 tempfile = "3.15.0"
 typetag = "0.2.19"
 indexmap = { version = "2.7.1", features = ["serde"]}
 chrono = { version = "0.4.41", features = ["serde"] }
+crossbeam = "0.8.4"
 
 datafusion = "49.0.0"
 object_store = { version = "0.12.3", features = ["aws"] }
@@ -39,4 +39,3 @@ parquet = { version = "^55.2.0", features = ["async"] }
 geoarrow = { version = "=0.4.0" }
 geoarrow-array = "=0.4.0"
 geoparquet = "0.4.0"
-pprof = { version = "0.15", features = ["flamegraph"] }

beacon-api/Cargo.toml

@@ -5,6 +5,7 @@ edition = "2021"
 
 [target.'cfg(not(windows))'.dependencies]
 tikv-jemallocator = "0.6.0"
+pprof = { version = "0.15", features = ["flamegraph"] }
 
 [dependencies]
 axum = { version = "0.8.1", features = ["tracing"] }
@@ -27,7 +28,6 @@ arrow-schema = { workspace = true}
 uuid = { version = "1.16.0" }
 tracing = { workspace = true}
 tracing-subscriber = { workspace = true}
-pprof = { workspace = true }
 
 # Local dependencies
 beacon-config = { path = "../beacon-config" }

beacon-arrow-zarr/Cargo.toml

@@ -0,0 +1,22 @@
+[package]
+name = "beacon-arrow-zarr"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+zarrs = {"version" = "0.22.2", features = ["async"]}
+zarrs_object_store = "0.5.0"
+zarrs_storage = "0.4.0"
+object_store = { workspace = true }
+tokio = { workspace = true }
+arrow = { workspace = true }
+nd-arrow-array = { git = "https://github.com/maris-development/nd-arrow-array.git", branch = "main", version = "1.2.2" }
+serde_json = { workspace = true }
+indexmap = { workspace = true }
+tracing = { workspace = true }
+hifitime = "4.0.2"
+regex = "1.11.1"
+futures = { workspace = true }
+pin-project = "1.1.10"
+parking_lot = { workspace = true }
+crossbeam = { workspace = true }

beacon-arrow-zarr/src/array_slice_pushdown.rs

@@ -0,0 +1,60 @@
+#[derive(Clone, Debug)]
+pub struct ArraySlicePushDown {
+    pub dimension: String,
+    pub start: Option<usize>,
+    pub end: Option<usize>,
+}
+
+impl ArraySlicePushDown {
+    pub fn new(dimension: String, start: Option<usize>, end: Option<usize>) -> Self {
+        Self {
+            dimension,
+            start,
+            end,
+        }
+    }
+
+    pub fn dimension(&self) -> &str {
+        &self.dimension
+    }
+
+    pub fn start(&self) -> Option<usize> {
+        self.start
+    }
+
+    pub fn end(&self) -> Option<usize> {
+        self.end
+    }
+
+    pub fn overlapping_range(
+        &self,
+        array_range: std::ops::Range<usize>,
+    ) -> Option<std::ops::Range<usize>> {
+        match (self.start, self.end) {
+            (Some(push_start), Some(push_end)) => {
+                let push_end = push_end + 1; // Make end inclusive
+                let start = array_range.start.max(push_start);
+                let end = array_range.end.min(push_end);
+                if start < end { Some(start..end) } else { None }
+            }
+            (Some(push_start), None) => {
+                let start = array_range.start.max(push_start);
+                if start < array_range.end {
+                    Some(start..array_range.end)
+                } else {
+                    None
+                }
+            }
+            (None, Some(push_end)) => {
+                let push_end = push_end + 1; // Make end inclusive
+                let end = array_range.end.min(push_end);
+                if array_range.start < end {
+                    Some(array_range.start..end)
+                } else {
+                    None
+                }
+            }
+            (None, None) => Some(array_range),
+        }
+    }
+}

beacon-arrow-zarr/src/attributes.rs

@@ -0,0 +1,78 @@
+use std::sync::Arc;
+
+use nd_arrow_array::NdArrowArray;
+
+pub enum AttributeValue {
+    String(String),
+    Float64(f64),
+    Bool(bool),
+}
+
+impl AttributeValue {
+    pub fn from_json_value(value: &serde_json::Value) -> Option<Self> {
+        match value {
+            serde_json::Value::String(s) => Some(AttributeValue::String(s.clone())),
+            serde_json::Value::Number(n) => n.as_f64().map(AttributeValue::Float64),
+            serde_json::Value::Bool(b) => Some(AttributeValue::Bool(*b)),
+            _ => None,
+        }
+    }
+
+    pub fn as_str(&self) -> Option<&str> {
+        match self {
+            AttributeValue::String(s) => Some(s.as_str()),
+            _ => None,
+        }
+    }
+
+    pub fn as_f64(&self) -> Option<f64> {
+        match self {
+            AttributeValue::Float64(f) => Some(*f),
+            _ => None,
+        }
+    }
+
+    pub fn as_bool(&self) -> Option<bool> {
+        match self {
+            AttributeValue::Bool(b) => Some(*b),
+            _ => None,
+        }
+    }
+
+    pub fn arrow_data_type(&self) -> arrow::datatypes::DataType {
+        match self {
+            AttributeValue::String(_) => arrow::datatypes::DataType::Utf8,
+            AttributeValue::Float64(_) => arrow::datatypes::DataType::Float64,
+            AttributeValue::Bool(_) => arrow::datatypes::DataType::Boolean,
+        }
+    }
+
+    pub fn as_nd_arrow_array(&self) -> NdArrowArray {
+        match self {
+            AttributeValue::String(s) => {
+                let array = arrow::array::StringArray::from(vec![s.as_str()]);
+                NdArrowArray::new(
+                    Arc::new(array) as arrow::array::ArrayRef,
+                    nd_arrow_array::dimensions::Dimensions::Scalar,
+                )
+                .unwrap()
+            }
+            AttributeValue::Float64(f) => {
+                let array = arrow::array::Float64Array::from(vec![*f]);
+                NdArrowArray::new(
+                    Arc::new(array) as arrow::array::ArrayRef,
+                    nd_arrow_array::dimensions::Dimensions::Scalar,
+                )
+                .unwrap()
+            }
+            AttributeValue::Bool(b) => {
+                let array = arrow::array::BooleanArray::from(vec![*b]);
+                NdArrowArray::new(
+                    Arc::new(array) as arrow::array::ArrayRef,
+                    nd_arrow_array::dimensions::Dimensions::Scalar,
+                )
+                .unwrap()
+            }
+        }
+    }
+}

beacon-arrow-zarr/src/data_types.rs

@@ -0,0 +1,20 @@
+pub fn try_zarrs_dtype_to_arrow(
+    data_type: zarrs::array::DataType,
+) -> Result<arrow::datatypes::DataType, String> {
+    match data_type {
+        zarrs::array::DataType::Bool => Ok(arrow::datatypes::DataType::Boolean),
+        zarrs::array::DataType::Int8 => Ok(arrow::datatypes::DataType::Int8),
+        zarrs::array::DataType::Int16 => Ok(arrow::datatypes::DataType::Int16),
+        zarrs::array::DataType::Int32 => Ok(arrow::datatypes::DataType::Int32),
+        zarrs::array::DataType::Int64 => Ok(arrow::datatypes::DataType::Int64),
+        zarrs::array::DataType::UInt8 => Ok(arrow::datatypes::DataType::UInt8),
+        zarrs::array::DataType::UInt16 => Ok(arrow::datatypes::DataType::UInt16),
+        zarrs::array::DataType::UInt32 => Ok(arrow::datatypes::DataType::UInt32),
+        zarrs::array::DataType::UInt64 => Ok(arrow::datatypes::DataType::UInt64),
+        zarrs::array::DataType::Float32 => Ok(arrow::datatypes::DataType::Float32),
+        zarrs::array::DataType::Float64 => Ok(arrow::datatypes::DataType::Float64),
+        zarrs::array::DataType::String => Ok(arrow::datatypes::DataType::Utf8),
+        zarrs::array::DataType::Bytes => Ok(arrow::datatypes::DataType::Binary),
+        _ => Err(format!("Unsupported Zarrs data type: {:?}", data_type)),
+    }
+}