genegraph_storage/traits/
lance.rs

1use std::path::Path;
2use std::sync::Arc;
3
4use arrow::array::{ArrayRef, PrimitiveArray};
5use arrow::datatypes::{ArrowPrimitiveType, Field, Schema};
6use arrow::record_batch::RecordBatch;
7use log::{debug, info};
8
9use crate::metadata::FileInfo;
10use crate::traits::backend::StorageBackend;
11use crate::traits::metadata::Metadata;
12use crate::{StorageError, StorageResult};
13
14#[cfg(feature = "official-lance")]
15use arrow::record_batch::RecordBatchIterator;
16#[cfg(feature = "official-lance")]
17use futures::StreamExt;
18#[cfg(feature = "official-lance")]
19use lance::Dataset;
20#[cfg(feature = "official-lance")]
21use lance::dataset::{WriteMode, WriteParams};
22
23/// Resolves a `file://` URI (as produced by `path_to_uri`) to a local path.
24#[cfg(not(feature = "official-lance"))]
25fn uri_to_path(uri: &str) -> StorageResult<std::path::PathBuf> {
26    let url = url::Url::parse(uri)
27        .map_err(|e| StorageError::Invalid(format!("bad dataset URI `{uri}`: {e}")))?;
28    url.to_file_path().map_err(|_| {
29        StorageError::Invalid(format!("dataset URI is not a local file path: `{uri}`"))
30    })
31}
32
33pub trait LanceStorage {
34    /// Async helper: write a RecordBatch to a Lance dataset.
35    ///
36    /// With default features this runs the in-house v2.1 implementation
37    /// (`lancefmt`) on the blocking pool; the `official-lance` feature opts
38    /// back into the official crate (see #75 M5).
39    #[cfg(feature = "official-lance")]
40    async fn write_lance_batch_async(&self, uri: String, batch: RecordBatch) -> StorageResult<()> {
41        info!("Writing Lance dataset to {}", uri);
42
43        let schema = batch.schema();
44        let batches = vec![batch];
45        let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema);
46
47        let params = WriteParams {
48            mode: WriteMode::Create,
49            ..WriteParams::default()
50        };
51
52        Dataset::write(reader, &uri, Some(params))
53            .await
54            .map_err(|e| StorageError::Lance(e.to_string()))?;
55
56        info!("Successfully wrote Lance dataset to {}", uri);
57        Ok(())
58    }
59
60    #[cfg(not(feature = "official-lance"))]
61    async fn write_lance_batch_async(&self, uri: String, batch: RecordBatch) -> StorageResult<()> {
62        info!("Writing Lance dataset (in-house v2.1) to {}", uri);
63        let path = uri_to_path(&uri)?;
64        tokio::task::spawn_blocking(move || crate::lancefmt::write_dataset(&batch, &path))
65            .await
66            .map_err(|e| StorageError::Io(format!("lancefmt writer task failed: {e}")))?
67    }
68
69    /// Async helper: read and concatenate all RecordBatches from a Lance dataset.
70    #[cfg(feature = "official-lance")]
71    async fn read_lance_all_batches_async(&self, uri: String) -> StorageResult<RecordBatch> {
72        info!("Reading Lance dataset from {}", uri);
73
74        let dataset = Dataset::open(&uri)
75            .await
76            .map_err(|e| StorageError::Lance(e.to_string()))?;
77        let scanner = dataset.scan();
78        let mut stream = scanner
79            .try_into_stream()
80            .await
81            .map_err(|e| StorageError::Lance(e.to_string()))?;
82
83        let mut batches = Vec::new();
84        while let Some(batch_result) = stream.next().await {
85            let batch = batch_result.map_err(|e| StorageError::Lance(e.to_string()))?;
86            batches.push(batch);
87        }
88
89        if batches.is_empty() {
90            return Err(StorageError::Invalid("Empty Lance dataset".into()));
91        }
92
93        let schema = batches[0].schema();
94        let combined = arrow::compute::concat_batches(&schema, &batches)
95            .map_err(|e| StorageError::Lance(format!("Failed to concatenate batches: {}", e)))?;
96
97        debug!(
98            "Combined Lance batch for {:?} has {} rows",
99            uri,
100            combined.num_rows()
101        );
102        Ok(combined)
103    }
104
105    #[cfg(not(feature = "official-lance"))]
106    async fn read_lance_all_batches_async(&self, uri: String) -> StorageResult<RecordBatch> {
107        info!("Reading Lance dataset (in-house v2.1) from {}", uri);
108        let path = uri_to_path(&uri)?;
109        let combined = tokio::task::spawn_blocking(move || crate::lancefmt::scan_all(&path))
110            .await
111            .map_err(|e| StorageError::Io(format!("lancefmt reader task failed: {e}")))??;
112        debug!(
113            "Combined Lance batch for {:?} has {} rows",
114            uri,
115            combined.num_rows()
116        );
117        Ok(combined)
118    }
119
120    /// Writes a single-column primitive vector as a Lance dataset named
121    /// `<name>_<key>.lance` and registers it in the metadata files map under
122    /// `key` with filetype "vector" and shape (`<len>`, 1).
123    ///
124    /// Shared implementation for the scalar `save_*` methods (lambdas,
125    /// vectors, indices, norms, centroid maps, cluster assignments).
126    async fn save_primitive_column<T: ArrowPrimitiveType>(
127        &self,
128        key: &str,
129        field_name: &str,
130        values: Vec<T::Native>,
131        md_path: &Path,
132    ) -> StorageResult<()>
133    where
134        Self: StorageBackend,
135    {
136        self.validate_initialized(md_path)?;
137        let path = self.file_path(key);
138        let len = values.len();
139        info!("Saving {} values for {} (field {})", len, key, field_name);
140
141        let schema = Schema::new(vec![Field::new(field_name, T::DATA_TYPE, false)]);
142        let batch = RecordBatch::try_new(
143            Arc::new(schema),
144            vec![Arc::new(PrimitiveArray::<T>::from_iter_values(values)) as ArrayRef],
145        )
146        .map_err(|e| StorageError::Lance(e.to_string()))?;
147
148        let mut metadata = self.load_metadata().await?;
149        metadata = metadata.add_file(
150            key,
151            FileInfo::new(
152                format!("{}_{}.lance", self.get_name(), key),
153                "vector",
154                (len, 1),
155                None,
156                None,
157            )?,
158        );
159        self.save_metadata(&metadata).await?;
160
161        let uri = Self::path_to_uri(&path)?;
162        self.write_lance_batch_async(uri, batch).await?;
163        info!("Vector {} saved successfully", key);
164        Ok(())
165    }
166}
genegraph_storage/traits/lance.rs

genegraph_storage/traits/
lance.rs
