feat: enhance read_table to accept custom table providers and update documentation

Author: kosiew

docs/source/user-guide/io/table_provider.rst

@@ -56,3 +56,13 @@ to the ``SessionContext``.
     ctx.register_table_provider("my_table", provider)
 
     ctx.table("my_table").show()
+
+If you already have a provider instance you can also use
+``SessionContext.read_table`` to obtain a :class:`~datafusion.DataFrame`
+directly without registering it first:
+
+.. code-block:: python
+
+    provider = MyTableProvider()
+    df = ctx.read_table(provider)
+    df.show()

python/datafusion/context.py

@@ -1163,14 +1163,23 @@ def read_avro(
             self.ctx.read_avro(str(path), schema, file_partition_cols, file_extension)
         )
 
-    def read_table(self, table: Table) -> DataFrame:
+    def read_table(
+        self, table: Table | TableProviderExportable
+    ) -> DataFrame:
         """Creates a :py:class:`~datafusion.dataframe.DataFrame` from a table.
 
-        For a :py:class:`~datafusion.catalog.Table` such as a
-        :py:class:`~datafusion.catalog.ListingTable`, create a
-        :py:class:`~datafusion.dataframe.DataFrame`.
+        Args:
+            table: Either a :py:class:`~datafusion.catalog.Table` (such as a
+                :py:class:`~datafusion.catalog.ListingTable`) or an object that
+                implements ``__datafusion_table_provider__`` and returns a
+                PyCapsule describing a custom table provider.
+
+        Returns:
+            A :py:class:`~datafusion.dataframe.DataFrame` backed by the
+            provided table provider.
         """
-        return DataFrame(self.ctx.read_table(table.table))
+        provider = table.table if isinstance(table, Table) else table
+        return DataFrame(self.ctx.read_table(provider))
 
     def execute(self, plan: ExecutionPlan, partitions: int) -> RecordBatchStream:
         """Execute the ``plan`` and return the results."""

python/tests/test_context.py

@@ -17,6 +17,7 @@
 import datetime as dt
 import gzip
 import pathlib
+from uuid import uuid4
 
 import pyarrow as pa
 import pyarrow.dataset as ds
@@ -113,6 +114,28 @@ def test_register_record_batches(ctx):
     assert result[0].column(1) == pa.array([-3, -3, -3])
 
 
+def test_read_table_accepts_table_provider(ctx):
+    batch = pa.RecordBatch.from_arrays(
+        [pa.array([1, 2]), pa.array(["x", "y"])],
+        names=["value", "label"],
+    )
+
+    ctx.register_record_batches("capsule_provider", [[batch]])
+
+    table = ctx.catalog().schema().table("capsule_provider")
+    provider = table.table
+
+    expected = pa.Table.from_batches([batch])
+
+    provider_result = pa.Table.from_batches(
+        ctx.read_table(provider).collect()
+    )
+    assert provider_result.equals(expected)
+
+    table_result = pa.Table.from_batches(ctx.read_table(table).collect())
+    assert table_result.equals(expected)
+
+
 def test_create_dataframe_registers_unique_table_name(ctx):
     # create a RecordBatch and register it as memtable
     batch = pa.RecordBatch.from_arrays(
@@ -484,8 +507,6 @@ def test_table_exist(ctx):
 
 
 def test_table_not_found(ctx):
-    from uuid import uuid4
-
     with pytest.raises(KeyError):
         ctx.table(f"not-found-{uuid4()}")
 

src/catalog.rs

@@ -17,7 +17,7 @@
 
 use crate::dataset::Dataset;
 use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionError, PyDataFusionResult};
-use crate::utils::{validate_pycapsule, wait_for_future};
+use crate::utils::{get_tokio_runtime, validate_pycapsule, wait_for_future};
 use async_trait::async_trait;
 use datafusion::catalog::{MemoryCatalogProvider, MemorySchemaProvider};
 use datafusion::common::DataFusionError;
@@ -34,6 +34,7 @@ use pyo3::types::PyCapsule;
 use pyo3::IntoPyObjectExt;
 use std::any::Any;
 use std::collections::HashSet;
+use std::ffi::CString;
 use std::sync::Arc;
 
 #[pyclass(name = "RawCatalog", module = "datafusion.catalog", subclass)]
@@ -261,6 +262,23 @@ impl PyTable {
         }
     }
 
+    fn __datafusion_table_provider__<'py>(
+        &self,
+        py: Python<'py>,
+    ) -> PyResult<Bound<'py, PyCapsule>> {
+        let name = CString::new("datafusion_table_provider").unwrap();
+        let runtime = get_tokio_runtime().0.handle().clone();
+
+        let provider = Arc::clone(&self.table);
+        let provider_ptr = Arc::into_raw(provider);
+        let provider: Arc<dyn TableProvider + Send> =
+            unsafe { Arc::from_raw(provider_ptr as *const (dyn TableProvider + Send)) };
+
+        let provider = FFI_TableProvider::new(provider, false, Some(runtime));
+
+        PyCapsule::new(py, provider, Some(name.clone()))
+    }
+
     fn __repr__(&self) -> PyResult<String> {
         let kind = self.kind();
         Ok(format!("Table(kind={kind})"))

src/context.rs

@@ -1102,9 +1102,29 @@ impl PySessionContext {
         Ok(PyDataFrame::new(df))
     }
 
-    pub fn read_table(&self, table: &PyTable) -> PyDataFusionResult<PyDataFrame> {
-        let df = self.ctx.read_table(table.table())?;
-        Ok(PyDataFrame::new(df))
+    pub fn read_table(&self, table: Bound<'_, PyAny>) -> PyDataFusionResult<PyDataFrame> {
+        if table.hasattr("__datafusion_table_provider__")? {
+            let capsule = table.getattr("__datafusion_table_provider__")?.call0()?;
+            let capsule = capsule.downcast::<PyCapsule>().map_err(py_datafusion_err)?;
+            validate_pycapsule(capsule, "datafusion_table_provider")?;
+
+            let provider = unsafe { capsule.reference::<FFI_TableProvider>() };
+            let provider: ForeignTableProvider = provider.into();
+
+            let df = self.ctx.read_table(Arc::new(provider))?;
+            Ok(PyDataFrame::new(df))
+        } else {
+            match table.extract::<PyTable>() {
+                Ok(py_table) => {
+                    let df = self.ctx.read_table(py_table.table())?;
+                    Ok(PyDataFrame::new(df))
+                }
+                Err(_) => Err(crate::errors::PyDataFusionError::Common(
+                    "Object must be a datafusion.Table or expose __datafusion_table_provider__()."
+                        .to_string(),
+                )),
+            }
+        }
     }
 
     fn __repr__(&self) -> PyResult<String> {