feat: expose variety of features from DF54 update

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

@@ -17,9 +17,8 @@
 
 use std::sync::Arc;
 
-use datafusion_catalog::{TableFunctionImpl, TableProvider};
+use datafusion_catalog::{TableFunctionArgs, TableFunctionImpl, TableProvider};
 use datafusion_common::error::Result as DataFusionResult;
-use datafusion_expr::Expr;
 use datafusion_ffi::udtf::FFI_TableFunction;
 use datafusion_python_util::ffi_logical_codec_from_pycapsule;
 use pyo3::types::PyCapsule;
@@ -59,7 +58,7 @@ impl MyTableFunction {
 }
 
 impl TableFunctionImpl for MyTableFunction {
-    fn call(&self, _args: &[Expr]) -> DataFusionResult<Arc<dyn TableProvider>> {
+    fn call_with_args(&self, _args: TableFunctionArgs) -> DataFusionResult<Arc<dyn TableProvider>> {
         let provider = MyTableProvider::new(4, 3, 2).create_table()?;
         Ok(Arc::new(provider))
     }

python/datafusion/context.py

@@ -86,13 +86,16 @@
 
     import pandas as pd
     import polars as pl  # type: ignore[import]
+    from _typeshed import CapsuleType as _PyCapsule
 
     from datafusion.catalog import CatalogProvider, Table
     from datafusion.common import DFSchema
     from datafusion.expr import Expr, SortKey
     from datafusion.plan import ExecutionPlan, LogicalPlan
     from datafusion.user_defined import (
         AggregateUDF,
+        LogicalExtensionCodecExportable,
+        PhysicalExtensionCodecExportable,
         ScalarUDF,
         TableFunction,
         WindowUDF,
@@ -1744,11 +1747,15 @@ def __datafusion_logical_extension_codec__(self) -> Any:
         """Access the PyCapsule FFI_LogicalExtensionCodec."""
         return self.ctx.__datafusion_logical_extension_codec__()
 
-    def with_logical_extension_codec(self, codec: Any) -> SessionContext:
+    def with_logical_extension_codec(
+        self, codec: LogicalExtensionCodecExportable | _PyCapsule
+    ) -> SessionContext:
         """Create a new session context with specified codec.
 
         This only supports codecs that have been implemented using the
-        FFI interface.
+        FFI interface. ``codec`` must either be a raw ``FFI_LogicalExtensionCodec``
+        ``PyCapsule`` or an object exposing
+        ``__datafusion_logical_extension_codec__``.
         """
         new_internal = self.ctx.with_logical_extension_codec(codec)
         new = SessionContext.__new__(SessionContext)
@@ -1759,11 +1766,15 @@ def __datafusion_physical_extension_codec__(self) -> Any:
         """Access the PyCapsule FFI_PhysicalExtensionCodec."""
         return self.ctx.__datafusion_physical_extension_codec__()
 
-    def with_physical_extension_codec(self, codec: Any) -> SessionContext:
+    def with_physical_extension_codec(
+        self, codec: PhysicalExtensionCodecExportable | _PyCapsule
+    ) -> SessionContext:
         """Create a new session context with the specified physical codec.
 
         This only supports codecs that have been implemented using the
-        FFI interface.
+        FFI interface. ``codec`` must either be a raw
+        ``FFI_PhysicalExtensionCodec`` ``PyCapsule`` or an object exposing
+        ``__datafusion_physical_extension_codec__``.
         """
         new_internal = self.ctx.with_physical_extension_codec(codec)
         new = SessionContext.__new__(SessionContext)

python/datafusion/user_defined.py

@@ -113,6 +113,18 @@ def _is_pycapsule(value: object) -> TypeGuard[_PyCapsule]:
     return value.__class__.__name__ == "PyCapsule"
 
 
+class LogicalExtensionCodecExportable(Protocol):
+    """Type hint for objects exposing ``__datafusion_logical_extension_codec__``."""
+
+    def __datafusion_logical_extension_codec__(self) -> object: ...  # noqa: D105
+
+
+class PhysicalExtensionCodecExportable(Protocol):
+    """Type hint for objects exposing ``__datafusion_physical_extension_codec__``."""
+
+    def __datafusion_physical_extension_codec__(self) -> object: ...  # noqa: D105
+
+
 class ScalarUDF:
     """Class for performing scalar user-defined functions (UDF).