Harden conversion of datafusion_table_provider capsules

Enhance the safety of capsule conversions by verifying
destructors, pointer alignment, and necessary function
pointers pre-instantiation of ForeignTableProvider.

Introduced RawTable.from_table_provider_capsule for
reusability in capsule normalization within
SessionContext.read_table. Added a regression test
to ensure invalid capsules raise a ValueError instead
of causing segmentation faults.

Author: kosiew

python/datafusion/catalog.py

@@ -178,6 +178,14 @@ def from_dataset(dataset: pa.dataset.Dataset) -> Table:
         """Turn a :mod:`pyarrow.dataset` ``Dataset`` into a :class:`Table`."""
         return Table(dataset)
 
+    @staticmethod
+    def from_table_provider_capsule(capsule: object) -> Table:
+        """Wrap a validated table provider :class:`PyCapsule` as a :class:`Table`."""
+
+        return Table(
+            df_internal.catalog.RawTable.from_table_provider_capsule(capsule)
+        )
+
     @property
     def schema(self) -> pa.Schema:
         """Returns the schema associated with this table."""

python/datafusion/context.py

@@ -29,7 +29,7 @@
 
 import pyarrow as pa
 
-from datafusion.catalog import Catalog
+from datafusion.catalog import Catalog, Table
 from datafusion.dataframe import DataFrame
 from datafusion.expr import sort_list_to_raw_sort_list
 from datafusion.record_batch import RecordBatchStream
@@ -1181,6 +1181,16 @@ def read_table(
         self, table: Table | TableProviderExportable | DataFrame | pa.dataset.Dataset
     ) -> DataFrame:
         """Creates a :py:class:`~datafusion.dataframe.DataFrame` from a table."""
+        if not isinstance(table, Table):
+            capsule_obj: object | None = None
+            if hasattr(table, "__datafusion_table_provider__"):
+                capsule_obj = table.__datafusion_table_provider__()
+            elif table.__class__.__name__ == "PyCapsule":
+                capsule_obj = table
+
+            if capsule_obj is not None:
+                table = Table.from_table_provider_capsule(capsule_obj)
+
         return DataFrame(self.ctx.read_table(table))
 
     def execute(self, plan: ExecutionPlan, partitions: int) -> RecordBatchStream:

python/tests/test_context.py

@@ -14,6 +14,7 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+import ctypes
 import datetime as dt
 import gzip
 import pathlib
@@ -341,6 +342,35 @@ def test_read_table_from_dataset(ctx):
     assert result[0].column(1) == pa.array([4, 5, 6])
 
 
+def test_read_table_rejects_invalid_table_provider_capsule(ctx):
+    class CapsuleContainer:
+        def __init__(self) -> None:
+            self._buffer = ctypes.create_string_buffer(b"x")
+
+        def __datafusion_table_provider__(self) -> object:
+            pycapsule_new = ctypes.pythonapi.PyCapsule_New
+            pycapsule_new.restype = ctypes.py_object
+            pycapsule_new.argtypes = [
+                ctypes.c_void_p,
+                ctypes.c_char_p,
+                ctypes.c_void_p,
+            ]
+            dummy_ptr = ctypes.cast(self._buffer, ctypes.c_void_p)
+            return pycapsule_new(
+                dummy_ptr, b"datafusion_table_provider", None
+            )
+
+    container = CapsuleContainer()
+
+    with pytest.raises(ValueError, match="missing a destructor"):
+        ctx.read_table(container)
+
+    with pytest.raises(ValueError, match="missing a destructor"):
+        Table.from_table_provider_capsule(
+            container.__datafusion_table_provider__()
+        )
+
+
 def test_deregister_table(ctx, database):
     default = ctx.catalog()
     public = default.schema("public")

src/table.rs

@@ -18,11 +18,13 @@
 use arrow::pyarrow::ToPyArrow;
 use datafusion::datasource::{TableProvider, TableType};
 use pyo3::prelude::*;
+use pyo3::types::PyCapsule;
 use std::sync::Arc;
 
 use crate::dataframe::PyDataFrame;
 use crate::dataset::Dataset;
