Skip to content

feat: expand ListingSchemaProvider to support register and deregister table #3150

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

feat: expand ListingSchemaProvider to support register and deregister table #3150

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
62 changes: 44 additions & 18 deletions crates/core/src/data_catalog/storage/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,9 @@ use futures::TryStreamExt;
use object_store::ObjectStore;

use crate::errors::DeltaResult;
use crate::open_table_with_storage_options;
use crate::storage::*;
use crate::table::builder::ensure_table_uri;
use crate::DeltaTableBuilder;
use crate::{storage::*, DeltaTable};

const DELTA_LOG_FOLDER: &str = "_delta_log";

Expand All @@ -36,7 +36,7 @@ pub struct ListingSchemaProvider {
/// Underlying object store
store: Arc<dyn ObjectStore>,
/// A map of table names to a fully quilfied storage location
tables: DashMap<String, String>,
tables: DashMap<String, Arc<dyn TableProvider>>,
/// Options used to create underlying object stores
storage_options: StorageOptions,
}
Expand Down Expand Up @@ -73,17 +73,42 @@ impl ListingSchemaProvider {
parent = p;
}
}

for table in tables.into_iter() {
let table_name = normalize_table_name(table)?;
let table_path = table
.to_str()
.ok_or_else(|| DataFusionError::Internal("Cannot parse file name!".to_string()))?
.to_string();
if !self.table_exist(&table_name) {
let table_url = format!("{}/{table_path}", self.authority);
self.tables.insert(table_name.to_string(), table_url);
let table_url = format!("{}/{}", self.authority, table_path);
let Ok(delta_table) = DeltaTableBuilder::from_uri(table_url)
.with_storage_options(self.storage_options.0.clone())
.build()
else {
continue;
};
let _ = self.register_table(table_name, Arc::new(delta_table));
}
}
Ok(())
}

/// Tables are not initialized but have a reference setup. To initialize the delta
/// table, the `load()` function must be called on the delta table. This function helps with
/// that and ensures the DashMap is updated
pub async fn load_table(&self, table_name: &str) -> datafusion::common::Result<()> {
if let Some(mut table) = self.tables.get_mut(&table_name.to_string()) {
if let Some(delta_table) = table.value().as_any().downcast_ref::<DeltaTable>() {
// If table has not yet been loaded, we remove it from the tables map and add it again
if delta_table.state.is_none() {
let mut delta_table = delta_table.clone();
delta_table.load().await?;
*table = Arc::from(delta_table);
}
}
}

Ok(())
}
Comment on lines +96 to 113
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the table function is the right place to create an up to date snapshot. see comment below.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So an idea could be to move some of the logic down to the table function and then ensure it is updated?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I like that idea - if this is what you meant, @roeap

Copy link
Contributor Author

@Nordalf Nordalf Apr 9, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@roeap - what do you think of just moving the code to the table func and remove the state.is_none() check. Then it will become:

    async fn table(
        &self,
        name: &str,
    ) -> datafusion::common::Result<Option<Arc<dyn TableProvider>>> {
        if let Some(mut table) = self.tables.get_mut(&name.to_string()) {
            if let Some(delta_table) = table.value().as_any().downcast_ref::<DeltaTable>() {

                let mut delta_table = delta_table.clone();
                delta_table.update().await?;
                *table = Arc::from(delta_table);
            }
             return Ok(Some(Arc::clone(table.value())));

        }
        Ok(None)
    }

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This should work ... eventually we may just want to do that in the scan method of the table provider, since right now we are kind-of locking in the version for every provider.

But that would go a bit to far for this PR ...

}
Expand Down Expand Up @@ -112,31 +137,31 @@ impl SchemaProvider for ListingSchemaProvider {
}

async fn table(&self, name: &str) -> datafusion_common::Result<Option<Arc<dyn TableProvider>>> {
let Some(location) = self.tables.get(name).map(|t| t.clone()) else {
let Some(provider) = self.tables.get(name).map(|t| t.clone()) else {
return Ok(None);
};
let provider =
open_table_with_storage_options(location, self.storage_options.0.clone()).await?;
Ok(Some(Arc::new(provider) as Arc<dyn TableProvider>))
Ok(Some(provider))
Comment on lines +140 to +143
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We need to make sure we are returning an up to date table here. Datafusion internally will not make any calls to update the log data o.a. so if any operations happend (including the ones we may have done in the same DF session) we would serve a stale snapshot.

}

fn register_table(
&self,
_name: String,
_table: Arc<dyn TableProvider>,
name: String,
table: Arc<dyn TableProvider>,
) -> datafusion_common::Result<Option<Arc<dyn TableProvider>>> {
Err(DataFusionError::Execution(
"schema provider does not support registering tables".to_owned(),
))
if !self.table_exist(name.as_str()) {
self.tables.insert(name, table.clone());
}
Ok(Some(table))
Comment on lines +151 to +154
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the sematics here are a little bit different, i.e. we should raise if the table already exists.

https://github.com/apache/datafusion/blob/784df33f8930f91eada0d67aa5acc25a4c25cea2/datafusion/catalog/src/schema.rs#L57-L69

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You are right on this one. I will raise the error but should the table be returned if registered then?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I will change the code later to:

        if self.table_exist(name.as_str()) {
            return exec_err!("The table {name} already exists");
        };
        Ok(self.tables.insert(name, table))

}

fn deregister_table(
&self,
_name: &str,
name: &str,
) -> datafusion_common::Result<Option<Arc<dyn TableProvider>>> {
Err(DataFusionError::Execution(
"schema provider does not support deregistering tables".to_owned(),
))
if let Some(table) = self.tables.remove(name) {
return Ok(Some(table.1));
}
Ok(None)
}

fn table_exist(&self, name: &str) -> bool {
Expand Down Expand Up @@ -177,6 +202,7 @@ mod tests {
async fn test_query_table() {
let schema = Arc::new(ListingSchemaProvider::try_new("../test/tests/data/", None).unwrap());
schema.refresh().await.unwrap();
schema.load_table("simple_table").await.unwrap();

let ctx = SessionContext::new();
let catalog = Arc::new(MemoryCatalogProvider::default());
Expand Down