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 object_store::{memory::InMemory, path::Path};
7use url::Url;
8
9use crate::object_store::{
10    ObjectStore, ObjectStoreParams, ObjectStoreProvider, StorageOptions, DEFAULT_LOCAL_BLOCK_SIZE,
11    DEFAULT_MAX_IOP_SIZE,
12};
13use lance_core::{error::Result, utils::tokio::get_num_compute_intensive_cpus};
14
15/// Provides a fresh in-memory object store for each call to `new_store`.
16#[derive(Default, Debug)]
17pub struct MemoryStoreProvider;
18
19#[async_trait::async_trait]
20impl ObjectStoreProvider for MemoryStoreProvider {
21    async fn new_store(&self, _base_path: Url, params: &ObjectStoreParams) -> Result<ObjectStore> {
22        let block_size = params.block_size.unwrap_or(DEFAULT_LOCAL_BLOCK_SIZE);
23        let storage_options = StorageOptions(params.storage_options.clone().unwrap_or_default());
24        let download_retry_count = storage_options.download_retry_count();
25        Ok(ObjectStore {
26            inner: Arc::new(InMemory::new()),
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: get_num_compute_intensive_cpus(),
33            download_retry_count,
34        })
35    }
36
37    fn extract_path(&self, url: &Url) -> Path {
38        let mut output = String::new();
39        if let Some(domain) = url.domain() {
40            output.push_str(domain);
41        }
42        output.push_str(url.path());
43        Path::from(output)
44    }
45}
46
47#[cfg(test)]
48mod tests {
49    use super::*;
50
51    #[test]
52    fn test_memory_store_path() {
53        let provider = MemoryStoreProvider;
54
55        let url = Url::parse("memory://path/to/file").unwrap();
56        let path = provider.extract_path(&url);
57        let expected_path = Path::from("path/to/file");
58        assert_eq!(path, expected_path);
59    }
60}
lance_io/object_store/providers/memory.rs

lance_io/object_store/providers/
memory.rs
