feat: Add FFI_TableProviderFactory support (#1396)

* feat: Add FFI_TableProviderFactory support

This wraps the new FFI_TableProviderFactory APIs in datafusion-ffi.

* Address PR comments

* Add support for Python based TableProviderFactory

This adds the ability to register Python based TableProviderFactory
instances to the SessionContext.

* Correction after rebase

---------

Co-authored-by: Tim Saucer <timsaucer@gmail.com>

Author: davisp

crates/core/src/context.rs

@@ -27,7 +27,7 @@ use arrow::pyarrow::FromPyArrow;
 use datafusion::arrow::datatypes::{DataType, Schema, SchemaRef};
 use datafusion::arrow::pyarrow::PyArrowType;
 use datafusion::arrow::record_batch::RecordBatch;
-use datafusion::catalog::{CatalogProvider, CatalogProviderList};
+use datafusion::catalog::{CatalogProvider, CatalogProviderList, TableProviderFactory};
 use datafusion::common::{ScalarValue, TableReference, exec_err};
 use datafusion::datasource::file_format::file_compression_type::FileCompressionType;
 use datafusion::datasource::file_format::parquet::ParquetFormat;
@@ -51,6 +51,7 @@ use datafusion_ffi::catalog_provider::FFI_CatalogProvider;
 use datafusion_ffi::catalog_provider_list::FFI_CatalogProviderList;
 use datafusion_ffi::execution::FFI_TaskContextProvider;
 use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
+use datafusion_ffi::table_provider_factory::FFI_TableProviderFactory;
 use datafusion_proto::logical_plan::DefaultLogicalExtensionCodec;
 use datafusion_python_util::{
     create_logical_extension_capsule, ffi_logical_codec_from_pycapsule, get_global_ctx,
@@ -81,7 +82,7 @@ use crate::record_batch::PyRecordBatchStream;
 use crate::sql::logical::PyLogicalPlan;
 use crate::sql::util::replace_placeholders_with_strings;
 use crate::store::StorageContexts;
-use crate::table::PyTable;
+use crate::table::{PyTable, RustWrappedPyTableProviderFactory};
 use crate::udaf::PyAggregateUDF;
 use crate::udf::PyScalarUDF;
 use crate::udtf::PyTableFunction;
@@ -659,6 +660,43 @@ impl PySessionContext {
         Ok(())
     }
 
+    pub fn register_table_factory(
+        &self,
+        format: &str,
+        mut factory: Bound<'_, PyAny>,
+    ) -> PyDataFusionResult<()> {
+        if factory.hasattr("__datafusion_table_provider_factory__")? {
+            let py = factory.py();
+            let codec_capsule = create_logical_extension_capsule(py, self.logical_codec.as_ref())?;
+            factory = factory
+                .getattr("__datafusion_table_provider_factory__")?
+                .call1((codec_capsule,))?;
+        }
+
+        let factory: Arc<dyn TableProviderFactory> =
+            if let Ok(capsule) = factory.cast::<PyCapsule>().map_err(py_datafusion_err) {
+                validate_pycapsule(capsule, "datafusion_table_provider_factory")?;
+
+                let data: NonNull<FFI_TableProviderFactory> = capsule
+                    .pointer_checked(Some(c_str!("datafusion_table_provider_factory")))?
+                    .cast();
+                let factory = unsafe { data.as_ref() };
+                factory.into()
+            } else {
+                Arc::new(RustWrappedPyTableProviderFactory::new(
+                    factory.into(),
+                    self.logical_codec.clone(),
+                ))
+            };
+
+        let st = self.ctx.state_ref();
+        let mut lock = st.write();
+        lock.table_factories_mut()
+            .insert(format.to_owned(), factory);
+
+        Ok(())
+    }
+
     pub fn register_catalog_provider_list(
         &self,
         mut provider: Bound<PyAny>,

crates/core/src/table.rs

@@ -21,19 +21,24 @@ use std::sync::Arc;
 use arrow::datatypes::SchemaRef;
 use arrow::pyarrow::ToPyArrow;
 use async_trait::async_trait;
-use datafusion::catalog::Session;
+use datafusion::catalog::{Session, TableProviderFactory};
 use datafusion::common::Column;
 use datafusion::datasource::{TableProvider, TableType};
-use datafusion::logical_expr::{Expr, LogicalPlanBuilder, TableProviderFilterPushDown};
+use datafusion::logical_expr::{
+    CreateExternalTable, Expr, LogicalPlanBuilder, TableProviderFilterPushDown,
+};
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion::prelude::DataFrame;
-use datafusion_python_util::table_provider_from_pycapsule;
+use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
+use datafusion_python_util::{create_logical_extension_capsule, table_provider_from_pycapsule};
 use pyo3::IntoPyObjectExt;
 use pyo3::prelude::*;
 
 use crate::context::PySessionContext;
 use crate::dataframe::PyDataFrame;
 use crate::dataset::Dataset;
+use crate::errors;
+use crate::expr::create_external_table::PyCreateExternalTable;
 
 /// This struct is used as a common method for all TableProviders,
 /// whether they refer to an FFI provider, an internally known
@@ -206,3 +211,51 @@ impl TableProvider for TempViewTable {
         Ok(vec![TableProviderFilterPushDown::Exact; filters.len()])
     }
 }
+
+#[derive(Debug)]
+pub(crate) struct RustWrappedPyTableProviderFactory {
+    pub(crate) table_provider_factory: Py<PyAny>,
+    pub(crate) codec: Arc<FFI_LogicalExtensionCodec>,
+}
+
+impl RustWrappedPyTableProviderFactory {
+    pub fn new(table_provider_factory: Py<PyAny>, codec: Arc<FFI_LogicalExtensionCodec>) -> Self {
+        Self {
+            table_provider_factory,
+            codec,
+        }
+    }
+
+    fn create_inner(
+        &self,
+        cmd: CreateExternalTable,
+        codec: Bound<PyAny>,
+    ) -> PyResult<Arc<dyn TableProvider>> {
+        Python::attach(|py| {
+            let provider = self.table_provider_factory.bind(py);
+            let cmd = PyCreateExternalTable::from(cmd);
+
+            provider
+                .call_method1("create", (cmd,))
+                .and_then(|t| PyTable::new(t, Some(codec)))
+                .map(|t| t.table())
+        })
+    }
+}
+
+#[async_trait]
+impl TableProviderFactory for RustWrappedPyTableProviderFactory {
+    async fn create(
+        &self,
+        _: &dyn Session,
+        cmd: &CreateExternalTable,
+    ) -> datafusion::common::Result<Arc<dyn TableProvider>> {
+        Python::attach(|py| {
+            let codec = create_logical_extension_capsule(py, self.codec.as_ref())
+                .map_err(errors::to_datafusion_err)?;
+
+            self.create_inner(cmd.clone(), codec.into_any())
+                .map_err(errors::to_datafusion_err)
+        })
+    }
+}

examples/datafusion-ffi-example/python/tests/_test_table_provider_factory.py

@@ -0,0 +1,41 @@
+# 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.
+
+from __future__ import annotations
+
+from datafusion import SessionContext
+from datafusion_ffi_example import MyTableProviderFactory
+
+
+def test_table_provider_factory_ffi() -> None:
+    ctx = SessionContext()
+    table = MyTableProviderFactory()
+
+    ctx.register_table_factory("MY_FORMAT", table)
+
+    # Create a new external table
+    ctx.sql("""
+        CREATE EXTERNAL TABLE
+        foo
+        STORED AS my_format
+        LOCATION '';
+    """).collect()
+
+    # Query the pre-populated table
+    result = ctx.sql("SELECT * FROM foo;").collect()
+    assert len(result) == 2
+    assert result[0].num_columns == 2

examples/datafusion-ffi-example/src/lib.rs

@@ -22,20 +22,23 @@ use crate::catalog_provider::{FixedSchemaProvider, MyCatalogProvider, MyCatalogP
 use crate::scalar_udf::IsNullUDF;
 use crate::table_function::MyTableFunction;
 use crate::table_provider::MyTableProvider;
+use crate::table_provider_factory::MyTableProviderFactory;
 use crate::window_udf::MyRankUDF;
 
 pub(crate) mod aggregate_udf;
 pub(crate) mod catalog_provider;
 pub(crate) mod scalar_udf;
 pub(crate) mod table_function;
 pub(crate) mod table_provider;
+pub(crate) mod table_provider_factory;
 pub(crate) mod window_udf;
 
 #[pymodule]
 fn datafusion_ffi_example(m: &Bound<'_, PyModule>) -> PyResult<()> {
     pyo3_log::init();
 
     m.add_class::<MyTableProvider>()?;
+    m.add_class::<MyTableProviderFactory>()?;
     m.add_class::<MyTableFunction>()?;
     m.add_class::<MyCatalogProvider>()?;
     m.add_class::<MyCatalogProviderList>()?;

examples/datafusion-ffi-example/src/table_provider_factory.rs

@@ -0,0 +1,87 @@
+// 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::sync::Arc;
+
+use async_trait::async_trait;
+use datafusion_catalog::{Session, TableProvider, TableProviderFactory};
+use datafusion_common::error::Result as DataFusionResult;
+use datafusion_expr::CreateExternalTable;
+use datafusion_ffi::table_provider_factory::FFI_TableProviderFactory;
+use datafusion_python_util::ffi_logical_codec_from_pycapsule;
+use pyo3::types::PyCapsule;
+use pyo3::{Bound, PyAny, PyResult, Python, pyclass, pymethods};
+
+use crate::catalog_provider;
+
+#[derive(Debug)]
+pub(crate) struct ExampleTableProviderFactory {}
+
+impl ExampleTableProviderFactory {
+    fn new() -> Self {
+        Self {}
+    }
+}
+
+#[async_trait]
+impl TableProviderFactory for ExampleTableProviderFactory {
+    async fn create(
+        &self,
+        _state: &dyn Session,
+        _cmd: &CreateExternalTable,
+    ) -> DataFusionResult<Arc<dyn TableProvider>> {
+        Ok(catalog_provider::my_table())
+    }
+}
+
+#[pyclass(
+    name = "MyTableProviderFactory",
+    module = "datafusion_ffi_example",
+    subclass
+)]
+#[derive(Debug)]
+pub struct MyTableProviderFactory {
+    inner: Arc<ExampleTableProviderFactory>,
+}
+
+impl Default for MyTableProviderFactory {
+    fn default() -> Self {
+        let inner = Arc::new(ExampleTableProviderFactory::new());
+        Self { inner }
+    }
+}
+
+#[pymethods]
+impl MyTableProviderFactory {
+    #[new]
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn __datafusion_table_provider_factory__<'py>(
+        &self,
+        py: Python<'py>,
+        codec: Bound<PyAny>,
+    ) -> PyResult<Bound<'py, PyCapsule>> {
+        let name = cr"datafusion_table_provider_factory".into();
+        let codec = ffi_logical_codec_from_pycapsule(codec)?;
+        let factory = Arc::clone(&self.inner) as Arc<dyn TableProviderFactory + Send>;
+        let factory = FFI_TableProviderFactory::new_with_ffi_codec(factory, None, codec);
+
+        PyCapsule::new(py, factory, Some(name))
+    }
+}

