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