#![allow(
clippy::cast_precision_loss,
clippy::cast_possible_truncation,
clippy::cast_sign_loss,
clippy::cast_possible_wrap,
clippy::float_cmp,
clippy::approx_constant
)]
use parking_lot::RwLock;
use std::io;
use std::sync::Arc;
use super::traits::VectorStorage;
use super::MmapStorage;
pub async fn reserve_capacity_async(
storage: Arc<RwLock<MmapStorage>>,
vector_count: usize,
) -> io::Result<()> {
tokio::task::spawn_blocking(move || {
let mut guard = storage.write();
guard.reserve_capacity(vector_count)
})
.await
.map_err(|e| io::Error::other(format!("Task join error: {e}")))?
}
pub async fn compact_async(storage: Arc<RwLock<MmapStorage>>) -> io::Result<usize> {
tokio::task::spawn_blocking(move || {
let mut guard = storage.write();
guard.compact()
})
.await
.map_err(|e| io::Error::other(format!("Task join error: {e}")))?
}
pub async fn flush_async(storage: Arc<RwLock<MmapStorage>>) -> io::Result<()> {
tokio::task::spawn_blocking(move || {
let mut guard = storage.write();
guard.flush()
})
.await
.map_err(|e| io::Error::other(format!("Task join error: {e}")))?
}
pub async fn store_batch_async(
storage: Arc<RwLock<MmapStorage>>,
vectors: Vec<(u64, Vec<f32>)>,
) -> io::Result<usize> {
tokio::task::spawn_blocking(move || {
let mut guard = storage.write();
let mut count = 0;
for (id, vector) in vectors {
guard.store(id, &vector)?;
count += 1;
}
Ok(count)
})
.await
.map_err(|e| io::Error::other(format!("Task join error: {e}")))?
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
#[tokio::test]
async fn test_reserve_capacity_async() {
let dir = TempDir::new().unwrap();
let storage = Arc::new(RwLock::new(MmapStorage::new(dir.path(), 128).unwrap()));
let result = reserve_capacity_async(storage, 1000).await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_store_batch_async() {
let dir = TempDir::new().unwrap();
let storage = Arc::new(RwLock::new(MmapStorage::new(dir.path(), 4).unwrap()));
let vectors: Vec<(u64, Vec<f32>)> = (0..100).map(|i| (i, vec![i as f32; 4])).collect();
let result = store_batch_async(storage.clone(), vectors).await;
assert!(result.is_ok());
assert_eq!(result.unwrap(), 100);
}
#[tokio::test]
async fn test_flush_async() {
let dir = TempDir::new().unwrap();
let storage = Arc::new(RwLock::new(MmapStorage::new(dir.path(), 128).unwrap()));
let result = flush_async(storage).await;
assert!(result.is_ok());
}
}