Skip to main content

lance_io/object_store/providers/
local.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    ObjectStore, ObjectStoreParams, ObjectStoreProvider, StorageOptions, DEFAULT_LOCAL_BLOCK_SIZE,
8    DEFAULT_LOCAL_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE,
9};
10use lance_core::error::Result;
11use lance_core::Error;
12use object_store::{local::LocalFileSystem, path::Path};
13use snafu::location;
14use url::Url;
15
16#[derive(Default, Debug)]
17pub struct FileStoreProvider;
18
19#[async_trait::async_trait]
20impl ObjectStoreProvider for FileStoreProvider {
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().cloned().unwrap_or_default());
24        let download_retry_count = storage_options.download_retry_count();
25        Ok(ObjectStore {
26            inner: Arc::new(LocalFileSystem::new()),
27            scheme: base_path.scheme().to_owned(),
28            block_size,
29            max_iop_size: *DEFAULT_MAX_IOP_SIZE,
30            use_constant_size_upload_parts: false,
31            list_is_lexically_ordered: false,
32            io_parallelism: DEFAULT_LOCAL_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        })
38    }
39
40    fn extract_path(&self, url: &Url) -> Result<Path> {
41        if let Ok(file_path) = url.to_file_path() {
42            if let Ok(path) = Path::from_absolute_path(&file_path) {
43                return Ok(path);
44            }
45        }
46
47        Path::parse(url.path()).map_err(|e| {
48            Error::invalid_input(
49                format!("Failed to parse path '{}': {}", url.path(), e),
50                location!(),
51            )
52        })
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(url.scheme().to_string())
61    }
62}
63
64#[cfg(test)]
65mod tests {
66    use crate::object_store::uri_to_url;
67
68    use super::*;
69
70    #[test]
71    fn test_file_store_path() {
72        let provider = FileStoreProvider;
73
74        let cases = [
75            ("file:///", ""),
76            ("file:///usr/local/bin", "usr/local/bin"),
77            ("file-object-store:///path/to/file", "path/to/file"),
78            ("file:///path/to/foo/../bar", "path/to/bar"),
79        ];
80
81        for (uri, expected_path) in cases {
82            let url = uri_to_url(uri).unwrap();
83            let path = provider.extract_path(&url).unwrap();
84            assert_eq!(path.as_ref(), expected_path, "uri: '{}'", uri);
85        }
86    }
87
88    #[test]
89    fn test_calculate_object_store_prefix() {
90        let provider = FileStoreProvider;
91        assert_eq!(
92            "file",
93            provider
94                .calculate_object_store_prefix(&Url::parse("file:///etc").unwrap(), None)
95                .unwrap()
96        );
97    }
98
99    #[test]
100    fn test_calculate_object_store_prefix_for_file_object_store() {
101        let provider = FileStoreProvider;
102        assert_eq!(
103            "file-object-store",
104            provider
105                .calculate_object_store_prefix(
106                    &Url::parse("file-object-store:///etc").unwrap(),
107                    None
108                )
109                .unwrap()
110        );
111    }
112
113    #[test]
114    #[cfg(windows)]
115    fn test_file_store_path_windows() {
116        let provider = FileStoreProvider;
117
118        let cases = [
119            (
120                "C:\\Users\\ADMINI~1\\AppData\\Local\\",
121                "C:/Users/ADMINI~1/AppData/Local",
122            ),
123            (
124                "C:\\Users\\ADMINI~1\\AppData\\Local\\..\\",
125                "C:/Users/ADMINI~1/AppData",
126            ),
127        ];
128
129        for (uri, expected_path) in cases {
130            let url = uri_to_url(uri).unwrap();
131            let path = provider.extract_path(&url).unwrap();
132            assert_eq!(path.as_ref(), expected_path);
133        }
134    }
135}