feat(datafusion) Add optional parameter to IcebergTableProvider so downstream Executors can re-register the Table

crates/integrations/datafusion/src/catalog.rs

@@ -22,6 +22,7 @@ use datafusion::catalog::{CatalogProvider, SchemaProvider};
 use futures::future::try_join_all;
 use iceberg::{Catalog, NamespaceIdent, Result};
 
+use crate::IcebergCatalogConfig;
 use crate::schema::IcebergSchemaProvider;
 
 /// Provides an interface to manage and access multiple schemas
@@ -46,6 +47,24 @@ impl IcebergCatalogProvider {
     /// attempts to create a schema provider for each namespace, and
     /// collects these providers into a `HashMap`.
     pub async fn try_new(client: Arc<dyn Catalog>) -> Result<Self> {
+        Self::try_new_impl(client, None).await
+    }
+
+    /// Like [`try_new`](Self::try_new), but threads a serializable
+    /// [`IcebergCatalogConfig`] into every schema and table provider it creates,
+    /// so the catalog's tables can be queried by a distributed engine such as
+    /// Ballista. The `client` must already be built from the same `config`.
+    pub async fn try_new_with_config(
+        client: Arc<dyn Catalog>,
+        config: IcebergCatalogConfig,
+    ) -> Result<Self> {
+        Self::try_new_impl(client, Some(config)).await
+    }
+
+    async fn try_new_impl(
+        client: Arc<dyn Catalog>,
+        config: Option<IcebergCatalogConfig>,
+    ) -> Result<Self> {
         // TODO:
         // Schemas and providers should be cached and evicted based on time
         // As of right now; schemas might become stale.
@@ -62,6 +81,7 @@ impl IcebergCatalogProvider {
                 .map(|name| {
                     IcebergSchemaProvider::try_new(
                         client.clone(),
+                        config.clone(),
                         NamespaceIdent::new(name.clone()),
                     )
                 })
@@ -91,3 +111,93 @@ impl CatalogProvider for IcebergCatalogProvider {
         self.schemas.get(name).cloned()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder};
+    use iceberg::spec::{NestedField, PrimitiveType, Schema, Type};
+    use iceberg::{CatalogBuilder, TableCreation};
+    use tempfile::TempDir;
+
+    use super::*;
+    use crate::table::IcebergTableProvider;
+
+    async fn catalog_with_one_table() -> (Arc<dyn Catalog>, TempDir) {
+        let temp_dir = TempDir::new().unwrap();
+        let warehouse_path = temp_dir.path().to_str().unwrap().to_string();
+        let catalog = MemoryCatalogBuilder::default()
+            .load(
+                "memory",
+                HashMap::from([(MEMORY_CATALOG_WAREHOUSE.to_string(), warehouse_path.clone())]),
+            )
+            .await
+            .unwrap();
+
+        let namespace = NamespaceIdent::new("test_ns".to_string());
+        catalog
+            .create_namespace(&namespace, HashMap::new())
+            .await
+            .unwrap();
+
+        let schema = Schema::builder()
+            .with_schema_id(0)
+            .with_fields(vec![
+                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
+            ])
+            .build()
+            .unwrap();
+        catalog
+            .create_table(
+                &namespace,
+                TableCreation::builder()
+                    .name("t".to_string())
+                    .location(format!("{warehouse_path}/t"))
+                    .schema(schema)
+                    .properties(HashMap::new())
+                    .build(),
+            )
+            .await
+            .unwrap();
+
+        (Arc::new(catalog), temp_dir)
+    }
+
+    /// Downcasts the `t` table provider in the `test_ns` schema and returns
+    /// whether it carries an [`IcebergCatalogConfig`].
+    async fn table_provider_has_config(provider: &IcebergCatalogProvider) -> bool {
+        let schema = provider.schema("test_ns").expect("test_ns schema");
+        let table = schema.table("t").await.unwrap().expect("table provider");
+        table
+            .downcast_ref::<IcebergTableProvider>()
+            .expect("IcebergTableProvider")
+            .catalog_config()
+            .is_some()
+    }
+
+    #[tokio::test]
+    async fn test_try_new_with_config_propagates_to_tables() {
+        let (catalog, _temp_dir) = catalog_with_one_table().await;
+        let config = IcebergCatalogConfig::new("memory", "memory", HashMap::new());
+
+        let provider = IcebergCatalogProvider::try_new_with_config(catalog, config)
+            .await
+            .unwrap();
+
+        assert!(
+            table_provider_has_config(&provider).await,
+            "try_new_with_config should thread the config down to table providers"
+        );
+    }
+
+    #[tokio::test]
+    async fn test_try_new_leaves_tables_config_less() {
+        let (catalog, _temp_dir) = catalog_with_one_table().await;
+
+        let provider = IcebergCatalogProvider::try_new(catalog).await.unwrap();
+
+        assert!(
+            !table_provider_has_config(&provider).await,
+            "try_new should leave table providers without a config"
+        );
+    }
+}

crates/integrations/datafusion/src/catalog_config.rs

@@ -0,0 +1,56 @@
+// 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;
+
+/// A serializable description of the catalog (and storage) that backs an
+/// [`IcebergTableProvider`](crate::table::IcebergTableProvider).
+///
+/// This is the minimal, self-contained handle needed to *reconstruct* a catalog
+/// and its associated `FileIO` on a remote node. It deliberately holds only
+/// plain data (no live connections) so that distributed query engines such as
+/// Ballista can serialize it, ship it to executors, and rebuild the catalog
+/// there via a catalog loader (e.g. `iceberg-catalog-loader`) and the storage
+/// via `FileIOBuilder::with_props`.
+///
+/// The `props` map carries both the catalog connection properties (e.g. the
+/// REST catalog URI) and the storage/`FileIO` properties (e.g. S3 endpoint and
+/// credentials); in practice these live together in a single map.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct IcebergCatalogConfig {
+    /// The catalog type, e.g. `"rest"`, `"sql"`, `"glue"`.
+    pub r#type: String,
+    /// The catalog name it is registered under (used to look it up when
+    /// reconstructing the catalog on a remote node).
+    pub name: String,
+    /// Catalog connection and storage properties.
+    pub props: HashMap<String, String>,
+}
+
+impl IcebergCatalogConfig {
+    pub fn new(
+        r#type: impl Into<String>,
+        name: impl Into<String>,
+        props: HashMap<String, String>,
+    ) -> Self {
+        Self {
+            r#type: r#type.into(),
+            name: name.into(),
+            props,
+        }
+    }
+}

crates/integrations/datafusion/src/lib.rs

@@ -18,6 +18,9 @@
 mod catalog;
 pub use catalog::*;
 
+mod catalog_config;
+pub use catalog_config::*;
+
 mod error;
 pub use error::*;
 

crates/integrations/datafusion/src/physical_plan/commit.rs

@@ -40,13 +40,17 @@ use crate::to_datafusion_error;
 /// IcebergCommitExec is responsible for collecting the files written and use
 /// [`Transaction::fast_append`] to commit the data files written.
 #[derive(Debug)]
-pub(crate) struct IcebergCommitExec {
+pub struct IcebergCommitExec {
     table: Table,
     catalog: Arc<dyn Catalog>,
     input: Arc<dyn ExecutionPlan>,
     schema: ArrowSchemaRef,
     count_schema: ArrowSchemaRef,
     plan_properties: Arc<PlanProperties>,
+    /// Optional serializable catalog/storage config, populated when this node is
+    /// built through a config-backed provider so it can be reconstructed on a
+    /// remote node by a distributed engine.
+    catalog_config: Option<crate::IcebergCatalogConfig>,
 }
 
 impl IcebergCommitExec {
@@ -67,9 +71,29 @@ impl IcebergCommitExec {
             schema,
             count_schema,
             plan_properties,
+            catalog_config: None,
         }
     }
 
+    /// Attaches a serializable catalog/storage config to this node so that a
+    /// distributed engine can reconstruct it (including the catalog) on a remote
+    /// node.
+    pub fn with_catalog_config(
+        mut self,
+        catalog_config: Option<crate::IcebergCatalogConfig>,
+    ) -> Self {
+        self.catalog_config = catalog_config;
+        self
+    }
+
+    pub fn catalog_config(&self) -> Option<&crate::IcebergCatalogConfig> {
+        self.catalog_config.as_ref()
+    }
+
+    pub fn table(&self) -> &Table {
+        &self.table
+    }
+
     // Compute the plan properties for this execution plan
     fn compute_properties(schema: ArrowSchemaRef) -> Arc<PlanProperties> {
         Arc::new(PlanProperties::new(
@@ -155,12 +179,15 @@ impl ExecutionPlan for IcebergCommitExec {
             )));
         }
 
-        Ok(Arc::new(IcebergCommitExec::new(
-            self.table.clone(),
-            self.catalog.clone(),
-            children[0].clone(),
-            self.schema.clone(),
-        )))
+        Ok(Arc::new(
+            IcebergCommitExec::new(
+                self.table.clone(),
+                self.catalog.clone(),
+                children[0].clone(),
+                self.schema.clone(),
+            )
+            .with_catalog_config(self.catalog_config.clone()),
+        ))
     }
 
     fn execute(
@@ -659,6 +686,7 @@ mod tests {
 
         let iceberg_table_provider = IcebergTableProvider::try_new(
             catalog.clone(),
+            None,
             namespace.clone(),
             table_name.to_string(),
         )

crates/integrations/datafusion/src/physical_plan/metadata_scan.rs

@@ -45,6 +45,13 @@ impl IcebergMetadataScan {
             properties,
         }
     }
+
+    /// Returns the metadata-table provider this node scans, so a distributed
+    /// engine can serialize the catalog config + table identifier + metadata type
+    /// it carries and rebuild it on a remote node.
+    pub fn provider(&self) -> &IcebergMetadataTableProvider {
+        &self.provider
+    }
 }
 
 impl DisplayAs for IcebergMetadataScan {

crates/integrations/datafusion/src/physical_plan/mod.rs

@@ -26,6 +26,9 @@ pub(crate) mod write;
 
 pub(crate) const DATA_FILES_COL_NAME: &str = "data_files";
 
+pub use commit::IcebergCommitExec;
 pub use expr_to_predicate::convert_filters_to_predicate;
-pub use project::project_with_partition;
+pub use metadata_scan::IcebergMetadataScan;
+pub use project::{PartitionExpr, project_with_partition};
 pub use scan::IcebergTableScan;
+pub use write::IcebergWriteExec;