diff --git a/python/datafusion/context.py b/python/datafusion/context.py index bbf08e84e..a362fc29e 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -161,6 +161,29 @@ class PhysicalOptimizerRuleExportable(Protocol): def __datafusion_physical_optimizer_rule__(self) -> object: ... # noqa: D105 +class ExtensionOptionsExportable(Protocol): + """Type hint for object that has __datafusion_extension_options__ PyCapsule. + + The method returns a PyCapsule wrapping an ``FFI_ExtensionOptions``, + typically produced by a separate compiled extension and consumed by + :py:meth:`SessionConfig.with_extension`. + """ + + def __datafusion_extension_options__(self) -> object: ... # noqa: D105 + + +class TaskContextProviderExportable(Protocol): + """Type hint for object that has __datafusion_task_context_provider__ PyCapsule. + + The method returns a PyCapsule wrapping an ``FFI_TaskContextProvider``. + :py:class:`SessionContext` exposes one for its own task context; a + separate compiled extension can decode it (or one of its own) using + the matching Rust-side ``from_pycapsule`` helper. + """ + + def __datafusion_task_context_provider__(self) -> object: ... # noqa: D105 + + class SessionConfig: """Session configuration options.""" @@ -353,12 +376,14 @@ def set(self, key: str, value: str) -> SessionConfig: self.config_internal = self.config_internal.set(key, value) return self - def with_extension(self, extension: Any) -> SessionConfig: + def with_extension(self, extension: ExtensionOptionsExportable) -> SessionConfig: """Create a new configuration using an extension. Args: extension: A custom configuration extension object. These are - shared from another DataFusion extension library. + shared from another DataFusion extension library. It must expose + an ``__datafusion_extension_options__`` PyCapsule, see + :py:class:`ExtensionOptionsExportable`. Returns: A new :py:class:`SessionConfig` object with the updated setting. diff --git a/python/datafusion/user_defined.py b/python/datafusion/user_defined.py index f8d273177..5587ca73b 100644 --- a/python/datafusion/user_defined.py +++ b/python/datafusion/user_defined.py @@ -1180,6 +1180,15 @@ def adapter(*args: Any, session: Any, **kwargs: Any) -> Any: return adapter +class TableFunctionExportable(Protocol): + """Type hint for object that has __datafusion_table_function__ PyCapsule. + + https://datafusion.apache.org/python/user-guide/io/table_provider.html + """ + + def __datafusion_table_function__(self, session: Any) -> object: ... # noqa: D105 + + class TableFunction: """Class for performing user-defined table functions (UDTF). @@ -1190,7 +1199,7 @@ class TableFunction: def __init__( self, name: str, - func: Callable[..., Any], + func: Callable[..., Any] | TableFunctionExportable, ctx: SessionContext | None = None, *, with_session: bool = False, @@ -1249,6 +1258,10 @@ def udtf( with_session: bool = False, ) -> TableFunction: ... + @overload + @staticmethod + def udtf(func: TableFunctionExportable, name: str) -> TableFunction: ... + @staticmethod def udtf(*args: Any, with_session: bool = False, **kwargs: Any): """Create a new User-Defined Table Function (UDTF).