diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 644c7b445..14531a5f5 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -157,6 +157,29 @@ class QueryPlannerExportable(Protocol): def __datafusion_query_planner__(self, session: Any) -> 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.""" @@ -349,12 +372,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 eafcefdaf..afd78e445 100644 --- a/python/datafusion/user_defined.py +++ b/python/datafusion/user_defined.py @@ -1148,6 +1148,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). @@ -1158,7 +1167,7 @@ class TableFunction: def __init__( self, name: str, - func: Callable[..., Any], + func: Callable[..., Any] | TableFunctionExportable, ctx: SessionContext | None = None, *, with_session: bool = False, @@ -1217,6 +1226,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).