lance_io/object_store/providers/
memory.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::sync::Arc;
5
6use crate::object_store::{
7    ObjectStore, ObjectStoreParams, ObjectStoreProvider, StorageOptions,
8    DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_LOCAL_BLOCK_SIZE, DEFAULT_MAX_IOP_SIZE,
9};
10use lance_core::error::Result;
11use object_store::{memory::InMemory, path::Path};
12use url::Url;
13
14/// Provides a fresh in-memory object store for each call to `new_store`.
15#[derive(Default, Debug)]
16pub struct MemoryStoreProvider;
17
18#[async_trait::async_trait]
19impl ObjectStoreProvider for MemoryStoreProvider {
20    async fn new_store(&self, _base_path: Url, params: &ObjectStoreParams) -> Result<ObjectStore> {
21        let block_size = params.block_size.unwrap_or(DEFAULT_LOCAL_BLOCK_SIZE);
22        let storage_options = StorageOptions(params.storage_options.clone().unwrap_or_default());
23        let download_retry_count = storage_options.download_retry_count();
24        Ok(ObjectStore {
25            inner: Arc::new(InMemory::new()),
26            scheme: String::from("memory"),
27            block_size,
28            max_iop_size: *DEFAULT_MAX_IOP_SIZE,
29            use_constant_size_upload_parts: false,
30            list_is_lexically_ordered: true,
31            io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM,
32            download_retry_count,
33        })
34    }
35
36    fn extract_path(&self, url: &Url) -> Result<Path> {
37        let mut output = String::new();
38        if let Some(domain) = url.domain() {
39            output.push_str(domain);
40        }
41        output.push_str(url.path());
42        Ok(Path::from(output))
43    }
44}
45
46#[cfg(test)]
47mod tests {
48    use super::*;
49
50    #[test]
51    fn test_memory_store_path() {
52        let provider = MemoryStoreProvider;
53
54        let url = Url::parse("memory://path/to/file").unwrap();
55        let path = provider.extract_path(&url).unwrap();
56        let expected_path = Path::from("path/to/file");
57        assert_eq!(path, expected_path);
58    }
59}