lance-namespace-datafusion 8.0.0

Lance namespace integration with Apache DataFusion catalogs and schemas
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The Lance Authors

use std::any::Any;
use std::collections::HashSet;
use std::sync::Arc;

use dashmap::DashMap;
use datafusion::catalog::{CatalogProvider, CatalogProviderList, SchemaProvider};
use datafusion::error::Result;

#[allow(unused_imports)]
use crate::SessionBuilder;
use crate::namespace_level::NamespaceLevel;
use crate::schema::LanceSchemaProvider;

/// A dynamic [`CatalogProviderList`] that maps Lance namespaces to catalogs.
///
/// The underlying namespace must be a four-level namespace. It is explicitly configured
/// via [`SessionBuilder::with_root`], and each child namespace under this root is
/// automatically registered as a [`LanceCatalogProvider`].
///
/// This `CatalogProviderList` is optional when building a DataFusion `SessionContext`.
/// If not provided, you can still configure catalogs using
/// [`SessionBuilder::add_catalog`] or set a default catalog via
/// [`SessionBuilder::with_default_catalog`].
#[derive(Debug, Clone)]
pub struct LanceCatalogProviderList {
    /// Root Lance namespace used to resolve catalogs / schemas / tables.
    #[allow(dead_code)]
    ns_level: NamespaceLevel,
    /// Catalogs that have been loaded from the root namespace.
    ///
    /// Note: The values in this map may become stale over time, as there is currently
    /// no mechanism to automatically refresh or invalidate cached catalog providers.
    catalogs: DashMap<String, Arc<dyn CatalogProvider>>,
}

impl LanceCatalogProviderList {
    pub async fn try_new(namespace: NamespaceLevel) -> Result<Self> {
        let catalogs = DashMap::new();
        for child_namespace in namespace.children().await? {
            let catalog_name = child_namespace.name().to_string();
            let catalog_provider = Arc::new(LanceCatalogProvider::try_new(child_namespace).await?);
            catalogs.insert(catalog_name, catalog_provider as Arc<dyn CatalogProvider>);
        }

        Ok(Self {
            ns_level: namespace,
            catalogs,
        })
    }
}

impl CatalogProviderList for LanceCatalogProviderList {
    fn as_any(&self) -> &dyn Any {
        self
    }

    /// Adds a new catalog to this catalog list.
    /// If a catalog of the same name existed before, it is replaced in the list and returned.
    fn register_catalog(
        &self,
        name: String,
        catalog: Arc<dyn CatalogProvider>,
    ) -> Option<Arc<dyn CatalogProvider>> {
        self.catalogs.insert(name, catalog)
    }

    fn catalog_names(&self) -> Vec<String> {
        self.catalogs
            .iter()
            .map(|entry| entry.key().clone())
            .collect::<HashSet<_>>()
            .into_iter()
            .collect()
    }

    fn catalog(&self, name: &str) -> Option<Arc<dyn CatalogProvider>> {
        self.catalogs
            .get(name)
            .map(|entry| Arc::clone(entry.value()))
    }
}

/// A dynamic [`CatalogProvider`] that exposes the immediate child namespaces
/// of a Lance namespace as database schemas.
///
/// The underlying namespace must be a three-level namespace. It is either explicitly
/// registered via [`SessionBuilder::add_catalog`], or automatically created as part of
/// the catalog hierarchy when [`SessionBuilder::with_root`] is used.
/// Child namespaces are automatically loaded as [`LanceSchemaProvider`] instances.
#[derive(Debug, Clone)]
pub struct LanceCatalogProvider {
    #[allow(dead_code)]
    ns_level: NamespaceLevel,
    /// Note: The values in this map may become stale over time, as there is currently
    /// no mechanism to automatically refresh or invalidate cached schema providers.
    schemas: DashMap<String, Arc<dyn SchemaProvider>>,
}

impl LanceCatalogProvider {
    pub async fn try_new(namespace: NamespaceLevel) -> Result<Self> {
        let schemas = DashMap::new();
        for child_namespace in namespace.children().await? {
            let schema_name = child_namespace.name().to_string();
            let schema_provider = Arc::new(LanceSchemaProvider::try_new(child_namespace).await?);
            schemas.insert(schema_name, schema_provider as Arc<dyn SchemaProvider>);
        }

        Ok(Self {
            ns_level: namespace,
            schemas,
        })
    }
}

impl CatalogProvider for LanceCatalogProvider {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn schema_names(&self) -> Vec<String> {
        self.schemas
            .iter()
            .map(|entry| entry.key().clone())
            .collect::<HashSet<_>>()
            .into_iter()
            .collect()
    }

    fn schema(&self, schema_name: &str) -> Option<Arc<dyn SchemaProvider>> {
        self.schemas
            .get(schema_name)
            .map(|entry| Arc::clone(entry.value()))
    }

    fn register_schema(
        &self,
        name: &str,
        schema: Arc<dyn SchemaProvider>,
    ) -> Result<Option<Arc<dyn SchemaProvider>>> {
        Ok(self.schemas.insert(name.to_string(), schema))
    }
}