Add support for Python based TableProviderFactory

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

Author: davisp

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,15 @@ 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.
 

python/datafusion/context.py

@@ -37,6 +37,7 @@
     CatalogProviderExportable,
     CatalogProviderList,
     CatalogProviderListExportable,
+    TableProviderFactory,
     TableProviderFactoryExportable,
 )
 from datafusion.dataframe import DataFrame
@@ -832,7 +833,9 @@ def deregister_table(self, name: str) -> None:
         self.ctx.deregister_table(name)
 
     def register_table_factory(
-        self, format: str, factory: TableProviderFactoryExportable
+        self,
+        format: str,
+        factory: TableProviderFactory | TableProviderFactoryExportable,
     ) -> None:
         """Register a :py:class:`~datafusion.TableProviderFactoryExportable`.
 

python/tests/test_catalog.py

@@ -120,6 +120,12 @@ def register_catalog(
         self.catalogs[name] = catalog
 
 
+class CustomTableProviderFactory(dfn.catalog.TableProviderFactory):
+    def create(self, cmd: dfn.expr.CreateExternalTable):
+        assert cmd.name() == "test_table_factory"
+        return create_dataset()
+
+
 def test_python_catalog_provider_list(ctx: SessionContext):
     ctx.register_catalog_provider_list(CustomCatalogProviderList())
 
@@ -314,3 +320,24 @@ def my_table_function_udtf() -> Table:
     assert len(result[0]) == 1
     assert len(result[0][0]) == 1
     assert result[0][0][0].as_py() == 3
+
+
+def test_register_python_table_provider_factory(ctx: SessionContext):
+    ctx.register_table_factory("CUSTOM_FACTORY", CustomTableProviderFactory())
+
+    ctx.sql("""
+        CREATE EXTERNAL TABLE test_table_factory
+        STORED AS CUSTOM_FACTORY
+        LOCATION foo;
+    """).collect()
+
+    result = ctx.sql("SELECT * FROM test_table_factory;").collect()
+
+    expect = [
+        pa.RecordBatch.from_arrays(
+            [pa.array([1, 2, 3]), pa.array([4, 5, 6])],
+            names=["a", "b"],
+        )
+    ]
+
+    assert result == expect

src/context.rs

@@ -78,7 +78,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;
@@ -663,22 +663,31 @@ impl PySessionContext {
     pub fn register_table_factory(
         &self,
         format: &str,
-        factory: Bound<'_, PyAny>,
+        mut factory: Bound<'_, PyAny>,
     ) -> PyDataFusionResult<()> {
-        let py = factory.py();
-        let codec_capsule = create_logical_extension_capsule(py, self.logical_codec.as_ref())?;
-
-        let capsule = factory
-            .getattr("__datafusion_table_provider_factory__")?
-            .call1((codec_capsule,))?;
-        let capsule = capsule.cast::<PyCapsule>().map_err(py_datafusion_err)?;
-        validate_pycapsule(capsule, "datafusion_table_provider_factory")?;
-
-        let factory: NonNull<FFI_TableProviderFactory> = capsule
-            .pointer_checked(Some(c_str!("datafusion_table_provider_factory")))?
-            .cast();
-        let factory = unsafe { factory.as_ref() };
-        let factory: Arc<dyn TableProviderFactory> = factory.into();
+        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();

src/table.rs

@@ -21,19 +21,23 @@ 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_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
 use pyo3::IntoPyObjectExt;
 use pyo3::prelude::*;
 
 use crate::context::PySessionContext;
 use crate::dataframe::PyDataFrame;
 use crate::dataset::Dataset;
-use crate::utils::table_provider_from_pycapsule;
+use crate::expr::create_external_table::PyCreateExternalTable;
+use crate::{errors, utils};
 
 /// This struct is used as a common method for all TableProviders,
 /// whether they refer to an FFI provider, an internally known
@@ -91,7 +95,7 @@ impl PyTable {
                 Some(session) => session,
                 None => PySessionContext::global_ctx()?.into_bound_py_any(obj.py())?,
             };
-            table_provider_from_pycapsule(obj.clone(), session)?
+            utils::table_provider_from_pycapsule(obj.clone(), session)?
         } {
             Ok(PyTable::from(provider))
         } else {
@@ -206,3 +210,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 = utils::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)
+        })
+    }
+}