python/datafusion/catalog.py

@@ -29,6 +29,7 @@
 
     from datafusion import DataFrame, SessionContext
     from datafusion.context import TableProviderExportable
+    from datafusion.expr import CreateExternalTable
 
 try:
     from warnings import deprecated  # Python 3.13+
@@ -243,6 +244,24 @@ def kind(self) -> str:
         return self._inner.kind
 
 
+class TableProviderFactory(ABC):
+    """Abstract class for defining a Python based Table Provider Factory."""
+
+    @abstractmethod
+    def create(self, cmd: CreateExternalTable) -> Table:
+        """Create a table using the :class:`CreateExternalTable`."""
+        ...
+
+
+class TableProviderFactoryExportable(Protocol):
+    """Type hint for object that has __datafusion_table_provider_factory__ PyCapsule.
+
+    https://docs.rs/datafusion/latest/datafusion/catalog/trait.TableProviderFactory.html
+    """
+
+    def __datafusion_table_provider_factory__(self, session: Any) -> object: ...
+
+
 class CatalogProviderList(ABC):
     """Abstract class for defining a Python based Catalog Provider List."""
 

python/datafusion/context.py

@@ -37,6 +37,8 @@
     CatalogProviderExportable,
     CatalogProviderList,
     CatalogProviderListExportable,
+    TableProviderFactory,
+    TableProviderFactoryExportable,
 )
 from datafusion.dataframe import DataFrame
 from datafusion.expr import sort_list_to_raw_sort_list
@@ -830,6 +832,22 @@ def deregister_table(self, name: str) -> None:
         """Remove a table from the session."""
         self.ctx.deregister_table(name)
 
+    def register_table_factory(
+        self,
+        format: str,
+        factory: TableProviderFactory | TableProviderFactoryExportable,
+    ) -> None:
+        """Register a :py:class:`~datafusion.TableProviderFactoryExportable`.
+
+        The registered factory can be referenced from SQL DDL statements executed
+        against this context.
+
+        Args:
+            format: The value to be used in `STORED AS ${format}` clause.
+            factory: A PyCapsule that implements :class:`TableProviderFactoryExportable`
+        """
+        self.ctx.register_table_factory(format, factory)
+
     def catalog_names(self) -> set[str]:
         """Returns the list of catalogs in this context."""
         return self.ctx.catalog_names()