-use crate::utils::table_provider_from_pycapsule;
+use crate::errors::py_datafusion_err;
+use crate::utils::{table_provider_from_capsule, table_provider_from_pycapsule};
 
 /// This struct is used as a common method for all TableProviders,
 /// whether they refer to an FFI provider, an internally known
@@ -77,6 +79,13 @@ impl PyTable {
         }
     }
 
+    #[staticmethod]
+    pub fn from_table_provider_capsule(capsule: &Bound<'_, PyAny>) -> PyResult<Self> {
+        let capsule = capsule.downcast::<PyCapsule>().map_err(py_datafusion_err)?;
+        let provider = table_provider_from_capsule(&capsule)?;
+        Ok(PyTable::from(provider))
+    }
+
     /// Get a reference to the schema for this table
     #[getter]
     fn schema(&self, py: Python) -> PyResult<PyObject> {

src/utils.rs

@@ -27,7 +27,7 @@ use datafusion::{
 };
 use datafusion_ffi::table_provider::{FFI_TableProvider, ForeignTableProvider};
 use pyo3::prelude::*;
-use pyo3::{exceptions::PyValueError, types::PyCapsule};
+use pyo3::{exceptions::PyValueError, ffi, types::PyCapsule};
 use std::{
     future::Future,
     sync::{Arc, OnceLock},
@@ -124,18 +124,73 @@ pub(crate) fn validate_pycapsule(capsule: &Bound<PyCapsule>, name: &str) -> PyRe
     Ok(())
 }
 
+fn ensure_capsule_has_destructor(capsule: &Bound<PyCapsule>) -> PyResult<()> {
+    if unsafe { ffi::PyCapsule_GetDestructor(capsule.as_ptr()) }.is_none() {
+        return Err(PyValueError::new_err(
+            "Table provider capsule is missing a destructor; ensure it was created via datafusion_ffi's helpers.",
+        ));
+    }
+
+    Ok(())
+}
+
+fn ensure_capsule_pointer(capsule: &Bound<PyCapsule>) -> PyResult<()> {
+    let ptr = capsule.pointer();
+    if ptr.is_null() {
+        return Err(PyValueError::new_err(
+            "Table provider capsule contained a null pointer.",
+        ));
+    }
+
+    if (ptr as usize) % std::mem::align_of::<FFI_TableProvider>() != 0 {
+        return Err(PyValueError::new_err(
+            "Table provider capsule pointer was not aligned for FFI_TableProvider.",
+        ));
+    }
+
+    Ok(())
+}
+
+fn validate_foreign_table_provider(provider: &FFI_TableProvider) -> PyResult<()> {
+    if provider.schema as usize == 0
+        || provider.scan as usize == 0
+        || provider.table_type as usize == 0
+        || provider.clone as usize == 0
+        || provider.release as usize == 0
+        || provider.version as usize == 0
+        || provider.private_data.is_null()
+    {
+        return Err(PyValueError::new_err(
+            "Table provider capsule is missing required function pointers.",
+        ));
+    }
+
+    Ok(())
+}
+
+pub(crate) fn table_provider_from_capsule(
+    capsule: &Bound<PyCapsule>,
+) -> PyResult<Arc<dyn TableProvider>> {
+    validate_pycapsule(capsule, "datafusion_table_provider")?;
+    ensure_capsule_has_destructor(capsule)?;
+    ensure_capsule_pointer(capsule)?;
+
+    let provider = unsafe { capsule.reference::<FFI_TableProvider>() };
+    validate_foreign_table_provider(provider)?;
+    let provider: ForeignTableProvider = provider.into();
+
+    Ok(Arc::new(provider))
+}
+
 pub(crate) fn table_provider_from_pycapsule(
     obj: &Bound<PyAny>,
 ) -> PyResult<Option<Arc<dyn TableProvider>>> {
     if obj.hasattr("__datafusion_table_provider__")? {
         let capsule = obj.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 provider = table_provider_from_capsule(&capsule)?;
 
-        Ok(Some(Arc::new(provider)))
+        Ok(Some(provider))
     } else {
         Ok(None)
     }