renato2099 commented on code in PR #1738:
URL:
https://github.com/apache/datafusion-python/pull/1738#discussion_r4054204192
##########
crates/core/src/context.rs:
##########
@@ -1717,49 +1716,68 @@ impl PySessionContext {
})
}
- /// Re-export a planner a `__datafusion_session_planner__` hook returned as
- /// a capsule, so the next hook in the chain receives one either way.
+ /// Run the planner hooks and commit a `with_extensions` call.
///
- /// A hook may hand back an object exposing `__datafusion_query_planner__`
- /// or a raw capsule; the next hook wraps whatever it is given and should
- /// not have to branch on which. Importing here also surfaces a malformed
- /// planner at the hook that produced it rather than at the final install.
- /// Writes nothing.
- pub fn _export_query_planner<'py>(
- slf: &Bound<'py, Self>,
- planner: Bound<'py, PyAny>,
- ) -> PyDataFusionResult<Bound<'py, PyCapsule>> {
- let ffi = ffi_query_planner_from_pycapsule(&planner,
Some(slf.as_any()))?;
- Ok(create_query_planner_capsule(slf.py(), &ffi)?)
- }
-
- /// Commit the query planner for a `with_extensions` call.
+ /// The second phase, run on the handle carrying the completed chains —
+ /// `session` is that same handle as the Python-level wrapper, which is
+ /// what each `__datafusion_session_planner__` hook receives. The hooks
+ /// run first, **in argument order**, each handed the planner built so
+ /// far as a capsule; a hook may hand back an object exposing
+ /// `__datafusion_query_planner__` or a raw capsule, and each return is
+ /// imported here so a malformed planner surfaces at the hook that
+ /// produced it rather than at the install. Returning `None` contributes
+ /// no planner. All of that writes nothing, so a hook that raises leaves
+ /// the session exactly as it was.
///
- /// The second phase, run once every codec is installed and every planner
- /// hook has returned, so the planner is bound against the final chains.
- /// This is the one call in `with_extensions` that writes to the session,
- /// and it goes through this context's own `state_ref()`, so providers
- /// bound to it stay valid.
+ /// Binding the planner is the commit, and it cannot fail. Anything else a
+ /// call installs is committed by the Python caller once this returns,
+ /// through `register_udf`, `register_udaf` and `register_udwf`, which
+ /// cannot fail. Whatever can fail belongs before this call — see
+ /// docs/source/contributor-guide/ffi-internals.md, under
+ /// "Why `with_extensions` commits last".
///
- /// `None` means no bundle supplied a planner. That still rebuilds
- /// whichever planner the session already holds against the new chains,
- /// exactly as `with_logical_extension_codec` does, and writes nothing at
- /// all if the session has no FFI planner to rebuild.
- ///
- /// The caller skips this step entirely when the call installed no codec
- /// and no planner, the same way [`Self::with_python_udf_inlining`] returns
- /// early for a no-op toggle: there is nothing to rebind against, and the
- /// rebuild would drag a planner sitting on another handle's codecs onto
- /// this one's.
- #[pyo3(signature = (planner=None))]
- pub fn _install_extension_planner<'py>(
+ /// The planner is bound through this context's own `state_ref()`, so
+ /// providers bound to it stay valid. With no planner supplied the bind
+ /// still rebuilds whichever planner the session already holds against
+ /// the new chains, exactly as `with_logical_extension_codec` does —
+ /// unless `rebind_planner` is also false, meaning the call installed no
+ /// codec either. Then the bind is skipped entirely, the same way
+ /// [`Self::with_python_udf_inlining`] returns early for a no-op toggle:
+ /// there is nothing to rebind against, and the rebuild would drag a
+ /// planner sitting on another handle's codecs onto this one's.
+ pub fn _commit_extensions<'py>(
slf: &Bound<'py, Self>,
- planner: Option<Bound<'py, PyAny>>,
+ extensions: Vec<Bound<'py, PyAny>>,
+ session: Bound<'py, PyAny>,
+ rebind_planner: bool,
) -> PyDataFusionResult<()> {
- let planner = planner
- .map(|planner| ffi_query_planner_from_pycapsule(&planner,
Some(slf.as_any())))
- .transpose()?;
- slf.borrow().set_session_query_planner(planner);
+ let py = slf.py();
+ // Nest the planners, outermost last. `planner` stays `None` when no
+ // bundle supplies one, which leaves an already-installed planner in
+ // place rather than wrapping the session's default in an FFI hop.
+ let mut planner: Option<FFI_QueryPlanner> = None;
+ for extension in &extensions {
+ if !extension.hasattr("__datafusion_session_planner__")? {
+ continue;
+ }
+ let fallback = match &planner {
+ Some(ffi) => create_query_planner_capsule(py, ffi)?,
+ None => slf.borrow().__datafusion_query_planner__(py, None)?,
Review Comment:
I am not very familiar with datafusion code base, so please excuse me if
this is something obvious.
My question is if we need to worry about the lifetime/ownership of the
object coming from `__datafusion_query_planner__`, I mean is it guaranteed that
it is just an adapter around the current session planner?
##########
python/datafusion/context.py:
##########
@@ -1975,70 +2133,72 @@ def with_extensions(
>>> batches[0].column(0).to_pylist() # doctest: +SKIP
[1]
"""
- for extension in extensions:
- if not isinstance(
- extension, (SessionComponentsExportable,
SessionPlannerExportable)
- ):
- msg = (
- "Extension implements neither "
- "__datafusion_session_components__ nor "
- f"__datafusion_session_planner__: {extension!r}"
- )
- raise TypeError(msg)
-
- # Phase one: collect every bundle's codecs. Components are bound
+ # Phase one: collect every bundle's components. Components are bound
# against this context, not a context derived from it. There is one
# `Arc<SessionContext>` per session, so a component bound here holds a
# task-context provider that the returned handle keeps alive.
- logical_codecs: list[LogicalExtensionCodecExportable] = []
- physical_codecs: list[PhysicalExtensionCodecExportable] = []
- for extension in extensions:
- if not isinstance(extension, SessionComponentsExportable):
- continue
- components = extension.__datafusion_session_components__(self)
- if not isinstance(components, SessionExtensionComponents):
- msg = (
- "__datafusion_session_components__ must return "
- "SessionExtensionComponents, got "
- f"{type(components).__name__} from {extension!r}"
- )
- raise TypeError(msg)
- logical_codecs.extend(components.logical_extension_codecs)
- physical_codecs.extend(components.physical_extension_codecs)
+ contributed = _collect_contributions(extensions, self)
# Writes nothing: the chains belong to the new handle, so a failure
# above or below leaves this context as it was.
new = SessionContext.__new__(SessionContext)
- new.ctx = self.ctx._install_extension_codecs(logical_codecs,
physical_codecs)
-
- # Phase two: nest the planners, outermost last. Each hook runs against
- # `new`, which carries the final chains, so a planner captured here
- # never sees a partial codec set. `planner` stays None when no bundle
- # supplies one, which leaves an already-installed planner in place
- # rather than wrapping the session's default in an FFI hop.
- planner: _PyCapsule | None = None
- for extension in extensions:
- if not isinstance(extension, SessionPlannerExportable):
- continue
- fallback = (
- planner
- if planner is not None
- else new.ctx.__datafusion_query_planner__()
- )
- supplied = extension.__datafusion_session_planner__(new, fallback)
- if supplied is None:
- continue
- planner = new.ctx._export_query_planner(supplied)
-
- # Rebinding the session's planner is a side effect on state shared with
- # every other handle, so do not pay it for a call that installs nothing
- # -- the same guard `with_python_udf_inlining` carries. With no codec
- # installed the chains the planner would be rebuilt against are the
ones
- # it already holds, so the rebuild is unobservable except in the one
case
- # where it does harm: a planner sitting on some *other* handle's codecs
- # gets dragged onto this handle's, silently undoing that install.
- if planner is not None or logical_codecs or physical_codecs:
- new.ctx._install_extension_planner(planner)
+ new.ctx = self.ctx._install_extension_codecs(
+ contributed.logical_codecs, contributed.physical_codecs
+ )
+
+ # Resolve every declared function to the wrapper that registers it, and
+ # settle name collisions, while a failure still costs nothing. None of
+ # these getters take an argument, so unlike a provider they do not care
+ # which handle they are resolved against. `user_defined` imports this
+ # module, so the import waits until here, with the cycle long settled.
+ from datafusion import user_defined as _ud # noqa: PLC0415
+
+ udfs = _resolve_declared_functions(
+ contributed.udfs,
+ _ud.ScalarUDF,
+ "__datafusion_scalar_udf__",
+ _ud.udf,
+ "scalar function",
+ )
+ udafs = _resolve_declared_functions(
+ contributed.udafs,
+ _ud.AggregateUDF,
+ "__datafusion_aggregate_udf__",
+ _ud.udaf,
+ "aggregate function",
+ )
+ udwfs = _resolve_declared_functions(
+ contributed.udwfs,
+ _ud.WindowUDF,
+ "__datafusion_window_udf__",
+ _ud.udwf,
+ "window function",
+ )
+
+ # Phase two: run the planner hooks, then bind the planner. Each hook
+ # runs against `new`, which carries the final chains, so a planner
+ # captured there never sees a partial codec set. The hook loop and the
+ # guard that skips the rebind for a call that installs nothing live on
+ # the Rust side -- see `_commit_extensions`. The list is narrowed with
+ # the same check the validation above uses, so one predicate decides
+ # both what is admitted and what is called.
+ new.ctx._commit_extensions(
+ [e for e in extensions if isinstance(e, SessionPlannerExportable)],
Review Comment:
thinking a bit more about this, there is one advantage of making the rust
side trust its inputs: stricter/cleaner invariant, i.e., the rust side is only
about execution/commit
##########
python/tests/test_context.py:
##########
@@ -1185,6 +1191,32 @@ def __datafusion_session_planner__(self, ctx, fallback):
assert batches[0].column(0) == pa.array([1])
+def test_with_extensions_ignores_a_planner_attribute_set_to_none(ctx):
Review Comment:
do we need a test for a callable planner hook returning `None` between two
planner-producing hooks? The main goal would be to test that returning `None`
preserves the previously accumulated planner .
##########
python/datafusion/context.py:
##########
@@ -1975,70 +2133,72 @@ def with_extensions(
>>> batches[0].column(0).to_pylist() # doctest: +SKIP
[1]
"""
- for extension in extensions:
- if not isinstance(
- extension, (SessionComponentsExportable,
SessionPlannerExportable)
- ):
- msg = (
- "Extension implements neither "
- "__datafusion_session_components__ nor "
- f"__datafusion_session_planner__: {extension!r}"
- )
- raise TypeError(msg)
-
- # Phase one: collect every bundle's codecs. Components are bound
+ # Phase one: collect every bundle's components. Components are bound
# against this context, not a context derived from it. There is one
# `Arc<SessionContext>` per session, so a component bound here holds a
# task-context provider that the returned handle keeps alive.
- logical_codecs: list[LogicalExtensionCodecExportable] = []
- physical_codecs: list[PhysicalExtensionCodecExportable] = []
- for extension in extensions:
- if not isinstance(extension, SessionComponentsExportable):
- continue
- components = extension.__datafusion_session_components__(self)
- if not isinstance(components, SessionExtensionComponents):
- msg = (
- "__datafusion_session_components__ must return "
- "SessionExtensionComponents, got "
- f"{type(components).__name__} from {extension!r}"
- )
- raise TypeError(msg)
- logical_codecs.extend(components.logical_extension_codecs)
- physical_codecs.extend(components.physical_extension_codecs)
+ contributed = _collect_contributions(extensions, self)
# Writes nothing: the chains belong to the new handle, so a failure
# above or below leaves this context as it was.
new = SessionContext.__new__(SessionContext)
- new.ctx = self.ctx._install_extension_codecs(logical_codecs,
physical_codecs)
-
- # Phase two: nest the planners, outermost last. Each hook runs against
- # `new`, which carries the final chains, so a planner captured here
- # never sees a partial codec set. `planner` stays None when no bundle
- # supplies one, which leaves an already-installed planner in place
- # rather than wrapping the session's default in an FFI hop.
- planner: _PyCapsule | None = None
- for extension in extensions:
- if not isinstance(extension, SessionPlannerExportable):
- continue
- fallback = (
- planner
- if planner is not None
- else new.ctx.__datafusion_query_planner__()
- )
- supplied = extension.__datafusion_session_planner__(new, fallback)
- if supplied is None:
- continue
- planner = new.ctx._export_query_planner(supplied)
-
- # Rebinding the session's planner is a side effect on state shared with
- # every other handle, so do not pay it for a call that installs nothing
- # -- the same guard `with_python_udf_inlining` carries. With no codec
- # installed the chains the planner would be rebuilt against are the
ones
- # it already holds, so the rebuild is unobservable except in the one
case
- # where it does harm: a planner sitting on some *other* handle's codecs
- # gets dragged onto this handle's, silently undoing that install.
- if planner is not None or logical_codecs or physical_codecs:
- new.ctx._install_extension_planner(planner)
+ new.ctx = self.ctx._install_extension_codecs(
+ contributed.logical_codecs, contributed.physical_codecs
+ )
+
+ # Resolve every declared function to the wrapper that registers it, and
+ # settle name collisions, while a failure still costs nothing. None of
+ # these getters take an argument, so unlike a provider they do not care
+ # which handle they are resolved against. `user_defined` imports this
+ # module, so the import waits until here, with the cycle long settled.
+ from datafusion import user_defined as _ud # noqa: PLC0415
+
+ udfs = _resolve_declared_functions(
+ contributed.udfs,
+ _ud.ScalarUDF,
+ "__datafusion_scalar_udf__",
+ _ud.udf,
+ "scalar function",
+ )
+ udafs = _resolve_declared_functions(
+ contributed.udafs,
+ _ud.AggregateUDF,
+ "__datafusion_aggregate_udf__",
+ _ud.udaf,
+ "aggregate function",
+ )
+ udwfs = _resolve_declared_functions(
+ contributed.udwfs,
+ _ud.WindowUDF,
+ "__datafusion_window_udf__",
+ _ud.udwf,
+ "window function",
+ )
+
+ # Phase two: run the planner hooks, then bind the planner. Each hook
+ # runs against `new`, which carries the final chains, so a planner
+ # captured there never sees a partial codec set. The hook loop and the
+ # guard that skips the rebind for a call that installs nothing live on
+ # the Rust side -- see `_commit_extensions`. The list is narrowed with
+ # the same check the validation above uses, so one predicate decides
+ # both what is admitted and what is called.
+ new.ctx._commit_extensions(
+ [e for e in extensions if isinstance(e, SessionPlannerExportable)],
Review Comment:
clarification question: is this check `if isinstance(e,
SessionPlannerExportable)` also doing the same check as the rust side
```
if !!extension.hasattr("__datafusion_session_planner__")? {
continue;
}
```
not asking to remove either as I think it is just safer to keep both checks
as we don't know other possible consumers of the rust side would do the right
thing
--
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]
