genegraph_storage/traits/
lance.rs

1use crate::{StorageError, StorageResult};
2use arrow::record_batch::RecordBatch;
3use arrow::record_batch::RecordBatchIterator;
4use log::{debug, info};
5
6use futures::StreamExt;
7use lance::Dataset;
8use lance::dataset::{WriteMode, WriteParams};
9
10pub trait LanceStorage {
11    /// Async helper: write a RecordBatch to a Lance dataset.
12    async fn write_lance_batch_async(&self, uri: String, batch: RecordBatch) -> StorageResult<()> {
13        info!("Writing Lance dataset to {}", uri);
14
15        let schema = batch.schema();
16        let batches = vec![batch];
17        let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema);
18
19        let params = WriteParams {
20            mode: WriteMode::Create,
21            ..WriteParams::default()
22        };
23
24        Dataset::write(reader, &uri, Some(params))
25            .await
26            .map_err(|e| StorageError::Lance(e.to_string()))?;
27
28        info!("Successfully wrote Lance dataset to {}", uri);
29        Ok(())
30    }
31
32    /// Async helper: read and concatenate all RecordBatches from a Lance dataset.
33    async fn read_lance_all_batches_async(&self, uri: String) -> StorageResult<RecordBatch> {
34        info!("Reading Lance dataset from {}", uri);
35
36        let dataset = Dataset::open(&uri)
37            .await
38            .map_err(|e| StorageError::Lance(e.to_string()))?;
39        let scanner = dataset.scan();
40        let mut stream = scanner
41            .try_into_stream()
42            .await
43            .map_err(|e| StorageError::Lance(e.to_string()))?;
44
45        let mut batches = Vec::new();
46        while let Some(batch_result) = stream.next().await {
47            let batch = batch_result.map_err(|e| StorageError::Lance(e.to_string()))?;
48            batches.push(batch);
49        }
50
51        if batches.is_empty() {
52            return Err(StorageError::Invalid("Empty Lance dataset".into()));
53        }
54
55        let schema = batches[0].schema();
56        let combined = arrow::compute::concat_batches(&schema, &batches)
57            .map_err(|e| StorageError::Lance(format!("Failed to concatenate batches: {}", e)))?;
58
59        debug!(
60            "Combined Lance batch for {:?} has {} rows",
61            uri,
62            combined.num_rows()
63        );
64        Ok(combined)
65    }
66
67    /// Async helper: read the first RecordBatch from a Lance dataset.
68    async fn read_lance_first_batch_async(&self, uri: String) -> StorageResult<RecordBatch> {
69        info!("Reading first batch from Lance dataset {}", uri);
70
71        let dataset = Dataset::open(&uri)
72            .await
73            .map_err(|e| StorageError::Lance(e.to_string()))?;
74        let scanner = dataset.scan();
75        let mut stream = scanner
76            .try_into_stream()
77            .await
78            .map_err(|e| StorageError::Lance(e.to_string()))?;
79
80        let batch = stream
81            .next()
82            .await
83            .ok_or_else(|| StorageError::Lance("empty Lance dataset".to_string()))?
84            .map_err(|e| StorageError::Lance(e.to_string()))?;
85
86        debug!(
87            "Read first RecordBatch for path {:?} with {} rows",
88            uri,
89            batch.num_rows()
90        );
91        Ok(batch)
92    }
93}
genegraph_storage/traits/lance.rs

genegraph_storage/traits/
lance.rs
