lance_io/object_store/providers/
memory.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::{collections::HashMap, sync::Arc};
5
6use crate::object_store::{
7    DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_LOCAL_BLOCK_SIZE, DEFAULT_MAX_IOP_SIZE, ObjectStore,
8    ObjectStoreParams, ObjectStoreProvider, StorageOptions,
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().cloned().unwrap_or_default());
23        let download_retry_count = storage_options.download_retry_count();
24        Ok(ObjectStore {
25            inner: Arc::new(InMemory::new()),
26            local_dir_operations: None,
27            scheme: String::from("memory"),
28            block_size,
29            max_iop_size: *DEFAULT_MAX_IOP_SIZE,
30            use_constant_size_upload_parts: false,
31            list_is_lexically_ordered: true,
32            io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM,
33            download_retry_count,
34            io_tracker: Default::default(),
35            store_prefix: self
36                .calculate_object_store_prefix(&base_path, params.storage_options())?,
37            // Listed in full: the store is already in memory, so a page costs no less than
38            // the directory does.
39            paginated_lister: None,
40        })
41    }
42
43    fn extract_path(&self, url: &Url) -> Result<Path> {
44        let mut output = String::new();
45        if let Some(domain) = url.domain() {
46            output.push_str(domain);
47        }
48        output.push_str(url.path());
49        // The in-memory store uses the Path directly as a key with no HTTP layer,
50        // so there is no re-encoding step and thus no double-encoding to avoid.
51        // Path::from also tolerates the empty segments that local temp paths embed.
52        Ok(Path::from(output))
53    }
54
55    fn calculate_object_store_prefix(
56        &self,
57        _url: &Url,
58        _storage_options: Option<&HashMap<String, String>>,
59    ) -> Result<String> {
60        Ok("memory".to_string())
61    }
62}
63
64#[cfg(test)]
65mod tests {
66    use super::*;
67
68    #[test]
69    fn test_memory_store_path() {
70        let provider = MemoryStoreProvider;
71
72        let url = Url::parse("memory://path/to/file").unwrap();
73        let path = provider.extract_path(&url).unwrap();
74        let expected_path = Path::from("path/to/file");
75        assert_eq!(path, expected_path);
76    }
77
78    #[test]
79    fn test_calculate_object_store_prefix() {
80        let provider = MemoryStoreProvider;
81        assert_eq!(
82            "memory",
83            provider
84                .calculate_object_store_prefix(&Url::parse("memory://etc").unwrap(), None)
85                .unwrap()
86        );
87    }
88}
lance_io/object_store/providers/memory.rs

lance_io/object_store/providers/
memory.rs
