somatize_core/store/
mod.rs

1//! Data Store: abstraction for moving data between workers.
2//!
3//! Separates WHERE data lives from HOW it's processed.
4//! Workers use DataRef to reference data without materializing it.
5
6/// Maximum payload size for inline WebSocket transport.
7/// Payloads above this threshold are uploaded via HTTP bulk or DataStore.
8pub const INLINE_THRESHOLD_BYTES: usize = 10 * 1024 * 1024; // 10 MB
9
10// The S3 and Zarr backends live in `somatize-store`. They each own a
11// tokio runtime, and a contract crate must not hand one to everything
12// that depends on it.
13
14use crate::cache::CacheKey;
15use crate::error::{Result, SomaError};
16use crate::value::Value;
17use serde::{Deserialize, Serialize};
18
19/// Metadata about a stored value, queryable without loading data.
20#[derive(Debug, Clone, Serialize, Deserialize)]
21pub struct StoreMeta {
22    /// Total number of rows (`shape[0]` for tensors, 1 for scalar types).
23    pub total_rows: usize,
24    /// Remaining shape dimensions after the row axis (shape[1..] for tensors).
25    pub shape_tail: Vec<usize>,
26    /// Type tag: "tensor", "text", "json", "bytes", or "empty".
27    pub dtype: String,
28}
29
30impl StoreMeta {
31    /// Build metadata from an in-memory Value.
32    pub fn from_value(value: &Value) -> Self {
33        match value {
34            Value::Tensor { shape, .. } => Self {
35                total_rows: shape.first().copied().unwrap_or(0),
36                shape_tail: shape.get(1..).unwrap_or_default().to_vec(),
37                dtype: "tensor".into(),
38            },
39            Value::Text(_) => Self {
40                total_rows: 1,
41                shape_tail: vec![],
42                dtype: "text".into(),
43            },
44            Value::Json(_) => Self {
45                total_rows: 1,
46                shape_tail: vec![],
47                dtype: "json".into(),
48            },
49            Value::Bytes(b) | Value::Object(b) => Self {
50                total_rows: b.len(),
51                shape_tail: vec![],
52                dtype: "bytes".into(),
53            },
54            Value::Empty => Self {
55                total_rows: 0,
56                shape_tail: vec![],
57                dtype: "empty".into(),
58            },
59        }
60    }
61}
62
63/// Slice rows `[start..start+len)` from a tensor value.
64pub fn slice_tensor_rows(value: &Value, start: usize, len: usize) -> Result<Value> {
65    match value {
66        Value::Tensor { values, shape } => {
67            if shape.is_empty() {
68                return Err(SomaError::DataStore("cannot slice scalar tensor".into()));
69            }
70            let cols: usize = shape[1..].iter().product::<usize>().max(1);
71            let row_start = start * cols;
72            let row_end = (start + len) * cols;
73            if row_end > values.len() {
74                return Err(SomaError::DataStore(format!(
75                    "row range {start}..{} out of bounds (total rows: {})",
76                    start + len,
77                    shape[0]
78                )));
79            }
80            let mut new_shape = shape.clone();
81            new_shape[0] = len;
82            Ok(Value::tensor(
83                values[row_start..row_end].to_vec(),
84                new_shape,
85            ))
86        }
87        _ => Err(SomaError::DataStore(
88            "get_rows only works on Tensor values".into(),
89        )),
90    }
91}
92
93/// A reference to data that may live in different places.
94/// Workers exchange DataRefs instead of raw data.
95#[derive(Debug, Clone, Serialize, Deserialize)]
96#[serde(tag = "type")]
97#[non_exhaustive]
98pub enum DataRef {
99    /// Data in local filesystem
100    Local {
101        /// Absolute path to the file holding the serialized value.
102        path: String,
103    },
104    /// Data in S3-compatible object storage
105    S3 {
106        /// Bucket the object lives in.
107        bucket: String,
108        /// Object key within the bucket.
109        key: String,
110        /// AWS region, `None` for endpoints that don't need one.
111        region: Option<String>,
112    },
113    /// Data in Soma cache (content-addressable)
114    Cached {
115        /// Key the value is cached under.
116        cache_key: CacheKey,
117    },
118    /// Data available as a stream endpoint
119    Stream {
120        /// URL the stream can be read from.
121        endpoint: String,
122        /// Wire format of the streamed records.
123        format: StreamFormat,
124    },
125    /// Data materialized inline (small values only)
126    Inline {
127        /// The value itself, carried in the reference.
128        value: Value,
129    },
130    /// Data stored as a Zarr v3 array in object storage (chunked tensors).
131    Zarr {
132        /// Bucket the array lives in.
133        bucket: String,
134        /// Root path of the Zarr array (contains zarr.json + chunk objects).
135        array_path: String,
136        /// AWS region, `None` for endpoints that don't need one.
137        region: Option<String>,
138    },
139}
140
141/// Stream data format.
142#[derive(Debug, Clone, Serialize, Deserialize, Default)]
143#[serde(rename_all = "snake_case")]
144#[non_exhaustive]
145pub enum StreamFormat {
146    /// Newline-delimited JSON, one record per line (the default).
147    #[default]
148    JsonLines,
149    /// Comma-separated values.
150    Csv,
151    /// Apache Arrow IPC stream.
152    Arrow,
153    /// Length-prefixed protobuf messages.
154    Protobuf,
155}
156
157/// Storage configuration for an investigation/pipeline.
158#[derive(Debug, Clone, Serialize, Deserialize)]
159#[serde(tag = "type")]
160#[non_exhaustive]
161pub enum StorageConfig {
162    /// Local filesystem (NFS, mounted volume)
163    #[serde(rename = "local")]
164    Local {
165        /// Directory values are written under.
166        base_path: String,
167    },
168    /// S3-compatible object storage
169    #[serde(rename = "s3")]
170    S3 {
171        /// Bucket to store objects in.
172        bucket: String,
173        /// Key prefix all objects are written under.
174        prefix: String,
175        /// AWS region, `None` for endpoints that don't need one.
176        region: Option<String>,
177        /// Custom endpoint URL for non-AWS backends (MinIO, Ceph).
178        endpoint: Option<String>,
179    },
180    /// Zarr v3 chunked storage on S3-compatible backend.
181    #[serde(rename = "zarr")]
182    Zarr {
183        /// Bucket to store arrays in.
184        bucket: String,
185        /// Key prefix all arrays are written under.
186        prefix: String,
187        /// AWS region, `None` for endpoints that don't need one.
188        region: Option<String>,
189        /// Custom endpoint URL for non-AWS backends (MinIO, Ceph).
190        endpoint: Option<String>,
191        /// Rows per chunk (first dimension).
192        chunk_rows: usize,
193    },
194}
195
196impl Default for StorageConfig {
197    fn default() -> Self {
198        Self::Local {
199            base_path: "/tmp/soma-data".to_string(),
200        }
201    }
202}
203
204/// The DataStore trait: put/get/stream data across workers.
205///
206/// Unlike CacheStore (which stores Values by CacheKey),
207/// DataStore moves data between locations and supports streaming.
208pub trait DataStore: Send + Sync {
209    /// Store data and return a reference to it.
210    fn put(&self, key: &CacheKey, data: &Value) -> Result<DataRef>;
211
212    /// Retrieve data from a reference.
213    fn get(&self, data_ref: &DataRef) -> Result<Value>;
214
215    /// Check if data exists at a reference.
216    fn exists(&self, data_ref: &DataRef) -> Result<bool>;
217
218    /// Delete data at a reference.
219    fn remove(&self, data_ref: &DataRef) -> Result<()>;
220
221    /// Get the storage config.
222    fn config(&self) -> &StorageConfig;
223
224    /// Read a range of rows `[start..start+len)` from a tensor.
225    /// Returns a `Value::Tensor` with `shape[0] == len`.
226    /// Default impl downloads the full value and slices in memory.
227    fn get_rows(&self, data_ref: &DataRef, start: usize, len: usize) -> Result<Value> {
228        let value = self.get(data_ref)?;
229        slice_tensor_rows(&value, start, len)
230    }
231
232    /// Get metadata about a stored value without reading the data.
233    /// Default impl downloads the full value to extract metadata.
234    fn meta(&self, data_ref: &DataRef) -> Result<StoreMeta> {
235        let value = self.get(data_ref)?;
236        Ok(StoreMeta::from_value(&value))
237    }
238}
239
240/// Local filesystem data store.
241pub struct LocalDataStore {
242    config: StorageConfig,
243    base_path: std::path::PathBuf,
244}
245
246impl LocalDataStore {
247    /// Create a store rooted at `base_path`, creating the directory if
248    /// needed. Creation failure is deliberately ignored here — the
249    /// first `put` will surface it as a [`SomaError::DataStore`].
250    pub fn new(base_path: impl Into<std::path::PathBuf>) -> Self {
251        let base = base_path.into();
252        std::fs::create_dir_all(&base).ok();
253        Self {
254            config: StorageConfig::Local {
255                base_path: base.to_string_lossy().to_string(),
256            },
257            base_path: base,
258        }
259    }
260}
261
262impl DataStore for LocalDataStore {
263    fn put(&self, key: &CacheKey, data: &Value) -> Result<DataRef> {
264        let path = self.base_path.join(key.to_hex());
265        let bytes = serde_json::to_vec(data)
266            .map_err(|e| crate::error::SomaError::DataStore(e.to_string()))?;
267        std::fs::write(&path, &bytes)
268            .map_err(|e| crate::error::SomaError::DataStore(e.to_string()))?;
269        Ok(DataRef::Local {
270            path: path.to_string_lossy().to_string(),
271        })
272    }
273
274    fn get(&self, data_ref: &DataRef) -> Result<Value> {
275        match data_ref {
276            DataRef::Local { path } => {
277                let bytes = std::fs::read(path)
278                    .map_err(|e| crate::error::SomaError::DataStore(e.to_string()))?;
279                serde_json::from_slice(&bytes)
280                    .map_err(|e| crate::error::SomaError::DataStore(e.to_string()))
281            }
282            DataRef::Cached { cache_key } => {
283                let path = self.base_path.join(cache_key.to_hex());
284                let bytes = std::fs::read(&path)
285                    .map_err(|e| crate::error::SomaError::DataStore(e.to_string()))?;
286                serde_json::from_slice(&bytes)
287                    .map_err(|e| crate::error::SomaError::DataStore(e.to_string()))
288            }
289            DataRef::Inline { value } => Ok(value.clone()),
290            _ => Err(crate::error::SomaError::DataStore(
291                "Cannot get non-local DataRef from LocalDataStore".into(),
292            )),
293        }
294    }
295
296    fn exists(&self, data_ref: &DataRef) -> Result<bool> {
297        match data_ref {
298            DataRef::Local { path } => Ok(std::path::Path::new(path).exists()),
299            DataRef::Cached { cache_key } => Ok(self.base_path.join(cache_key.to_hex()).exists()),
300            DataRef::Inline { .. } => Ok(true),
301            _ => Ok(false),
302        }
303    }
304
305    fn remove(&self, data_ref: &DataRef) -> Result<()> {
306        if let DataRef::Local { path } = data_ref {
307            std::fs::remove_file(path).ok();
308        }
309        Ok(())
310    }
311
312    fn config(&self) -> &StorageConfig {
313        &self.config
314    }
315}
316
317#[cfg(test)]
318mod tests {
319    use super::*;
320
321    #[test]
322    fn local_data_store_roundtrip() {
323        let dir = std::env::temp_dir().join("soma-ds-test");
324        let store = LocalDataStore::new(&dir);
325
326        let key = CacheKey::hash_data(b"test_data");
327        let value = Value::tensor(vec![1.0, 2.0, 3.0], vec![3]);
328
329        let data_ref = store.put(&key, &value).unwrap();
330        assert!(store.exists(&data_ref).unwrap());
331
332        let retrieved = store.get(&data_ref).unwrap();
333        let (data, _) = retrieved.as_tensor().unwrap();
334        assert_eq!(data, &[1.0, 2.0, 3.0]);
335
336        store.remove(&data_ref).unwrap();
337        assert!(!store.exists(&data_ref).unwrap());
338
339        let _ = std::fs::remove_dir_all(&dir);
340    }
341
342    #[test]
343    fn inline_data_ref() {
344        let dir = std::env::temp_dir().join("soma-ds-test-inline");
345        let store = LocalDataStore::new(&dir);
346
347        let data_ref = DataRef::Inline {
348            value: Value::tensor(vec![42.0], vec![1]),
349        };
350
351        assert!(store.exists(&data_ref).unwrap());
352        let v = store.get(&data_ref).unwrap();
353        let (data, _) = v.as_tensor().unwrap();
354        assert_eq!(data, &[42.0]);
355
356        let _ = std::fs::remove_dir_all(&dir);
357    }
358
359    #[test]
360    fn storage_config_serde() {
361        let s3 = StorageConfig::S3 {
362            bucket: "my-lab".into(),
363            prefix: "experiments/".into(),
364            region: Some("eu-west-1".into()),
365            endpoint: None,
366        };
367        let json = serde_json::to_string(&s3).unwrap();
368        assert!(json.contains("my-lab"));
369
370        let local = StorageConfig::Local {
371            base_path: "/data".into(),
372        };
373        let json = serde_json::to_string(&local).unwrap();
374        assert!(json.contains("/data"));
375    }
376
377    #[test]
378    fn data_ref_serde() {
379        let refs = vec![
380            DataRef::Local {
381                path: "/tmp/x".into(),
382            },
383            DataRef::S3 {
384                bucket: "b".into(),
385                key: "k".into(),
386                region: None,
387            },
388            DataRef::Cached {
389                cache_key: CacheKey::hash_data(b"x"),
390            },
391            DataRef::Inline {
392                value: Value::Empty,
393            },
394            DataRef::Zarr {
395                bucket: "b".into(),
396                array_path: "data/abc".into(),
397                region: None,
398            },
399        ];
400        for r in &refs {
401            let json = serde_json::to_string(r).unwrap();
402            let _: DataRef = serde_json::from_str(&json).unwrap();
403        }
404    }
405
406    #[test]
407    fn slice_tensor_rows_basic() {
408        // 4 rows × 3 cols
409        let v = Value::tensor(
410            vec![
411                1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0,
412            ],
413            vec![4, 3],
414        );
415        // Rows 1..3 → [[4,5,6], [7,8,9]]
416        let sliced = slice_tensor_rows(&v, 1, 2).unwrap();
417        let (data, shape) = sliced.as_tensor().unwrap();
418        assert_eq!(shape, &[2, 3]);
419        assert_eq!(data, &[4.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
420    }
421
422    #[test]
423    fn slice_tensor_rows_single() {
424        let v = Value::tensor(vec![10.0, 20.0, 30.0], vec![3]);
425        let sliced = slice_tensor_rows(&v, 1, 1).unwrap();
426        let (data, shape) = sliced.as_tensor().unwrap();
427        assert_eq!(shape, &[1]);
428        assert_eq!(data, &[20.0]);
429    }
430
431    #[test]
432    fn slice_tensor_rows_out_of_bounds() {
433        let v = Value::tensor(vec![1.0, 2.0, 3.0], vec![3]);
434        assert!(slice_tensor_rows(&v, 2, 5).is_err());
435    }
436
437    #[test]
438    fn store_meta_from_tensor() {
439        let v = Value::tensor(vec![0.0; 12], vec![4, 3]);
440        let meta = StoreMeta::from_value(&v);
441        assert_eq!(meta.total_rows, 4);
442        assert_eq!(meta.shape_tail, vec![3]);
443        assert_eq!(meta.dtype, "tensor");
444    }
445
446    #[test]
447    fn store_meta_from_json() {
448        let v = Value::json(serde_json::json!({"a": 1}));
449        let meta = StoreMeta::from_value(&v);
450        assert_eq!(meta.dtype, "json");
451        assert_eq!(meta.total_rows, 1);
452    }
453
454    #[test]
455    fn default_get_rows_on_local_store() {
456        let dir = std::env::temp_dir().join("soma-ds-test-getrows");
457        let store = LocalDataStore::new(&dir);
458
459        let key = CacheKey::hash_data(b"rows_test");
460        let value = Value::tensor(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![3, 2]);
461        let data_ref = store.put(&key, &value).unwrap();
462
463        // Read rows 1..2 via default impl (full get + slice)
464        let sliced = store.get_rows(&data_ref, 1, 2).unwrap();
465        let (data, shape) = sliced.as_tensor().unwrap();
466        assert_eq!(shape, &[2, 2]);
467        assert_eq!(data, &[3.0, 4.0, 5.0, 6.0]);
468
469        let _ = std::fs::remove_dir_all(&dir);
470    }
471
472    #[test]
473    fn default_meta_on_local_store() {
474        let dir = std::env::temp_dir().join("soma-ds-test-meta");
475        let store = LocalDataStore::new(&dir);
476
477        let key = CacheKey::hash_data(b"meta_test");
478        let value = Value::tensor(vec![0.0; 20], vec![5, 4]);
479        let data_ref = store.put(&key, &value).unwrap();
480
481        let meta = store.meta(&data_ref).unwrap();
482        assert_eq!(meta.total_rows, 5);
483        assert_eq!(meta.shape_tail, vec![4]);
484        assert_eq!(meta.dtype, "tensor");
485
486        let _ = std::fs::remove_dir_all(&dir);
487    }
488
489    #[test]
490    fn zarr_storage_config_serde() {
491        let zarr = StorageConfig::Zarr {
492            bucket: "soma-research".into(),
493            prefix: "data/".into(),
494            region: None,
495            endpoint: Some("s3.eu-central-003.backblazeb2.com".into()),
496            chunk_rows: 1024,
497        };
498        let json = serde_json::to_string(&zarr).unwrap();
499        assert!(json.contains("soma-research"));
500        assert!(json.contains("1024"));
501        let _: StorageConfig = serde_json::from_str(&json).unwrap();
502    }
503}
somatize_core/store/mod.rs

somatize_core/store/
mod.rs
