Skip to main content

genegraph_storage/
metadata.rs

1#![allow(dead_code)]
2
3use log::{debug, info};
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6use std::path::PathBuf;
7
8use crate::StorageError;
9use crate::StorageResult;
10use crate::traits::backend::StorageBackend;
11use crate::traits::metadata::Metadata;
12
13/// Represent a single file spec in the persistence directory
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct FileInfo {
16    /// name of the file, can be equal to filetype if there is only one per type
17    pub filename: String,
18    /// see `Self::which_filetype(..)`: "rawinput" | "sub_centroids" | "lambdas" | "..."
19    pub filetype: String,
20    /// see `Self::which_format(..)`
21    pub storage_format: String,
22    pub rows: usize,
23    pub cols: usize,
24    pub nnz: Option<usize>,
25    pub size_bytes: Option<u64>,
26}
27
28impl FileInfo {
29    /// Create a file spec to add to the persistence directory
30    ///
31    /// Fails with `StorageError::UnsupportedFormat` for unrecognised
32    /// filetypes instead of panicking (issue #45).
33    pub fn new(
34        filename: String,
35        filetype: &str,
36        data_shape: (usize, usize),
37        nnz: Option<usize>,
38        size_bytes: Option<u64>,
39    ) -> StorageResult<Self> {
40        debug!(
41            "FileInfo::new: filename={}, filetype={}, shape={}x{}, nnz={:?}",
42            filename, filetype, data_shape.0, data_shape.1, nnz
43        );
44        Ok(Self {
45            filename,
46            filetype: filetype.into(),
47            storage_format: Self::which_format(filetype)?,
48            rows: data_shape.0,
49            cols: data_shape.1,
50            nnz,
51            size_bytes,
52        })
53    }
54
55    /// Assign the right format to the file type
56    pub fn which_format(filetype: &str) -> StorageResult<String> {
57        match filetype {
58            "dense" => Ok(String::from("lance fixed-row")),
59            "sparse" => Ok(String::from("lance row-major")),
60            "vector" => Ok(String::from("lance row-major")),
61            other => Err(StorageError::UnsupportedFormat(other.to_string())),
62        }
63    }
64
65    /// Assign the right filetype to the keyname of the file
66    pub fn which_filetype(filetype: &str) -> StorageResult<String> {
67        match filetype {
68            "rawinput" | "sub_centroids" | "dense" => Ok(String::from("dense")),
69            "adjacency" | "laplacian" | "signals" | "sparse" => Ok(String::from("sparse")),
70            "lambdas" | "item_norms" | "norms" | "vector" => Ok(String::from("vector")),
71            other => Err(StorageError::UnsupportedFiletype(other.to_string())),
72        }
73    }
74}
75
76/// Metadata for an ArrowSpace index persisted to Lance storage.
77///
78/// Tracks dataset dimensions, builder configuration, file locations, and pipeline context.
79#[derive(Debug, Clone, Serialize, Deserialize)]
80pub struct GeneMetadata {
81    pub name_id: String,
82    pub nrows: usize,
83    pub ncols: usize,
84    pub base: String,
85    pub files: HashMap<String, FileInfo>,
86    pub created_at: String,
87}
88
89impl GeneMetadata {
90    /// Read metadata file from JSON
91    pub async fn read(path: PathBuf) -> Result<Self, StorageError> {
92        info!("Reading metadata from {:?}", path);
93        let s = tokio::fs::read_to_string(path)
94            .await
95            .map_err(|e| StorageError::Io(e.to_string()))?;
96        let md: GeneMetadata = serde_json::from_str(&s).map_err(StorageError::Serde)?;
97        info!("Metadata read successfully");
98        Ok(md)
99    }
100}
101
102impl Metadata for GeneMetadata {
103    /// Empty metadata object
104    /// do not use in test, use seed_metadata_eigen instead
105    fn new(name_id: &str) -> Self {
106        info!("GeneMetadata::new: creating metadata for '{}'", name_id);
107        Self {
108            name_id: name_id.to_string(),
109            nrows: 0,
110            ncols: 0,
111            base: String::from(""),
112            files: HashMap::new(),
113            created_at: chrono::Utc::now().to_rfc3339(),
114        }
115    }
116
117    fn new_fileinfo(
118        &self,
119        key: &str,
120        filetype: &str,
121        data_shape: (usize, usize),
122        nnz: Option<usize>,
123        size_bytes: Option<u64>,
124    ) -> StorageResult<FileInfo> {
125        FileInfo::new(
126            format!("{}_{}.lance", self.name_id, key),
127            filetype,
128            (data_shape.0, data_shape.1),
129            nnz,
130            size_bytes,
131        )
132    }
133
134    /// Standard pipeline object
135    async fn seed_metadata<B: StorageBackend>(
136        name_id: &str,
137        nitems: usize,
138        nfeatures: usize,
139        storage: &B,
140    ) -> Result<GeneMetadata, StorageError> {
141        info!(
142            "GeneMetadata::seed_metadata: seeding metadata for '{}' with nitems={}, nfeatures={}",
143            name_id, nitems, nfeatures
144        );
145
146        let md = Self::new(name_id)
147            .with_base(storage.base_path())
148            .with_dimensions(nitems, nfeatures);
149
150        debug!("GeneMetadata::seed_metadata: saving metadata to storage");
151        storage.save_metadata(&md).await?;
152
153        info!(
154            "GeneMetadata::seed_metadata: metadata seeded successfully for '{}'",
155            name_id
156        );
157        Ok(md)
158    }
159
160    fn with_base(mut self, base_path: PathBuf) -> Self {
161        self.base = base_path.to_string_lossy().to_string();
162        self
163    }
164
165    fn with_dimensions(mut self, rows: usize, cols: usize) -> Self {
166        debug!(
167            "GeneMetadata::with_dimensions: setting dimensions to {}x{}",
168            rows, cols
169        );
170        self.nrows = rows;
171        self.ncols = cols;
172        self
173    }
174
175    fn add_file(mut self, key: &str, info: FileInfo) -> Self {
176        debug!(
177            "GeneMetadata::add_file: adding file '{}' ({})",
178            key, info.filename
179        );
180        self.files.insert(key.to_string(), info);
181        self
182    }
183}