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 std::sync::Arc;
8use url::Url;
9
10fn path_to_parent(path: &Path) -> Result<(Path, String)> {
11    let mut parts = path.parts().collect::<Vec<_>>();
12    if parts.is_empty() {
13        return Err(Error::invalid_input(format!(
14            "Path {} is not a valid path to a file",
15            path
16        )));
17    }
18    let filename = parts.pop().unwrap().as_ref().to_owned();
19    Ok((Path::from_iter(parts), filename))
20}
21
22/// Get an object store and a path from a source string.
23pub(crate) async fn get_object_store_and_path(source: &String) -> Result<(Arc<ObjectStore>, Path)> {
24    if let Ok(mut url) = Url::parse(source)
25        && url.scheme().len() > 1
26    {
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.clone().join(filename);
39        return Ok((object_store, child_path));
40    }
41    let path = Path::from_filesystem_path(source)?;
42    let object_store = Arc::new(ObjectStore::local());
43    Ok((object_store, path))
44}
45
46#[cfg(test)]
47mod tests {
48    use super::*;
49
50    #[test]
51    fn test_path_to_parent() {
52        let (parent_path, filename) =
53            path_to_parent(&object_store::path::Path::parse("/a/b/c").unwrap()).unwrap();
54        assert_eq!("c", filename);
55        let parts: Vec<_> = parent_path.parts().collect();
56        assert_eq!(2, parts.len());
57        assert_eq!("a", parts.first().unwrap().as_ref());
58        assert_eq!("b", parts.get(1).unwrap().as_ref());
59    }
60}
lance_tools/util.rs

lance_tools/
util.rs
