You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "timsaucer (via GitHub)" <gi...@apache.org> on 2026/08/07 19:21:31 UTC
[PR] Add atomic SessionContext.with_extensions API [datafusion-python]
timsaucer opened a new pull request, #1679:
URL: https://github.com/apache/datafusion-python/pull/1679
# Which issue does this PR close?
Part 3 of 3 in the split of #1672 (part 1: FFI query planner support; part 2: composable extension codecs). Stacked on part 2. #1676 tracks extending this API to additional extension points.
# Rationale for this change
Working on the Ballista integration showed that chaining the low-level `with_*` methods is easy to get wrong: FFI codecs and planners carry a *weak* task-context provider bound to the context they were created against, so components can end up bound to an intermediate context that is later garbage collected. Queries then fail with `TaskContextProvider went out of scope over FFI boundary`, or worse, silently read stale session state.
# What changes are included in this PR?
- `SessionContext.with_extensions(*extensions)` installs one or more extension bundles atomically. Each bundle implements the new `__datafusion_session_extension__(ctx)` protocol: it receives the destination context, creates fresh components bound to that exact context, and returns them as a `SessionExtensionComponents` (new public dataclass; `SessionExtensionExportable` is the matching typing protocol).
- Internally one destination context is created, factories bind against it, all codecs are installed, and the planner is bound against the final codec chains in a single state write through that context's own `state_ref()`. No context is derived after any provider is created, so every weak provider targets the returned context.
- Validation happens before any state change: a failing factory or invalid capsule leaves the source context untouched.
- If no bundle supplies a planner, an existing FFI planner on the source context is rebound to the new codec chains; at most one bundle may supply a planner.
- The ownership contract is documented and tested: the returned context is the strong owner, and a `DataFrame` outliving it fails with a clean out-of-scope error rather than crashing.
- `MyPlannerExtension` in the example crate is a complete Rust implementation of the protocol, including extracting the host's task-context provider from the supplied context. Its codecs record the config value they resolve through the weak provider, letting tests prove the provider targets the returned context rather than the source.
- `docs/source/contributor-guide/ffi.md` documents `with_extensions` as the preferred API for extension bundles, keeps low-level chaining as advanced usage, and includes a full three-library registration recipe.
# Are there any user-facing changes?
New public APIs: `SessionContext.with_extensions`, `SessionExtensionComponents`, and the `SessionExtensionExportable` / `__datafusion_session_extension__` protocol. The context-outlives-DataFrame ownership contract is now documented. No breaking changes to existing APIs.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add method on SessionContext to add all extensions from one library [datafusion-python]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer merged PR #1679:
URL: https://github.com/apache/datafusion-python/pull/1679
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add method on SessionContext to add all extensions from one library [datafusion-python]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on PR #1679:
URL: https://github.com/apache/datafusion-python/pull/1679#issuecomment-5604129834
Ok, I did a lot of cleanup and I think this solution is *far* better than the original version. The big line number diff in the PR is because I revamped the documentation site to have 3 different audiences: datafusion-python contributors, users of the package, and people writing extensions. That split generated a lot of new documentation, which I think is a good 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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add atomic SessionContext.with_extensions API [datafusion-python]
Posted by "ntjohnson1 (via GitHub)" <gi...@apache.org>.
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add atomic SessionContext.with_extensions API [datafusion-python]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #1679:
URL: https://github.com/apache/datafusion-python/pull/1679#issuecomment-5375229711
Following up on comment https://github.com/apache/datafusion-ballista/pull/2252#pullrequestreview-4981918462 and follow up on this PR. `FFILogicalCodec::encode|decode_file_format` break our intent to have fully working distributed execution.
I might be wrong but `encode|decode_file_format` might not be an easy fix, at least not in short term and not in datafusion 55 timeframe, hence i have a proposal to make.
`DistributedExec` in ballista is nothing but a GRPC wrapper, it takes a logical plan, calls grpc endpoint and returns a stream of record batches. Would it make sense to create a `CallbackPlanner` (we might need a `CallbackExec`) in datafusion-py which would take a python closure `LogicalPlan -> Stream<RecordBatches>` (or `LogicalPlanBlob -> Stream<RecordBatches>`). Ballista would provide a closure which implements grcp logic in python code.
As `CallbackPlanner` have same library marker as df python `FFILogicalCodec` will not be triggered and we should have possibility to fully integrate distributed execution.
Basically we could implement query planner in python (limited but working)
wdyt @timsaucer and @ntjohnson1 ?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add atomic SessionContext.with_extensions API [datafusion-python]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on PR #1679:
URL: https://github.com/apache/datafusion-python/pull/1679#issuecomment-5544543804
Thanks for chasing this down — the `write_*` failures are real, but I think the fix is much smaller than a new planner API, and it is not in this PR.
**Where it actually breaks.** `FFI_LogicalExtensionCodec` has no vtable entries for file formats at all. Both hooks are hardcoded stubs in `datafusion-ffi` (`src/proto/logical_extension_codec.rs`, `try_decode_file_format` / `try_encode_file_format` → `not_impl_err!("FFI does not support ...")`). So *any* codec that crosses the FFI boundary loses file-format support, and `df.write_csv/write_parquet/write_json` build a `LogicalPlan::Copy` that carries a `FileFormatFactory`. That is the whole failure.
**This PR already survives it on the datafusion-python side.** `PythonLogicalCodec::try_{encode,decode}_file_format` go through `chain_encode`/`chain_decode`, which collect a chained codec's error and fall through to the terminal `DefaultLogicalExtensionCodec`. Built-in formats encode and decode fine there. The direction that fails is the mirror image: Ballista holds an `FFI_LogicalExtensionCodec` wrapping the Python session's codec and calls it directly, with no Default fallback behind it.
**Two fixes, both much cheaper than a new API:**
1. *Ballista side, today, no upstream change.* Compose the FFI codec you get from Python with a `DefaultLogicalExtensionCodec` fallback on the two file-format hooks — the same shape `PythonLogicalCodec` uses here. One file, unblocks DF 55.
2. *Upstream `datafusion-ffi`, ~8 lines.* Make those two hooks delegate to a local `DefaultLogicalExtensionCodec` instead of returning `not_impl_err!`. No new vtable entries, so no ABI break, so it can ship in a 55.x patch. It works because built-in file formats are fully self-describing protobuf and nothing has to cross the boundary: `DefaultLogicalExtensionCodec::try_encode_file_format` in `datafusion-proto` covers csv/json/arrow/avro/parquet by downcasting the factory locally. What is left unsupported is a *custom* `FileFormatFactory` owned by the foreign library, and I agree that one needs a real `FFI_FileFormatFactory` and will not make the 55 window — but it is also not what is breaking `write_*`.
**On `CallbackPlanner`.** I do not think the concept holds up, and I want to be concrete about why rather than just deferring it.
A `QueryPlanner` returning a `CallbackExec` leaf replaces the *entire* plan with one opaque node: no children, no partitioning DataFusion can see, no pushdown, no repartition, no limit, no statistics. That is not planning, it is interception at the root. It is coherent for exactly one architecture — thin client, full delegation — which is what Ballista's `DistributedExec` already is; the local plan is degenerate because none of the work is local. But it has no answer for a hybrid plan, a local table joined against a remote one, because there is nothing left to split on. All or nothing.
It also buys very little over what already ships:
```python
plan_bytes = df.logical_plan().to_bytes(ctx)
batches = my_grpc_client.execute(plan_bytes) # your gRPC logic, in Python
result = ctx.from_arrow(batches) # any __arrow_c_stream__ object
```
All three exist today, no new API. The only thing `CallbackPlanner` adds on top is transparency — the user keeps writing `ctx.sql(...).collect()` and never sees the interception. That is real, but it is a thin return for a permanent public API, and a second planner mechanism sitting next to the FFI query planner means two ways to do one thing.
If the goal is *partial* delegation rather than whole-plan delegation, the design that answers it already exists: a table provider exported from Python (`__datafusion_table_provider__`, see `docs/source/user-guide/io/table_provider.md`). Remote data appears as a table, DataFusion plans around it, filter and projection pushdown work, and local and remote sources mix in one query. That composes; a root-level callback does not.
The part that makes me most hesitant is the stated rationale: *"As `CallbackPlanner` have same library marker as df python, `FFILogicalCodec` will not be triggered."* That shapes a public API around which cdylib the code is compiled into, in order to route around an eight-line gap in `datafusion-ffi`. Fixing the gap is the smaller and more durable change.
**One genuinely open question, independent of all of the above** — and I think it is the same "chaining of codecs" point you flagged. Once Ballista's codec is installed on the Python session it becomes a chain entry, and chain entries write a framed payload (`DFPYCHN` + codec id) rather than bare bytes. A plain Rust `BallistaCodec` in the scheduler will not strip that envelope. So how does the scheduler decode — does it link `PythonLogicalCodec`, or does Ballista's codec stay off the chain? I am happy to add whatever hook makes the first option workable; I would rather solve that than work around it.
Proposal: keep #1679 as is, since the file-format gap predates it and is orthogonal, and I will open the upstream issue for (2). If you still want the callback route after the above, let us give it its own issue so the design can be argued on its own terms.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org