diff --git a/Cargo.lock b/Cargo.lock index ae367e7d8c..15e3a41bbb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3964,9 +3964,12 @@ dependencies = [ name = "iceberg-examples" version = "0.10.0" dependencies = [ + "async-trait", + "datafusion", "futures", "iceberg", "iceberg-catalog-rest", + "iceberg-datafusion", "iceberg-storage-opendal", "tokio", ] diff --git a/crates/examples/Cargo.toml b/crates/examples/Cargo.toml index 0492fed343..5b05b68e3c 100644 --- a/crates/examples/Cargo.toml +++ b/crates/examples/Cargo.toml @@ -26,12 +26,19 @@ rust-version = { workspace = true } version = { workspace = true } [dependencies] +async-trait = { workspace = true } +datafusion = { workspace = true } futures = { workspace = true } iceberg = { workspace = true } iceberg-catalog-rest = { workspace = true } +iceberg-datafusion = { workspace = true } iceberg-storage-opendal = { workspace = true, optional = true } tokio = { workspace = true, features = ["full"] } +[[example]] +name = "datafusion-session-catalog" +path = "src/datafusion_session_catalog.rs" + [[example]] name = "rest-catalog-namespace" path = "src/rest_catalog_namespace.rs" diff --git a/crates/examples/README.md b/crates/examples/README.md index 335d2ea287..40801c69e8 100644 --- a/crates/examples/README.md +++ b/crates/examples/README.md @@ -17,5 +17,15 @@ ~ under the License. --> -Example usage codes for `iceberg-rust`. Currently, these examples can't run directly since it requires setting up of -environments for catalogs, for example, rest catalog server. \ No newline at end of file +Example usage code for `iceberg-rust`. + +The [`datafusion-session-catalog` example](src/datafusion_session_catalog.rs) is +self-contained. It demonstrates how to attach application-specific request metadata to a +DataFusion session, resolve it into an Iceberg `SessionContext`, and query a +session-catalog-backed `IcebergCatalogProvider`: + +```shell +cargo run -p iceberg-examples --example datafusion-session-catalog +``` + +The REST catalog examples require a catalog server and its supporting environment. diff --git a/crates/examples/src/datafusion_session_catalog.rs b/crates/examples/src/datafusion_session_catalog.rs new file mode 100644 index 0000000000..a52c665275 --- /dev/null +++ b/crates/examples/src/datafusion_session_catalog.rs @@ -0,0 +1,303 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Connects a session-aware Iceberg catalog to DataFusion. +//! +//! Run with: +//! +//! ```text +//! cargo run -p iceberg-examples --example datafusion-session-catalog +//! ``` +//! +//! The adapter at the bottom only makes the example self-contained. Applications +//! should pass their own `SessionCatalog` implementation to the provider. + +use std::collections::HashMap; +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion::catalog::Session as DataFusionSession; +use datafusion::error::{DataFusionError, Result as DataFusionResult}; +use datafusion::prelude::{SessionConfig, SessionContext as DataFusionSessionContext}; +use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalog, MemoryCatalogBuilder}; +use iceberg::spec::{NestedField, PrimitiveType, Schema, Type}; +use iceberg::table::Table; +use iceberg::{ + Catalog, CatalogBuilder, Namespace, NamespaceIdent, Result, SessionCatalog, + SessionContext as IcebergSessionContext, TableCommit, TableCreation, TableIdent, +}; +use iceberg_datafusion::{IcebergCatalogProvider, SessionContextResolver}; + +/// User metadata stored as an application-specific DataFusion extension. +#[derive(Debug)] +struct UserContext { + name: String, +} + +/// Maps the application's DataFusion user context to an Iceberg session. +#[derive(Debug)] +struct UserSessionContextResolver; + +impl SessionContextResolver for UserSessionContextResolver { + fn resolve(&self, session: &dyn DataFusionSession) -> DataFusionResult { + let user = session + .config() + .get_extension::() + .ok_or_else(|| { + DataFusionError::Configuration( + "the DataFusion session has no RequestContext extension".to_string(), + ) + })?; + + Ok(IcebergSessionContext::builder() + // Reusing the DataFusion session ID gives the catalog a stable key + // for session-scoped caches. + .session_id(session.session_id().to_string()) + .identity(user.name.to_string()) + .build()) + } +} + +#[tokio::main] +async fn main() -> std::result::Result<(), Box> { + let catalog = + init_catalog_with_table(TableIdent::from_strs(&["datafusion", "example"])?).await?; + + let session_catalog = Arc::new(ExampleSessionCatalog::new(catalog)); + // Provider construction discovers namespaces and tables with one stable, + // anonymous fallback context. Session-dependent catalogs must make that + // discovery set available to the fallback context. + let provider = IcebergCatalogProvider::try_new_with_session_catalog( + session_catalog, + Arc::new(UserSessionContextResolver), + ) + .await?; + + let config = SessionConfig::new().with_extension(Arc::new(UserContext { + name: "user123".to_string(), + })); + let datafusion = DataFusionSessionContext::new_with_config(config); + datafusion.register_catalog("iceberg", Arc::new(provider)); + + // Planning the scan invokes the resolver and forwards its Iceberg context + // to the session catalog's `load_table` operation. + datafusion + .sql("SELECT COUNT(*) AS event_count FROM iceberg.datafusion.example") + .await? + .show() + .await?; + + Ok(()) +} + +/// A small session-aware wrapper around the in-memory catalog used by this +/// standalone example. +/// +/// Real session catalogs can use the context for authorization, credentials, +/// configuration, and caching. This wrapper logs it, then delegates catalog +/// operations so the example needs no external service. +#[derive(Debug)] +struct ExampleSessionCatalog { + inner: MemoryCatalog, +} + +impl ExampleSessionCatalog { + fn new(inner: MemoryCatalog) -> Self { + Self { inner } + } + + fn log_context(context: &IcebergSessionContext, operation: &str) { + let identity = context.identity().unwrap_or(""); + println!( + "{operation}: session_id={}, identity={identity}", + context.session_id() + ); + } +} + +#[async_trait] +impl SessionCatalog for ExampleSessionCatalog { + async fn list_namespaces( + &self, + context: &IcebergSessionContext, + parent: Option<&NamespaceIdent>, + ) -> Result> { + Self::log_context(context, "list_namespaces"); + self.inner.list_namespaces(parent).await + } + + async fn create_namespace( + &self, + context: &IcebergSessionContext, + namespace: &NamespaceIdent, + properties: HashMap, + ) -> Result { + Self::log_context(context, "create_namespace"); + self.inner.create_namespace(namespace, properties).await + } + + async fn get_namespace( + &self, + context: &IcebergSessionContext, + namespace: &NamespaceIdent, + ) -> Result { + Self::log_context(context, "get_namespace"); + self.inner.get_namespace(namespace).await + } + + async fn namespace_exists( + &self, + context: &IcebergSessionContext, + namespace: &NamespaceIdent, + ) -> Result { + Self::log_context(context, "namespace_exists"); + self.inner.namespace_exists(namespace).await + } + + async fn update_namespace( + &self, + context: &IcebergSessionContext, + namespace: &NamespaceIdent, + properties: HashMap, + ) -> Result<()> { + Self::log_context(context, "update_namespace"); + self.inner.update_namespace(namespace, properties).await + } + + async fn drop_namespace( + &self, + context: &IcebergSessionContext, + namespace: &NamespaceIdent, + ) -> Result<()> { + Self::log_context(context, "drop_namespace"); + self.inner.drop_namespace(namespace).await + } + + async fn list_tables( + &self, + context: &IcebergSessionContext, + namespace: &NamespaceIdent, + ) -> Result> { + Self::log_context(context, "list_tables"); + self.inner.list_tables(namespace).await + } + + async fn create_table( + &self, + context: &IcebergSessionContext, + namespace: &NamespaceIdent, + creation: TableCreation, + ) -> Result { + Self::log_context(context, "create_table"); + self.inner.create_table(namespace, creation).await + } + + async fn load_table( + &self, + context: &IcebergSessionContext, + table: &TableIdent, + ) -> Result
{ + Self::log_context(context, "load_table"); + self.inner.load_table(table).await + } + + async fn drop_table(&self, context: &IcebergSessionContext, table: &TableIdent) -> Result<()> { + Self::log_context(context, "drop_table"); + self.inner.drop_table(table).await + } + + async fn purge_table(&self, context: &IcebergSessionContext, table: &TableIdent) -> Result<()> { + Self::log_context(context, "purge_table"); + self.inner.purge_table(table).await + } + + async fn table_exists( + &self, + context: &IcebergSessionContext, + table: &TableIdent, + ) -> Result { + Self::log_context(context, "table_exists"); + self.inner.table_exists(table).await + } + + async fn rename_table( + &self, + context: &IcebergSessionContext, + src: &TableIdent, + dest: &TableIdent, + ) -> Result<()> { + Self::log_context(context, "rename_table"); + self.inner.rename_table(src, dest).await + } + + async fn register_table( + &self, + context: &IcebergSessionContext, + table: &TableIdent, + metadata_location: String, + ) -> Result
{ + Self::log_context(context, "register_table"); + self.inner.register_table(table, metadata_location).await + } + + async fn update_table( + &self, + context: &IcebergSessionContext, + commit: TableCommit, + ) -> Result
{ + Self::log_context(context, "update_table"); + self.inner.update_table(commit).await + } +} + +async fn init_catalog_with_table(table_ident: TableIdent) -> Result { + let catalog = MemoryCatalogBuilder::default() + .load( + "memory", + HashMap::from([( + MEMORY_CATALOG_WAREHOUSE.to_string(), + "memory://session-catalog-example".to_string(), + )]), + ) + .await?; + + catalog + .create_namespace(&table_ident.namespace(), HashMap::new()) + .await?; + catalog + .create_table( + &table_ident.namespace(), + TableCreation::builder() + .name(table_ident.name().to_string()) + .schema( + Schema::builder() + .with_fields(vec![ + NestedField::required( + 1, + "event_id", + Type::Primitive(PrimitiveType::Long), + ) + .into(), + ]) + .build()?, + ) + .build(), + ) + .await?; + + Ok(catalog) +} diff --git a/crates/integrations/datafusion/public-api.txt b/crates/integrations/datafusion/public-api.txt index e197c2057d..780ad26341 100644 --- a/crates/integrations/datafusion/public-api.txt +++ b/crates/integrations/datafusion/public-api.txt @@ -76,7 +76,7 @@ impl core::fmt::Debug for iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_catalog::table::TableProvider for iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::insert_into<'life0, 'life1, 'async_trait>(&'life0 self, state: &'life1 dyn datafusion_session::session::Session, input: alloc::sync::Arc, _insert_op: datafusion_expr::logical_plan::dml::InsertOp) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait -pub fn iceberg_datafusion::IcebergTableProvider::scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, _state: &'life1 dyn datafusion_session::session::Session, projection: core::option::Option<&'life2 alloc::vec::Vec>, filters: &'life3 [datafusion_expr::expr::Expr], limit: core::option::Option) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait +pub fn iceberg_datafusion::IcebergTableProvider::scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, state: &'life1 dyn datafusion_session::session::Session, projection: core::option::Option<&'life2 alloc::vec::Vec>, filters: &'life3 [datafusion_expr::expr::Expr], limit: core::option::Option) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait pub fn iceberg_datafusion::IcebergTableProvider::schema(&self) -> arrow_schema::schema::SchemaRef pub fn iceberg_datafusion::IcebergTableProvider::supports_filters_pushdown(&self, filters: &[&datafusion_expr::expr::Expr]) -> datafusion_common::error::Result> pub fn iceberg_datafusion::IcebergTableProvider::table_type(&self) -> datafusion_expr::table_source::TableType @@ -93,7 +93,8 @@ impl datafusion_catalog::table::TableProviderFactory for iceberg_datafusion::tab pub fn iceberg_datafusion::table_provider_factory::IcebergTableProviderFactory::create<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, _state: &'life1 dyn datafusion_session::session::Session, cmd: &'life2 datafusion_expr::logical_plan::ddl::CreateExternalTable) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait pub struct iceberg_datafusion::IcebergCatalogProvider impl iceberg_datafusion::IcebergCatalogProvider -pub async fn iceberg_datafusion::IcebergCatalogProvider::try_new(client: alloc::sync::Arc) -> iceberg::error::Result +pub async fn iceberg_datafusion::IcebergCatalogProvider::try_new(catalog: alloc::sync::Arc) -> iceberg::error::Result +pub async fn iceberg_datafusion::IcebergCatalogProvider::try_new_with_session_catalog(catalog: alloc::sync::Arc, resolver: alloc::sync::Arc) -> iceberg::error::Result impl core::fmt::Debug for iceberg_datafusion::IcebergCatalogProvider pub fn iceberg_datafusion::IcebergCatalogProvider::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_catalog::catalog::CatalogProvider for iceberg_datafusion::IcebergCatalogProvider @@ -120,7 +121,7 @@ impl core::fmt::Debug for iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_catalog::table::TableProvider for iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::insert_into<'life0, 'life1, 'async_trait>(&'life0 self, state: &'life1 dyn datafusion_session::session::Session, input: alloc::sync::Arc, _insert_op: datafusion_expr::logical_plan::dml::InsertOp) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait -pub fn iceberg_datafusion::IcebergTableProvider::scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, _state: &'life1 dyn datafusion_session::session::Session, projection: core::option::Option<&'life2 alloc::vec::Vec>, filters: &'life3 [datafusion_expr::expr::Expr], limit: core::option::Option) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait +pub fn iceberg_datafusion::IcebergTableProvider::scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, state: &'life1 dyn datafusion_session::session::Session, projection: core::option::Option<&'life2 alloc::vec::Vec>, filters: &'life3 [datafusion_expr::expr::Expr], limit: core::option::Option) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait pub fn iceberg_datafusion::IcebergTableProvider::schema(&self) -> arrow_schema::schema::SchemaRef pub fn iceberg_datafusion::IcebergTableProvider::supports_filters_pushdown(&self, filters: &[&datafusion_expr::expr::Expr]) -> datafusion_common::error::Result> pub fn iceberg_datafusion::IcebergTableProvider::table_type(&self) -> datafusion_expr::table_source::TableType @@ -134,5 +135,7 @@ impl core::fmt::Debug for iceberg_datafusion::table_provider_factory::IcebergTab pub fn iceberg_datafusion::table_provider_factory::IcebergTableProviderFactory::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_catalog::table::TableProviderFactory for iceberg_datafusion::table_provider_factory::IcebergTableProviderFactory pub fn iceberg_datafusion::table_provider_factory::IcebergTableProviderFactory::create<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, _state: &'life1 dyn datafusion_session::session::Session, cmd: &'life2 datafusion_expr::logical_plan::ddl::CreateExternalTable) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +pub trait iceberg_datafusion::SessionContextResolver: core::fmt::Debug + core::marker::Send + core::marker::Sync +pub fn iceberg_datafusion::SessionContextResolver::resolve(&self, session: &dyn datafusion_session::session::Session) -> datafusion_common::error::Result pub fn iceberg_datafusion::from_datafusion_error(error: datafusion_common::error::DataFusionError) -> iceberg::error::Error pub fn iceberg_datafusion::to_datafusion_error(error: iceberg::error::Error) -> datafusion_common::error::DataFusionError diff --git a/crates/integrations/datafusion/src/catalog.rs b/crates/integrations/datafusion/src/catalog.rs deleted file mode 100644 index 2c6e1ff002..0000000000 --- a/crates/integrations/datafusion/src/catalog.rs +++ /dev/null @@ -1,93 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::collections::HashMap; -use std::sync::Arc; - -use datafusion::catalog::{CatalogProvider, SchemaProvider}; -use futures::future::try_join_all; -use iceberg::{Catalog, NamespaceIdent, Result}; - -use crate::schema::IcebergSchemaProvider; - -/// Provides an interface to manage and access multiple schemas -/// within an Iceberg [`Catalog`]. -/// -/// Acts as a centralized catalog provider that aggregates -/// multiple [`SchemaProvider`], each associated with distinct namespaces. -#[derive(Debug)] -pub struct IcebergCatalogProvider { - /// A `HashMap` where keys are namespace names - /// and values are dynamic references to objects implementing the - /// [`SchemaProvider`] trait. - schemas: HashMap>, -} - -impl IcebergCatalogProvider { - /// Asynchronously tries to construct a new [`IcebergCatalogProvider`] - /// using the given client to fetch and initialize schema providers for - /// each namespace in the Iceberg [`Catalog`]. - /// - /// This method retrieves the list of namespace names - /// attempts to create a schema provider for each namespace, and - /// collects these providers into a `HashMap`. - pub async fn try_new(client: Arc) -> Result { - // TODO: - // Schemas and providers should be cached and evicted based on time - // As of right now; schemas might become stale. - let schema_names: Vec<_> = client - .list_namespaces(None) - .await? - .iter() - .flat_map(|ns| ns.as_ref().clone()) - .collect(); - - let providers = try_join_all( - schema_names - .iter() - .map(|name| { - IcebergSchemaProvider::try_new( - client.clone(), - NamespaceIdent::new(name.clone()), - ) - }) - .collect::>(), - ) - .await?; - - let schemas: HashMap> = schema_names - .into_iter() - .zip(providers) - .map(|(name, provider)| { - let provider = Arc::new(provider) as Arc; - (name, provider) - }) - .collect(); - - Ok(IcebergCatalogProvider { schemas }) - } -} - -impl CatalogProvider for IcebergCatalogProvider { - fn schema_names(&self) -> Vec { - self.schemas.keys().cloned().collect() - } - - fn schema(&self, name: &str) -> Option> { - self.schemas.get(name).cloned() - } -} diff --git a/crates/integrations/datafusion/src/catalog_access.rs b/crates/integrations/datafusion/src/catalog_access.rs new file mode 100644 index 0000000000..1a995de112 --- /dev/null +++ b/crates/integrations/datafusion/src/catalog_access.rs @@ -0,0 +1,192 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::collections::HashMap; +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion::catalog::Session; +use datafusion::error::Result as DFResult; +use iceberg::table::Table; +use iceberg::{ + Catalog, Namespace, NamespaceIdent, Result, SessionCatalog, SessionContext, TableCommit, + TableCreation, TableIdent, +}; + +use crate::SessionContextResolver; + +/// Describes how the DataFusion integration accesses an Iceberg catalog. +/// +/// A catalog can either be accessed directly through [`Catalog`] or through a +/// [`SessionCatalog`] bound to an Iceberg [`SessionContext`]. Operations that +/// receive a DataFusion session resolve and bind its context once. Operations +/// for which DataFusion provides no session reuse one anonymous fallback +/// context so session-scoped catalog state remains stable across those calls. +#[derive(Clone, Debug)] +pub(crate) enum CatalogAccess { + /// A catalog accessed directly through the [`Catalog`] API. + Direct(Arc), + + /// A session-aware catalog, its resolver, and the stable context used by + /// DataFusion APIs that do not expose a session. + SessionAware { + catalog: Arc, + resolver: Arc, + fallback_context: SessionContext, + }, +} + +impl CatalogAccess { + pub(crate) fn with_session(&self, session: &dyn Session) -> DFResult> { + match self { + CatalogAccess::Direct(catalog) => Ok(Arc::clone(catalog)), + CatalogAccess::SessionAware { + catalog, + resolver: session_context_resolver, + .. + } => Ok(Arc::new(SessionBoundCatalog::new( + session_context_resolver.resolve(session)?, + Arc::clone(catalog), + ))), + } + } + + pub(crate) fn without_session(&self) -> Arc { + match self { + CatalogAccess::Direct(catalog) => Arc::clone(catalog), + CatalogAccess::SessionAware { + catalog, + fallback_context, + .. + } => Arc::new(SessionBoundCatalog::new( + fallback_context.clone(), + Arc::clone(catalog), + )), + } + } +} + +/// Adapts a [`SessionCatalog`] to [`Catalog`] by binding one [`SessionContext`]. +/// +/// Every catalog operation is forwarded to the inner session-aware catalog +/// with the same context. The binding is fixed for the lifetime of this +/// adapter; create another adapter to use a different session context. +#[derive(Debug)] +struct SessionBoundCatalog { + context: SessionContext, + inner: Arc, +} + +impl SessionBoundCatalog { + /// Creates a catalog view of `inner` bound to `context`. + /// + /// The inner catalog receives this context for every operation performed + /// through the returned adapter. + fn new(context: SessionContext, inner: Arc) -> Self { + Self { context, inner } + } +} + +#[async_trait] +impl Catalog for SessionBoundCatalog { + async fn list_namespaces( + &self, + parent: Option<&NamespaceIdent>, + ) -> Result> { + self.inner.list_namespaces(&self.context, parent).await + } + + async fn create_namespace( + &self, + namespace: &NamespaceIdent, + properties: HashMap, + ) -> Result { + self.inner + .create_namespace(&self.context, namespace, properties) + .await + } + + async fn get_namespace(&self, namespace: &NamespaceIdent) -> Result { + self.inner.get_namespace(&self.context, namespace).await + } + + async fn namespace_exists(&self, ns: &NamespaceIdent) -> Result { + self.inner.namespace_exists(&self.context, ns).await + } + + async fn update_namespace( + &self, + namespace: &NamespaceIdent, + properties: HashMap, + ) -> Result<()> { + self.inner + .update_namespace(&self.context, namespace, properties) + .await + } + + async fn drop_namespace(&self, namespace: &NamespaceIdent) -> Result<()> { + self.inner.drop_namespace(&self.context, namespace).await + } + + async fn list_tables(&self, namespace: &NamespaceIdent) -> Result> { + self.inner.list_tables(&self.context, namespace).await + } + + async fn create_table( + &self, + namespace: &NamespaceIdent, + creation: TableCreation, + ) -> Result
{ + self.inner + .create_table(&self.context, namespace, creation) + .await + } + + async fn load_table(&self, table_ident: &TableIdent) -> Result
{ + self.inner.load_table(&self.context, table_ident).await + } + + async fn drop_table(&self, table: &TableIdent) -> Result<()> { + self.inner.drop_table(&self.context, table).await + } + + async fn purge_table(&self, table: &TableIdent) -> Result<()> { + self.inner.purge_table(&self.context, table).await + } + + async fn table_exists(&self, table: &TableIdent) -> Result { + self.inner.table_exists(&self.context, table).await + } + + async fn rename_table(&self, src: &TableIdent, dest: &TableIdent) -> Result<()> { + self.inner.rename_table(&self.context, src, dest).await + } + + async fn register_table( + &self, + table_ident: &TableIdent, + metadata_location: String, + ) -> Result
{ + self.inner + .register_table(&self.context, table_ident, metadata_location) + .await + } + + async fn update_table(&self, commit: TableCommit) -> Result
{ + self.inner.update_table(&self.context, commit).await + } +} diff --git a/crates/integrations/datafusion/src/catalog_provider.rs b/crates/integrations/datafusion/src/catalog_provider.rs new file mode 100644 index 0000000000..5aedfad134 --- /dev/null +++ b/crates/integrations/datafusion/src/catalog_provider.rs @@ -0,0 +1,133 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::collections::HashMap; +use std::sync::Arc; + +use datafusion::catalog::{CatalogProvider, SchemaProvider}; +use futures::future::try_join_all; +use iceberg::{Catalog, NamespaceIdent, Result, SessionCatalog, SessionContext}; + +use crate::SessionContextResolver; +use crate::catalog_access::CatalogAccess; +use crate::schema_provider::IcebergSchemaProvider; + +/// Provides a DataFusion interface to schemas in an Iceberg [`Catalog`] or +/// [`SessionCatalog`]. +/// +/// Acts as a centralized catalog provider that aggregates +/// multiple [`SchemaProvider`], each associated with distinct namespaces. +#[derive(Debug)] +pub struct IcebergCatalogProvider { + /// A `HashMap` where keys are namespace names + /// and values are dynamic references to objects implementing the + /// [`SchemaProvider`] trait. + schemas: HashMap>, +} + +impl IcebergCatalogProvider { + /// Asynchronously constructs an [`IcebergCatalogProvider`] from a + /// [`Catalog`], fetching and initializing a schema provider for each + /// namespace. + /// + /// This method retrieves the namespace names and collects an initialized + /// schema provider for each namespace into a `HashMap`. + pub async fn try_new(catalog: Arc) -> Result { + let direct = CatalogAccess::Direct(catalog); + Self::try_new_with_access(direct).await + } + + /// Creates an [`IcebergCatalogProvider`] backed by a [`SessionCatalog`]. + /// + /// The [`SessionContextResolver`] derives an Iceberg session context for + /// scans and inserts. Provider initialization, metadata-table lookup, + /// table registration, and table deregistration do not receive a + /// DataFusion session; they share one anonymous fallback context instead. + /// + /// Namespace and table discovery is performed once during construction and + /// shared by all DataFusion sessions. Catalogs with session-dependent + /// visibility must make the intended discovery set available to the + /// anonymous fallback; discovery is not repeated per DataFusion session. + pub async fn try_new_with_session_catalog( + catalog: Arc, + resolver: Arc, + ) -> Result { + let session_aware = CatalogAccess::SessionAware { + catalog, + resolver, + // One session context that's shared for all query-unrelated catalog operations. + fallback_context: SessionContext::empty(), + }; + Self::try_new_with_access(session_aware).await + } + + async fn try_new_with_access(catalog: CatalogAccess) -> Result { + // TODO: + // Schemas and providers should be cached and evicted based on time + // As of right now; schemas might become stale. + let schema_names: Vec<_> = catalog + .without_session() + .list_namespaces(None) + .await? + .iter() + .flat_map(|ns| ns.as_ref().clone()) + .collect(); + + Ok(IcebergCatalogProvider { + schemas: load_schema_providers(catalog, schema_names).await?, + }) + } +} + +impl CatalogProvider for IcebergCatalogProvider { + fn schema_names(&self) -> Vec { + self.schemas.keys().cloned().collect() + } + + fn schema(&self, name: &str) -> Option> { + self.schemas.get(name).cloned() + } +} + +async fn load_schema_providers( + catalog_access: CatalogAccess, + schema_names: Vec, +) -> Result>> { + let iceberg_providers = try_join_all( + schema_names + .iter() + .map(|name| { + IcebergSchemaProvider::try_new( + catalog_access.clone(), + NamespaceIdent::new(name.clone()), + ) + }) + .collect::>(), + ) + .await?; + + let provider_map = schema_names + .into_iter() + .zip(iceberg_providers) + .map(|(name, iceberg_provider)| { + let provider = Arc::new(iceberg_provider) as Arc; + (name, provider) + }) + .collect(); + + Ok(provider_map) +} diff --git a/crates/integrations/datafusion/src/lib.rs b/crates/integrations/datafusion/src/lib.rs index 4b0ea8606d..c7ad5f8124 100644 --- a/crates/integrations/datafusion/src/lib.rs +++ b/crates/integrations/datafusion/src/lib.rs @@ -15,16 +15,36 @@ // specific language governing permissions and limitations // under the License. -mod catalog; -pub use catalog::*; - mod error; pub use error::*; +mod catalog_access; +mod catalog_provider; +pub use catalog_provider::*; pub mod physical_plan; -mod schema; +mod schema_provider; pub mod table; pub use table::table_provider_factory::IcebergTableProviderFactory; pub use table::*; pub(crate) mod task_writer; + +use std::fmt; + +use datafusion::catalog::Session as DFSession; +use datafusion::error::Result as DFResult; +use iceberg::SessionContext; + +/// Resolves an Iceberg [`SessionContext`] from a DataFusion session. +/// +/// The DataFusion integration calls the resolver once while planning each +/// session-aware scan or insert. The returned context is bound to that +/// operation and, for inserts, is retained through transaction commit. +/// Implementations should therefore return a stable Iceberg session identity +/// for repeated operations from the same DataFusion session. +pub trait SessionContextResolver: fmt::Debug + Send + Sync { + /// Returns the Iceberg context associated with `session`. + /// + /// Returning an error aborts planning before the catalog is accessed. + fn resolve(&self, session: &dyn DFSession) -> DFResult; +} diff --git a/crates/integrations/datafusion/src/physical_plan/commit.rs b/crates/integrations/datafusion/src/physical_plan/commit.rs index 9ae8b845ce..482414672f 100644 --- a/crates/integrations/datafusion/src/physical_plan/commit.rs +++ b/crates/integrations/datafusion/src/physical_plan/commit.rs @@ -42,6 +42,7 @@ use crate::to_datafusion_error; #[derive(Debug)] pub(crate) struct IcebergCommitExec { table: Table, + /// Catalog already bound to the session resolved during insert planning. catalog: Arc, input: Arc, schema: ArrowSchemaRef, @@ -50,7 +51,7 @@ pub(crate) struct IcebergCommitExec { } impl IcebergCommitExec { - pub fn new( + pub(crate) fn new( table: Table, catalog: Arc, input: Arc, @@ -287,6 +288,7 @@ mod tests { use iceberg::{Catalog, CatalogBuilder, NamespaceIdent, TableCreation, TableIdent}; use super::*; + use crate::catalog_access::CatalogAccess; use crate::physical_plan::DATA_FILES_COL_NAME; use crate::table::IcebergTableProvider; @@ -658,7 +660,7 @@ mod tests { ctx.register_table("source_table", source_table)?; let iceberg_table_provider = IcebergTableProvider::try_new( - catalog.clone(), + CatalogAccess::Direct(catalog), namespace.clone(), table_name.to_string(), ) diff --git a/crates/integrations/datafusion/src/schema.rs b/crates/integrations/datafusion/src/schema_provider.rs similarity index 92% rename from crates/integrations/datafusion/src/schema.rs rename to crates/integrations/datafusion/src/schema_provider.rs index 545863f8c6..0871f3b9f9 100644 --- a/crates/integrations/datafusion/src/schema.rs +++ b/crates/integrations/datafusion/src/schema_provider.rs @@ -23,23 +23,24 @@ use datafusion::catalog::SchemaProvider; use datafusion::datasource::TableProvider; use datafusion::error::{DataFusionError, Result as DFResult}; use datafusion::execution::TaskContext; -use datafusion::prelude::SessionContext; +use datafusion::prelude::SessionContext as DFSessionContext; use futures::StreamExt; use futures::future::try_join_all; use iceberg::arrow::arrow_schema_to_schema_auto_assign_ids; use iceberg::inspect::MetadataTableType; use iceberg::spec::FormatVersion; -use iceberg::{Catalog, Error, ErrorKind, NamespaceIdent, Result, TableCreation, TableIdent}; +use iceberg::{Error, ErrorKind, NamespaceIdent, Result, TableCreation, TableIdent}; +use crate::catalog_access::CatalogAccess; use crate::table::IcebergTableProvider; use crate::to_datafusion_error; -/// Represents a [`SchemaProvider`] for the Iceberg [`Catalog`], managing -/// access to table providers within a specific namespace. +/// Represents a [`SchemaProvider`] for an Iceberg catalog, managing table +/// providers within a specific namespace. #[derive(Debug)] pub(crate) struct IcebergSchemaProvider { - /// Reference to the Iceberg catalog - catalog: Arc, + /// Access policy for the Iceberg catalog. + catalog_access: CatalogAccess, /// The namespace this schema represents namespace: NamespaceIdent, /// A concurrent map where keys are table names @@ -54,18 +55,18 @@ impl IcebergSchemaProvider { /// using the given client to fetch and initialize table providers for /// the provided namespace in the Iceberg [`Catalog`]. /// - /// This method retrieves a list of table names - /// attempts to create a table provider for each table name, and - /// collects these providers into a `HashMap`. + /// This method retrieves a list of table names, attempts to create a table + /// provider for each name, and collects the providers into a [`DashMap`]. pub(crate) async fn try_new( - client: Arc, + catalog_access: CatalogAccess, namespace: NamespaceIdent, ) -> Result { // TODO: // Tables and providers should be cached based on table_name // if we have a cache miss; we update our internal cache & check again // As of right now; tables might become stale. - let table_names: Vec<_> = client + let table_names: Vec<_> = catalog_access + .without_session() .list_tables(&namespace) .await? .iter() @@ -75,7 +76,9 @@ impl IcebergSchemaProvider { let providers = try_join_all( table_names .iter() - .map(|name| IcebergTableProvider::try_new(client.clone(), namespace.clone(), name)) + .map(|name| { + IcebergTableProvider::try_new(catalog_access.clone(), namespace.clone(), name) + }) .collect::>(), ) .await?; @@ -86,7 +89,7 @@ impl IcebergSchemaProvider { } Ok(IcebergSchemaProvider { - catalog: client, + catalog_access, namespace, tables, }) @@ -171,7 +174,7 @@ impl SchemaProvider for IcebergSchemaProvider { .format_version(format_version) .build(); - let catalog = self.catalog.clone(); + let catalog_access = self.catalog_access.clone(); let namespace = self.namespace.clone(); let tables = self.tables.clone(); let name_clone = name.clone(); @@ -186,14 +189,15 @@ impl SchemaProvider for IcebergSchemaProvider { .await .map_err(to_datafusion_error)?; - catalog + catalog_access + .without_session() .create_table(&namespace, table_creation) .await .map_err(to_datafusion_error)?; - // Create a new table provider using the catalog reference + // Create a new table provider using the catalog access let table_provider = IcebergTableProvider::try_new( - catalog.clone(), + catalog_access, namespace.clone(), name_clone.clone(), ) @@ -220,7 +224,7 @@ impl SchemaProvider for IcebergSchemaProvider { return Ok(None); } - let catalog = self.catalog.clone(); + let catalog = self.catalog_access.without_session(); let namespace = self.namespace.clone(); let tables = self.tables.clone(); let table_name = name.to_string(); @@ -254,7 +258,7 @@ impl SchemaProvider for IcebergSchemaProvider { /// Verifies that a table provider contains no data by scanning with LIMIT 1. /// Returns an error if the table has any rows. async fn ensure_table_is_empty(table: &Arc) -> Result<()> { - let session_ctx = SessionContext::new(); + let session_ctx = DFSessionContext::new(); let exec_plan = table .scan(&session_ctx.state(), None, &[], Some(1)) .await @@ -317,9 +321,10 @@ mod tests { .await .unwrap(); - let provider = IcebergSchemaProvider::try_new(Arc::new(catalog), namespace) - .await - .unwrap(); + let provider = + IcebergSchemaProvider::try_new(CatalogAccess::Direct(Arc::new(catalog)), namespace) + .await + .unwrap(); (provider, temp_dir) } diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index 9de7bcb9c2..d53933cb00 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -45,9 +45,10 @@ use iceberg::arrow::schema_to_arrow_schema; use iceberg::inspect::MetadataTableType; use iceberg::spec::TableProperties; use iceberg::table::Table; -use iceberg::{Catalog, Error, ErrorKind, NamespaceIdent, Result, TableIdent}; +use iceberg::{Error, ErrorKind, NamespaceIdent, Result, TableIdent}; use metadata_table::IcebergMetadataTableProvider; +use crate::catalog_access::CatalogAccess; use crate::error::to_datafusion_error; use crate::physical_plan::commit::IcebergCommitExec; use crate::physical_plan::project::project_with_partition; @@ -58,16 +59,18 @@ use crate::physical_plan::write::IcebergWriteExec; /// Catalog-backed table provider with automatic metadata refresh. /// -/// This provider loads fresh table metadata from the catalog on every scan and write -/// operation, ensuring you always see the latest table state. Use this when you need -/// write operations or want to see the most up-to-date data. +/// This provider loads fresh table metadata from the catalog on every scan and +/// write operation, ensuring you always see the latest table state. A +/// session-aware provider binds the current DataFusion session for those +/// operations. Initial schema loading and metadata-table lookup do not receive +/// a DataFusion session and use the provider's shared anonymous context. /// /// For read-only access to a specific snapshot without catalog overhead, use /// [`IcebergStaticTableProvider`] instead. #[derive(Debug, Clone)] pub struct IcebergTableProvider { - /// The catalog that manages this table - catalog: Arc, + /// Access to the catalog that manages this table + catalog_access: CatalogAccess, /// The table identifier (namespace + name) table_ident: TableIdent, /// A reference-counted arrow `Schema` (cached at construction) @@ -80,18 +83,21 @@ impl IcebergTableProvider { /// Loads the table once to get the initial schema, then stores the catalog /// reference for future metadata refreshes on each operation. pub(crate) async fn try_new( - catalog: Arc, + catalog_access: CatalogAccess, namespace: NamespaceIdent, name: impl Into, ) -> Result { let table_ident = TableIdent::new(namespace, name.into()); // Load table once to get initial schema - let table = catalog.load_table(&table_ident).await?; + let table = catalog_access + .without_session() + .load_table(&table_ident) + .await?; let schema = Arc::new(schema_to_arrow_schema(table.metadata().current_schema())?); Ok(IcebergTableProvider { - catalog, + catalog_access, table_ident, schema, }) @@ -102,7 +108,11 @@ impl IcebergTableProvider { r#type: MetadataTableType, ) -> Result { // Load fresh table metadata for metadata table access - let table = self.catalog.load_table(&self.table_ident).await?; + let table = self + .catalog_access + .without_session() + .load_table(&self.table_ident) + .await?; Ok(IcebergMetadataTableProvider { table, r#type }) } } @@ -119,14 +129,15 @@ impl TableProvider for IcebergTableProvider { async fn scan( &self, - _state: &dyn Session, + state: &dyn Session, projection: Option<&Vec>, filters: &[Expr], limit: Option, ) -> DFResult> { // Load fresh table metadata from catalog let table = self - .catalog + .catalog_access + .with_session(state)? .load_table(&self.table_ident) .await .map_err(to_datafusion_error)?; @@ -162,9 +173,10 @@ impl TableProvider for IcebergTableProvider { ))); } + let catalog = self.catalog_access.with_session(state)?; + // Load fresh table metadata from catalog - let table = self - .catalog + let table = catalog .load_table(&self.table_ident) .await .map_err(to_datafusion_error)?; @@ -225,7 +237,7 @@ impl TableProvider for IcebergTableProvider { Ok(Arc::new(IcebergCommitExec::new( table, - self.catalog.clone(), + catalog, coalesce_partitions, self.schema.clone(), ))) @@ -374,7 +386,7 @@ mod tests { static_table.into_table() } - async fn get_test_catalog_and_table() -> (Arc, NamespaceIdent, String, TempDir) { + async fn get_test_catalog_and_table() -> (CatalogAccess, NamespaceIdent, String, TempDir) { let temp_dir = TempDir::new().unwrap(); let warehouse_path = temp_dir.path().to_str().unwrap().to_string(); @@ -414,7 +426,7 @@ mod tests { .unwrap(); ( - Arc::new(catalog), + CatalogAccess::Direct(Arc::new(catalog)), namespace, "test_table".to_string(), temp_dir, @@ -625,7 +637,7 @@ mod tests { async fn get_partitioned_test_catalog_and_table( fanout_enabled: Option, - ) -> (Arc, NamespaceIdent, String, TempDir) { + ) -> (CatalogAccess, NamespaceIdent, String, TempDir) { use iceberg::spec::{Transform, UnboundPartitionSpec}; let temp_dir = TempDir::new().unwrap(); @@ -682,7 +694,7 @@ mod tests { .unwrap(); ( - Arc::new(catalog), + CatalogAccess::Direct(Arc::new(catalog)), namespace, "partitioned_table".to_string(), temp_dir,