Enhance table_provider_from_pycapsule handling

Updated the table_provider_from_pycapsule function to directly
recognize raw PyCapsule objects, ensuring proper validation
before resorting to attribute-based discovery. Additionally,
added a reusable helper for creating dummy table-provider
capsules and implemented an integration test to confirm
that RawTable.from_table_provider_capsule is exposed,
while also testing SessionContext.read_table with a raw capsule.

Author: kosiew

python/tests/test_context.py

@@ -33,6 +33,20 @@
     literal,
 )
 
+_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]
+
+
+def _make_invalid_table_provider_capsule():
+    backing = ctypes.create_string_buffer(b"x")
+    capsule = _PYCAPSULE_NEW(
+        ctypes.cast(backing, ctypes.c_void_p),
+        b"datafusion_table_provider",
+        None,
+    )
+    return capsule, backing
+
 
 def test_create_context_no_args():
     SessionContext()
@@ -345,20 +359,12 @@ def test_read_table_from_dataset(ctx):
 def test_read_table_rejects_invalid_table_provider_capsule(ctx):
     class CapsuleContainer:
         def __init__(self) -> None:
-            self._buffer = ctypes.create_string_buffer(b"x")
+            self._buffers: list[ctypes.Array[ctypes.c_char]] = []
 
         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
-            )
+            capsule, backing = _make_invalid_table_provider_capsule()
+            self._buffers.append(backing)
+            return capsule
 
     container = CapsuleContainer()
 
@@ -371,6 +377,18 @@ def __datafusion_table_provider__(self) -> object:
         )
 
 
+def test_read_table_with_raw_table_provider_capsule(ctx):
+    df_internal = pytest.importorskip("datafusion._internal")
+    assert hasattr(
+        df_internal.catalog.RawTable, "from_table_provider_capsule"
+    )
+
+    capsule, _backing = _make_invalid_table_provider_capsule()
+
+    with pytest.raises(ValueError, match="missing a destructor"):
+        ctx.read_table(capsule)
+
+
 def test_deregister_table(ctx, database):
     default = ctx.catalog()
     public = default.schema("public")

src/utils.rs

@@ -185,7 +185,11 @@ pub(crate) fn table_provider_from_capsule(
 pub(crate) fn table_provider_from_pycapsule(
     obj: &Bound<PyAny>,
 ) -> PyResult<Option<Arc<dyn TableProvider>>> {
-    if obj.hasattr("__datafusion_table_provider__")? {
+    if let Ok(capsule) = obj.downcast::<PyCapsule>() {
+        let provider = table_provider_from_capsule(&capsule)?;
+
+        Ok(Some(provider))
+    } else if obj.hasattr("__datafusion_table_provider__")? {
         let capsule = obj.getattr("__datafusion_table_provider__")?.call0()?;
         let capsule = capsule.downcast::<PyCapsule>().map_err(py_datafusion_err)?;
         let provider = table_provider_from_capsule(&capsule)?;