lance_namespace_datafusion/
namespace_level.rs1use 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#[derive(Debug, Clone)]
15pub struct NamespaceLevel {
16 root: Arc<dyn LanceNamespace>,
17 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 pub fn from_root(root: Arc<dyn LanceNamespace>) -> Self {
42 Self {
43 root,
44 namespace_id: None,
45 }
46 }
47
48 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 pub fn id(&self) -> Vec<String> {
58 self.namespace_id.clone().unwrap_or_default()
59 }
60
61 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 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 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 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}