genegraph_storage/
metadata.rs

1#![allow(dead_code)]
2
3use log::{debug, info};
4use serde::{Deserialize, Serialize};
5use std::collections::{BTreeMap, HashMap};
6use std::path::PathBuf;
7
8use crate::StorageError;
9use crate::StorageResult;
10use crate::traits::backend::StorageBackend;
11use crate::traits::metadata::Metadata;
12
13/// First-class collection kind (RFC #81-P1).
14///
15/// Every named collection in the metadata registry is either a vector space
16/// (dense vector collections and scalar artifacts), a graph (edge/adjacency
17/// artifacts), or a plain table (anything else).
18#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
19#[serde(rename_all = "kebab-case")]
20pub enum CollectionKind {
21    /// Dense vector collections (embeddings, maps, norms, lambdas, ...).
22    VectorSpace,
23    /// Graph data (edge lists, adjacency, laplacian).
24    Graph,
25    /// Anything else stored as a table.
26    Table,
27}
28
29impl CollectionKind {
30    /// Registry / descriptor string form (Generic Table API property value).
31    pub const fn as_str(&self) -> &'static str {
32        match self {
33            CollectionKind::VectorSpace => "vector-space",
34            CollectionKind::Graph => "graph",
35            CollectionKind::Table => "table",
36        }
37    }
38
39    /// Parses the string form ([`Self::as_str`]).
40    pub fn parse(s: &str) -> StorageResult<Self> {
41        match s {
42            "vector-space" => Ok(CollectionKind::VectorSpace),
43            "graph" => Ok(CollectionKind::Graph),
44            "table" => Ok(CollectionKind::Table),
45            other => Err(StorageError::Invalid(format!(
46                "unknown collection kind '{other}' (expected vector-space, graph or table)"
47            ))),
48        }
49    }
50
51    /// Derives the kind from a legacy artifact filetype; the compat shim
52    /// that turns the fixed artifact keys (`rawinput`, `adjacency`, ...)
53    /// into pre-seeded collections (RFC #81-P1).
54    pub fn for_filetype(filetype: &str) -> Option<Self> {
55        match filetype {
56            "dense" | "vectors" | "vector" => Some(CollectionKind::VectorSpace),
57            "sparse" | "graph" => Some(CollectionKind::Graph),
58            _ => None,
59        }
60    }
61}
62
63/// Represent a single file spec in the persistence directory
64#[derive(Debug, Clone, Serialize, Deserialize)]
65pub struct FileInfo {
66    /// name of the file, can be equal to filetype if there is only one per type
67    pub filename: String,
68    /// see `Self::which_filetype(..)`: "rawinput" | "sub_centroids" | "lambdas" | "..."
69    pub filetype: String,
70    /// see `Self::which_format(..)`
71    pub storage_format: String,
72    pub rows: usize,
73    pub cols: usize,
74    pub nnz: Option<usize>,
75    pub size_bytes: Option<u64>,
76    /// First-class collection kind (RFC #81-P1); `None` for entries written
77    /// before 0.28 (kind is then derived from `filetype`).
78    #[serde(default)]
79    pub kind: Option<CollectionKind>,
80    /// User properties attached to the collection (RFC #81-P1), e.g. the
81    /// `graph` linkage of a vector space (RFC #81-P4).
82    #[serde(default)]
83    pub properties: BTreeMap<String, String>,
84}
85
86impl FileInfo {
87    /// Create a file spec to add to the persistence directory
88    ///
89    /// Fails with `StorageError::UnsupportedFormat` for unrecognised
90    /// filetypes instead of panicking (issue #45).
91    pub fn new(
92        filename: String,
93        filetype: &str,
94        data_shape: (usize, usize),
95        nnz: Option<usize>,
96        size_bytes: Option<u64>,
97    ) -> StorageResult<Self> {
98        debug!(
99            "FileInfo::new: filename={}, filetype={}, shape={}x{}, nnz={:?}",
100            filename, filetype, data_shape.0, data_shape.1, nnz
101        );
102        Ok(Self {
103            filename,
104            filetype: filetype.into(),
105            storage_format: Self::which_format(filetype)?,
106            rows: data_shape.0,
107            cols: data_shape.1,
108            nnz,
109            size_bytes,
110            kind: CollectionKind::for_filetype(filetype),
111            properties: BTreeMap::new(),
112        })
113    }
114
115    /// Assign the right format to the file type
116    pub fn which_format(filetype: &str) -> StorageResult<String> {
117        match filetype {
118            "dense" | "vectors" => Ok(String::from("lance fixed-row")),
119            "sparse" => Ok(String::from("lance row-major")),
120            "vector" => Ok(String::from("lance row-major")),
121            "graph" => Ok(String::from("lance edge-list")),
122            other => Err(StorageError::UnsupportedFormat(other.to_string())),
123        }
124    }
125
126    /// Assign the right filetype to the keyname of the file
127    pub fn which_filetype(filetype: &str) -> StorageResult<String> {
128        match filetype {
129            "rawinput" | "sub_centroids" | "dense" => Ok(String::from("dense")),
130            "adjacency" | "laplacian" | "signals" | "sparse" => Ok(String::from("sparse")),
131            "lambdas" | "item_norms" | "norms" | "vector" => Ok(String::from("vector")),
132            "graph" | "vectors" => Ok(filetype.into()),
133            other => Err(StorageError::UnsupportedFiletype(other.to_string())),
134        }
135    }
136}
137
138/// Metadata for an ArrowSpace index persisted to Lance storage.
139///
140/// Tracks dataset dimensions, builder configuration, file locations, and pipeline context.
141#[derive(Debug, Clone, Serialize, Deserialize)]
142pub struct GeneMetadata {
143    pub name_id: String,
144    pub nrows: usize,
145    pub ncols: usize,
146    pub base: String,
147    pub files: HashMap<String, FileInfo>,
148    pub created_at: String,
149}
150
151impl GeneMetadata {
152    /// Read metadata file from JSON
153    pub async fn read(path: PathBuf) -> Result<Self, StorageError> {
154        info!("Reading metadata from {:?}", path);
155        let s = tokio::fs::read_to_string(path)
156            .await
157            .map_err(|e| StorageError::Io(e.to_string()))?;
158        let md: GeneMetadata = serde_json::from_str(&s).map_err(StorageError::Serde)?;
159        info!("Metadata read successfully");
160        Ok(md)
161    }
162}
163
164impl GeneMetadata {
165    /// Inherent constructor: mirrors [`Metadata::new`] so the common
166    /// construction chain works without importing the trait (0.25.0
167    /// ergonomics).
168    pub fn new(name_id: &str) -> Self {
169        <Self as Metadata>::new(name_id)
170    }
171
172    /// Inherent builder step: mirrors [`Metadata::with_base`].
173    pub fn with_base(self, base_path: PathBuf) -> Self {
174        <Self as Metadata>::with_base(self, base_path)
175    }
176
177    /// Inherent builder step: mirrors [`Metadata::with_dimensions`].
178    pub fn with_dimensions(self, rows: usize, cols: usize) -> Self {
179        <Self as Metadata>::with_dimensions(self, rows, cols)
180    }
181
182    /// Inherent builder step: mirrors [`Metadata::add_file`].
183    pub fn add_file(self, key: &str, info: FileInfo) -> Self {
184        <Self as Metadata>::add_file(self, key, info)
185    }
186
187    /// Inherent helper: mirrors [`Metadata::new_fileinfo`].
188    pub fn new_fileinfo(
189        &self,
190        key: &str,
191        filetype: &str,
192        data_shape: (usize, usize),
193        nnz: Option<usize>,
194        size_bytes: Option<u64>,
195    ) -> StorageResult<FileInfo> {
196        <Self as Metadata>::new_fileinfo(self, key, filetype, data_shape, nnz, size_bytes)
197    }
198}
199
200impl Metadata for GeneMetadata {
201    /// Empty metadata object
202    /// do not use in test, use seed_metadata_eigen instead
203    fn new(name_id: &str) -> Self {
204        info!("GeneMetadata::new: creating metadata for '{}'", name_id);
205        Self {
206            name_id: name_id.to_string(),
207            nrows: 0,
208            ncols: 0,
209            base: String::from(""),
210            files: HashMap::new(),
211            created_at: chrono::Utc::now().to_rfc3339(),
212        }
213    }
214
215    fn new_fileinfo(
216        &self,
217        key: &str,
218        filetype: &str,
219        data_shape: (usize, usize),
220        nnz: Option<usize>,
221        size_bytes: Option<u64>,
222    ) -> StorageResult<FileInfo> {
223        FileInfo::new(
224            crate::generations::artifact_file_name(&self.name_id, key),
225            filetype,
226            (data_shape.0, data_shape.1),
227            nnz,
228            size_bytes,
229        )
230    }
231
232    /// Standard pipeline object
233    async fn seed_metadata<B: StorageBackend>(
234        name_id: &str,
235        nitems: usize,
236        nfeatures: usize,
237        storage: &B,
238    ) -> Result<GeneMetadata, StorageError> {
239        info!(
240            "GeneMetadata::seed_metadata: seeding metadata for '{}' with nitems={}, nfeatures={}",
241            name_id, nitems, nfeatures
242        );
243
244        let md = Self::new(name_id)
245            .with_base(storage.base_path())
246            .with_dimensions(nitems, nfeatures);
247
248        debug!("GeneMetadata::seed_metadata: saving metadata to storage");
249        storage.save_metadata(&md).await?;
250
251        info!(
252            "GeneMetadata::seed_metadata: metadata seeded successfully for '{}'",
253            name_id
254        );
255        Ok(md)
256    }
257
258    fn with_base(mut self, base_path: PathBuf) -> Self {
259        self.base = base_path.to_string_lossy().to_string();
260        self
261    }
262
263    fn with_dimensions(mut self, rows: usize, cols: usize) -> Self {
264        debug!(
265            "GeneMetadata::with_dimensions: setting dimensions to {}x{}",
266            rows, cols
267        );
268        self.nrows = rows;
269        self.ncols = cols;
270        self
271    }
272
273    fn add_file(mut self, key: &str, info: FileInfo) -> Self {
274        debug!(
275            "GeneMetadata::add_file: adding file '{}' ({})",
276            key, info.filename
277        );
278        self.files.insert(key.to_string(), info);
279        self
280    }
281}
genegraph_storage/metadata.rs

genegraph_storage/
metadata.rs
