ntjohnson1 commented on code in PR #1679:
URL:
https://github.com/apache/datafusion-python/pull/1679#discussion_r3797950895
##########
python/datafusion/context.py:
##########
@@ -1808,6 +1853,108 @@ def with_query_planner(
new.ctx = new_internal
return new
+ def with_extensions(
+ self, *extensions: SessionExtensionExportable
+ ) -> SessionContext:
+ """Create a new session context with the given extension bundles.
+
+ This is the preferred way to install FFI extensions that need a
+ task-context provider (extension codecs and query planners). Each
+ extension's ``__datafusion_session_extension__`` method is called with
+ the destination context so it can bind its components to that exact
+ context, then all components are installed in one step. This avoids
+ the pitfalls of chaining :py:meth:`with_logical_extension_codec`,
+ :py:meth:`with_physical_extension_codec`, and
+ :py:meth:`with_query_planner` by hand, where components can end up
+ bound to an intermediate context that is later garbage collected.
+
+ Codecs compose with the existing chain and with each other: extensions
+ are processed left to right and prepend to the codec chain, so codecs
+ from later extensions are consulted first. At most one extension may
+ supply a query planner. If none does, an existing FFI planner on the
+ source context is rebound to the final codec chains.
+
+ If any extension raises or returns invalid components, the source
+ context's state is left unchanged and the partially built destination
+ is discarded. Extension factories must treat the context they receive
+ as configuration-only: catalogs are shared with the source context, so
+ registering tables or otherwise mutating the context during binding is
+ not rolled back on failure.
+
+ The returned context is the strong owner of the installed components'
+ task-context providers. Keep it alive for as long as DataFrames or
+ plans derived from it are in use; FFI operations after the context is
+ collected raise an error.
+
+ Args:
+ extensions: One or more objects implementing
+ ``__datafusion_session_extension__`` (see
+ :py:class:`SessionExtensionExportable`).
Review Comment:
This is redundant with the type hint
##########
python/datafusion/context.py:
##########
@@ -1808,6 +1853,108 @@ def with_query_planner(
new.ctx = new_internal
return new
+ def with_extensions(
+ self, *extensions: SessionExtensionExportable
+ ) -> SessionContext:
+ """Create a new session context with the given extension bundles.
+
+ This is the preferred way to install FFI extensions that need a
+ task-context provider (extension codecs and query planners). Each
+ extension's ``__datafusion_session_extension__`` method is called with
+ the destination context so it can bind its components to that exact
+ context, then all components are installed in one step. This avoids
+ the pitfalls of chaining :py:meth:`with_logical_extension_codec`,
+ :py:meth:`with_physical_extension_codec`, and
+ :py:meth:`with_query_planner` by hand, where components can end up
+ bound to an intermediate context that is later garbage collected.
+
+ Codecs compose with the existing chain and with each other: extensions
+ are processed left to right and prepend to the codec chain, so codecs
+ from later extensions are consulted first. At most one extension may
+ supply a query planner. If none does, an existing FFI planner on the
+ source context is rebound to the final codec chains.
+
+ If any extension raises or returns invalid components, the source
+ context's state is left unchanged and the partially built destination
+ is discarded. Extension factories must treat the context they receive
+ as configuration-only: catalogs are shared with the source context, so
+ registering tables or otherwise mutating the context during binding is
+ not rolled back on failure.
+
+ The returned context is the strong owner of the installed components'
+ task-context providers. Keep it alive for as long as DataFrames or
+ plans derived from it are in use; FFI operations after the context is
+ collected raise an error.
+
+ Args:
+ extensions: One or more objects implementing
+ ``__datafusion_session_extension__`` (see
+ :py:class:`SessionExtensionExportable`).
+
+ Returns:
+ A new context with all extension components installed.
+
+ Raises:
+ TypeError: If an argument does not implement the protocol or
+ returns something other than a
+ :py:class:`SessionExtensionComponents`.
+ ValueError: If no extensions are given or more than one extension
+ supplies a query planner.
+
+ Examples:
+ >>> from my_extension import DistributedEngineExtension #
doctest: +SKIP
Review Comment:
It might be nice to have a little sample we import as a part of the pytest
conf so these examples don't go stale. Through the stacked PRs lots of doctest
skip
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
