timsaucer commented on code in PR #1738:
URL:
https://github.com/apache/datafusion-python/pull/1738#discussion_r4221084117
##########
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:
_This reply was written by Claude (Claude Code) on behalf of @timsaucer._
It's a snapshot rather than a live adapter. `__datafusion_query_planner__`
clones the `Arc` of the session's *current* planner into an `FFI_QueryPlanner`
together with this handle's codecs, and the capsule owns that. So the fallback
keeps the old planner alive even after `set_session_query_planner` replaces it
on the session, which is exactly what lets the new planner delegate to it. The
only non-owning reference is the codecs' `FFI_TaskContextProvider`, which holds
the session weakly; that's why the returned context needs to be kept alive
(covered under "Keep your context alive" in the user guide).
##########
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:
_This reply was written by Claude (Claude Code) on behalf of @timsaucer._
Agreed with your second comment. I've removed the Rust-side `hasattr` in
4f308973. `_commit_extensions` is private with one caller, and the two checks
weren't equivalent anyway: `hasattr` accepts `__datafusion_session_planner__ =
None`, which `isinstance` against the runtime-checkable protocol rejects. Rust
now calls every entry it's handed, and
`test_with_extensions_ignores_a_planner_attribute_set_to_none` fails if the
Python filter is ever dropped.
##########
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:
_This reply was written by Claude (Claude Code) on behalf of @timsaucer._
Good catch. The existing test only put the `None` hook first, where
resetting and preserving the chain look the same. Pure Python can't tell the
two fallbacks apart (capsules are opaque), so I added it to the query-planner
example in 1c3c72b3: the nesting test now also runs with a `None`-returning
hook between the two planners and asserts the inner planner still runs.
Verified by mutating `_commit_extensions` to reset on `None`; the new case
fails with `plan_calls() == 0`.
--
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]
