Skip to main content

lance_namespace_datafusion/
namespace_level.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::sync::Arc;
5
6use lance::dataset::builder::DatasetBuilder;
7use lance::{Dataset, Result};
8use lance_namespace::LanceNamespace;
9use lance_namespace::models::{ListNamespacesRequest, ListTablesRequest};
10
11const DEFAULT_NAMESPACE_NAME: &str = "lance";
12
13/// Lightweight wrapper around a Lance namespace handle and identifier.
14#[derive(Debug, Clone)]
15pub struct NamespaceLevel {
16    root: Arc<dyn LanceNamespace>,
17    /// Full namespace identifier, e.g. [catalog, schema].
18    namespace_id: Option<Vec<String>>,
19}
20
21impl From<Arc<dyn LanceNamespace>> for NamespaceLevel {
22    fn from(lance_namespace: Arc<dyn LanceNamespace>) -> Self {
23        Self::from_root(Arc::clone(&lance_namespace))
24    }
25}
26
27impl From<(Arc<dyn LanceNamespace>, String)> for NamespaceLevel {
28    fn from(lance_namespace: (Arc<dyn LanceNamespace>, String)) -> Self {
29        Self::from_namespace(Arc::clone(&lance_namespace.0), vec![lance_namespace.1])
30    }
31}
32
33impl From<(Arc<dyn LanceNamespace>, Vec<String>)> for NamespaceLevel {
34    fn from(lance_namespace: (Arc<dyn LanceNamespace>, Vec<String>)) -> Self {
35        Self::from_namespace(Arc::clone(&lance_namespace.0), lance_namespace.1)
36    }
37}
38
39impl NamespaceLevel {
40    /// Construct a namespace rooted at the top-level Lance namespace.
41    pub fn from_root(root: Arc<dyn LanceNamespace>) -> Self {
42        Self {
43            root,
44            namespace_id: None,
45        }
46    }
47
48    /// Construct a namespace for a specific child identifier under the root.
49    pub fn from_namespace(root: Arc<dyn LanceNamespace>, namespace_id: Vec<String>) -> Self {
50        Self {
51            root,
52            namespace_id: Some(namespace_id),
53        }
54    }
55
56    /// Return the full namespace identifier.
57    pub fn id(&self) -> Vec<String> {
58        self.namespace_id.clone().unwrap_or_default()
59    }
60
61    /// Name for this namespace (last component or default).
62    pub fn name(&self) -> &str {
63        self.namespace_id
64            .as_deref()
65            .and_then(|v| v.last())
66            .map_or(DEFAULT_NAMESPACE_NAME, |relative_name| {
67                relative_name.as_str()
68            })
69    }
70
71    fn child_id(&self, child_name: String) -> Vec<String> {
72        match &self.namespace_id {
73            Some(namespace_id) => {
74                let mut child_namespace = namespace_id.clone();
75                child_namespace.push(child_name);
76                child_namespace
77            }
78            None => vec![child_name],
79        }
80    }
81
82    /// List direct child namespaces.
83    pub async fn children(&self) -> Result<Vec<Self>> {
84        let root = Arc::clone(&self.root);
85        let namespace_id = self.namespace_id.clone().unwrap_or_default();
86        let request = ListNamespacesRequest {
87            id: Some(namespace_id.clone()),
88            page_token: None,
89            limit: None,
90            ..Default::default()
91        };
92
93        let namespaces = root.list_namespaces(request).await?.namespaces;
94
95        Ok(namespaces
96            .into_iter()
97            .map(|relative_ns_id| {
98                Self::from_namespace(Arc::clone(&self.root), self.child_id(relative_ns_id))
99            })
100            .collect())
101    }
102
103    /// List table names under this namespace.
104    pub async fn tables(&self) -> Result<Vec<String>> {
105        let root = Arc::clone(&self.root);
106        let namespace_id = self.namespace_id.clone().unwrap_or_default();
107        let request = ListTablesRequest {
108            id: Some(namespace_id),
109            page_token: None,
110            limit: None,
111            ..Default::default()
112        };
113
114        root.list_tables(request).await.map(|resp| resp.tables)
115    }
116
117    /// Load a Lance dataset for the given table name in this namespace.
118    pub async fn load_dataset(&self, table_name: &str) -> Result<Dataset> {
119        DatasetBuilder::from_namespace(
120            Arc::clone(&self.root),
121            self.child_id(table_name.to_string()),
122        )
123        .await?
124        .load()
125        .await
126    }
127}