lance_tools/
util.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use lance_core::{Error, Result};
5use lance_io::object_store::{ObjectStore, ObjectStoreParams, ObjectStoreRegistry};
6use object_store::path::Path;
7use snafu::location;
8use std::sync::Arc;
9use url::Url;
10
11fn path_to_parent(path: &Path) -> Result<(Path, String)> {
12    let mut parts = path.parts().collect::<Vec<_>>();
13    if parts.is_empty() {
14        return Err(Error::invalid_input(
15            format!("Path {} is not a valid path to a file", path),
16            location!(),
17        ));
18    }
19    let filename = parts.pop().unwrap().as_ref().to_owned();
20    Ok((Path::from_iter(parts), filename))
21}
22
23/// Get an object store and a path from a source string.
24pub(crate) async fn get_object_store_and_path(source: &String) -> Result<(Arc<ObjectStore>, Path)> {
25    if let Ok(mut url) = Url::parse(source) {
26        if url.scheme().len() > 1 {
27            let path = object_store::path::Path::parse(url.path()).map_err(Error::from)?;
28            let (parent_path, filename) = path_to_parent(&path)?;
29            url.set_path(parent_path.as_ref());
30            let object_store_registry = Arc::new(ObjectStoreRegistry::default());
31            let object_store_params = ObjectStoreParams::default();
32            let (object_store, dir_path) = ObjectStore::from_uri_and_params(
33                object_store_registry,
34                url.as_str(),
35                &object_store_params,
36            )
37            .await?;
38            let child_path = dir_path.child(filename);
39            return Ok((object_store, child_path));
40        }
41    }
42    let path = Path::from_filesystem_path(source)?;
43    let object_store = Arc::new(ObjectStore::local());
44    Ok((object_store, path))
45}
46
47#[cfg(test)]
48mod tests {
49    use super::*;
50
51    #[test]
52    fn test_path_to_parent() {
53        let (parent_path, filename) =
54            path_to_parent(&object_store::path::Path::parse("/a/b/c").unwrap()).unwrap();
55        assert_eq!("c", filename);
56        let parts: Vec<_> = parent_path.parts().collect();
57        assert_eq!(2, parts.len());
58        assert_eq!("a", parts.first().unwrap().as_ref());
59        assert_eq!("b", parts.get(1).unwrap().as_ref());
60    }
61}
lance_tools/util.rs

lance_tools/
util.rs
