Report physical partitioning and resolve two panics escaping as PanicException - #1720
Open
timsaucer wants to merge 14 commits into
Open
Report physical partitioning and resolve two panics escaping as PanicException#1720timsaucer wants to merge 14 commits into
timsaucer wants to merge 14 commits into
Conversation
This was referenced Sep 9, 2026
timsaucer
commented
Sep 10, 2026
Comment on lines
-99
to
-108
| If a library ships codecs *and* a planner, prefer | ||
| `SessionContext.with_extensions(bundle)` over installing each piece by hand. It | ||
| installs every codec before it binds any planner, so a planner cannot end up | ||
| carrying a chain that a later `with_logical_extension_codec` call has grown. | ||
| The library exposes a bundle object implementing | ||
| `__datafusion_session_components__` for its codecs and | ||
| `__datafusion_session_planner__` for its planner — the latter is handed the | ||
| planner installed so far, so several libraries that each ship one nest instead | ||
| of displacing each other. See {ref}`extension_bundles`. | ||
|
|
Member
Author
There was a problem hiding this comment.
This is a drive by removal. There is no "upgrading" to this since this is a new feature.
timsaucer
commented
Sep 10, 2026
Comment on lines
-190
to
-191
| Tables created in memory from record batches are currently not | ||
| supported. |
Member
Author
There was a problem hiding this comment.
This was true for the logical side, but not the physical side so removed.
timsaucer
marked this pull request as ready for review
September 10, 2026 16:23
timsaucer
changed the title
Groundwork for a multi-library distributed-execution example. Each item here is something that example needs and cannot get today. `ExecutionPlan.output_partitioning` is new. `partition_count` already existed but discards everything except the count, so a driver deciding how to split work across workers could not tell hash-distributed output from merely counted output, nor read the hash keys. It returns a `PhysicalPartitioning`, named to keep it distinct from `datafusion.expr.Partitioning` — that one is the logical partitioning `repartition_by_hash` takes as a request, this one is what a built plan does. Physical expressions have no Python representation, so the hash keys are returned in their displayed form. `SessionContext.execute` now bounds-checks the partition index. The plan's leaves index their partition vector directly, so an out-of-range index reached `MemorySourceConfig` and panicked; the panic was caught as a tokio `JoinError` and arrived as `index out of bounds: the len is 2 but the index is 5`, naming neither the plan nor the index the caller passed. `SessionConfig.set` no longer routes through `SessionConfig::set_str`, which unwraps. An unknown namespace — `datafusion.runtime.*`, or a config extension not yet installed — aborted with a `PanicException`, which derives from `BaseException` and so escapes `except Exception`. `information_schema. df_settings` lists keys in both categories, so replaying settings onto a worker hit this first. Two docstrings on `ExecutionPlan` claimed that a table registered from record batches cannot be serialized. That is true of `LogicalPlan`, whose `try_encode_table_provider` has no arm for one, and false of the physical layer, which inlines the batches: verified by decoding on a context sharing nothing with the encoder and executing. A test pins it, since it is what lets a worker run a plan the driver encoded. Also documents, rather than fixes, the `ForeignExecutionPlan` arm in the example provider's physical codec. It claims every other library's nodes, which the extension guide tells authors not to do — but it is load-bearing: `EnsureCooperative` runs during a foreign planner's `create_physical_plan` and hands the library back a `ForeignExecutionPlan` wrapping the host's `CooperativeExec`, which has no reachable `try_to_proto`. Narrowing the arm makes 31 of the 51 tests in the query-planner example fail, all on that node. The comment now says so, and says a planner that controls its own physical optimizer rules needs no such arm. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`SessionContext.execute` took its second argument as `partitions`, which reads as a count when it is a single partition index. Rename it to `partition` and give the method a real docstring with a doctest covering both a full sweep over `partition_count` and the out-of-range `ValueError`. Every call site in the repo, docs, and examples passes it positionally, so add a short note to the upgrade guide for anyone passing it by keyword. Drop two changelog-shaped sections from the upgrade guide. `output_partitioning` is additive and `SessionContext.execute` / `SessionConfig.set` only trade a panic for a raise, so neither asks the reader to change anything. The `with_extensions` recommendation goes for the same reason; it is advice, and the extension guide already carries it under `extension_bundles`. Also remove the comment above the partition range check. It explained the check by way of a `MemorySourceConfig` panic, which reads as a complaint about DataFusion's leaves. The index arrives straight from Python and no planner has seen it, so validating it needs no more justification than any other argument check at the boundary. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The `ForeignExecutionPlan` arm's comment explained the symptom but left the reader no way to find out whether the workaround is still needed. Name the upstream umbrella issue, apache/datafusion#25152, and the cascade behind it: `FFI_PlanProperties` carries no `scheduling_type`, so `EnsureCooperative` reads every foreign leaf as non-cooperative and wraps it, and the resulting `ForeignExecutionPlan` then cannot serialize itself. Fixing either half retires the arm. Drop the out-of-range call from `SessionContext.execute`'s doctest. A docstring example shows a reader how to use the method, and this one put a wrong call in front of them; the message it asserted is already pinned by `test_execute_rejects_an_out_of_range_partition`. The `Raises:` section is the right home for that behaviour, so complete it: a negative or oversized index raises `OverflowError` from the `usize` conversion, not the `ValueError` the bounds check produces. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The docstring was wrong in three ways at once, and the panic fix earlier in this branch is what made it worth opening: it promised "a new SessionConfig object" when the method mutates in place and returns self, it mis-indented the `Args` entries so Sphinx rendered them as body text rather than a field list, and it had no `Raises` at all -- so the one behaviour this branch changed, an unknown key raising instead of aborting the interpreter, was documented only in the Rust source that no user reads. Give it a truthful `Returns`, a `Raises` covering both an unknown key and an unparsable value, and a doctest that reads the option back out through `information_schema.df_settings`. The `datafusion.runtime.*` trap goes to `configuration.md`, which is where a reader is when they need it: those keys appear in `df_settings` but are not settable, so replaying that table verbatim onto a worker fails on the first such row. The docstring states it in one sentence and points at the guide. That page also now records that the `with_*` methods modify in place, which is what makes the chained style it already demonstrates work. Note the same false `Returns` line appears on every other `with_*` method on this class; correcting those is a separate sweep, not this branch's business. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Trading a `PanicException` for an untyped `Exception` only half-fixed the problem: a caller replaying `information_schema.df_settings` still could not catch the failure without swallowing every other error this crate raises, and was left matching on message text. A rejected key or an unparsable value is an argument error, so it should raise `ValueError` -- the same thing an out-of-range partition index gets from `execute` two methods away. Map through `from_datafusion_error`, which already produces `PyValueError` and which `SessionContext.sql` already uses in this file, rather than propagating a `PyDataFusionError` and taking its blanket conversion. That conversion is deliberately untouched: reclassifying every error out of this crate is a much wider change with its own compatibility story. The test can now assert something. It previously ended in `assert isinstance(excinfo.value, Exception)` under a `pytest.raises(Exception, ...)` that had already proven exactly that, so the line could never fail. Assert the type instead, and add a case for a known key with a value of the wrong type, which reaches the same path by a route a settings-replay loop is just as likely to take. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The docstring referred the reader to `distributed_query_engines`, which is
wrong twice over: that page's premise is that you do *not* partition by hand
because the engine does it for you, and it documents work that is not yet
usable from datafusion-python. A reader following the link to find out what
a scheme means landed on a status page for a different road.
Nothing under docs/source/ discussed plan partitioning from the caller's
side, so give the claim a home on the page that already tells readers to
call `repartition` and `repartition_by_hash` to keep their cores busy --
and that, until now, gave them no way to check whether it worked.
The new section is worth more than a pointer. `repartition_by_hash(col("a"),
num=8)` followed by an aggregation reports `Hash([a@0], 16)`: the optimizer
discarded the requested repartition and inserted its own at
`target_partitions`, so neither the scheme nor the count is what was asked
for. Drop the aggregation and the repartition vanishes entirely, leaving
`UnknownPartitioning` over the source's partition count. Both were verified
against a Parquet source, and both are invisible without this accessor,
which is the concrete form of the have-versus-ask distinction the class
docstring asserts.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`From<PyPhysicalPartitioning> for Partitioning` had no callers. The type is a read-only report of what a built plan does, nothing accepts a partitioning as an argument, and a Rust caller wanting the `Partitioning` reads it off the plan, so there is no inbound direction to support. Removing the impl alone left a deprecation warning: pyo3 auto-derives `FromPyObject` for a `#[pyclass]` that implements `Clone` and now wants the choice made explicitly. Omission is not the way to decline it, so say `skip_from_py_object`, which is what the example crate's providers already use. Clippy is clean again, and `--all-targets` keeps it that way. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The `scheme` docstring listed four values as though a plan could report any of them. Two of the four needed opposite corrections. `RoundRobinBatch` is reachable, and now tested. It is easy to miss because it never survives at the root: the optimizer inserts one only above a source with fewer partitions than `target_partitions` and CPU work above it to parallelize, so a single-file Parquet scan under a filter and a grouped aggregate produces `RepartitionExec: partitioning=RoundRobinBatch(8)` three levels down. The test walks `children` and asserts all three schemes in that tree, which also pins that the accessor reads each node's own partitioning rather than the root's. `Range` is the other way: it is in the upstream enum, but nothing constructs one in a physical plan. `RangePartitioning` says optimizer and execution support is deliberately unimplemented, per apache/datafusion#22395. Say so, and say why the match arm exists anyway -- it keeps this getter compiling when that support lands, and dropping it would make the match non-exhaustive. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`scheme` returns one of exactly four strings, so annotate it `Literal` rather than `str`. The promise is safe to make: the Rust getter matches exhaustively over `Partitioning`, so a new upstream variant is a compile error here before it can be a lie in the type. Bind `output_partitioning` once in the scan half of its test. It was read three times, and each read clones the `Partitioning` and builds a fresh wrapper. Validate the partition index in `execute` before building the `TaskContext`, not after. Nothing observable changes; the rejected call just stops doing setup work it is about to throw away. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`SessionConfig.set` was fixed to raise instead of panicking, but the
constructor still routed its dictionary through `SessionConfig::set`, which
forwards to `set_str` and unwraps. So
`SessionConfig({"datafusion.runtime.memory_limit": "unlimited"})` aborted as a
`PanicException` -- a `BaseException`, past any `except Exception` -- while the
same key passed to `set` raised `ValueError`.
The constructor is the likelier of the two to meet a bad key: it takes a
`dict[str, str]`, which is the shape a replayed set of settings arrives in, and
`information_schema.df_settings` lists eight `datafusion.runtime.*` rows that
cannot be set from a session config at all.
Route each entry through `options_mut().set` and map with
`from_datafusion_error`, so both routes reject the same keys the same way. The
`ScalarValue::Utf8` wrapper went with it: the parameter is already
`HashMap<String, String>`, and upstream only called `to_string()` back on it.
Which entry a dictionary with several bad keys reports is unspecified, since
`HashMap` iteration order is arbitrary. Documented rather than sorted -- a
caller fixes the reported key and runs again.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Two docstrings claimed no plan reports `Partitioning::Range` because upstream support was unimplemented. Neither half holds for DataFusion 55: `repartition/mod.rs` routes range partitioning through `RangeExpr`, and `physical_planner.rs` has a test asserting a planned `RepartitionExec` reports it. The comment's reasoning was also inverted -- the match arm compiles today because the variant exists today, not in anticipation of it landing. What is true is narrower: nothing in this package's own API asks for one. `repartition` requests round-robin, `repartition_by_hash` requests hash, and SQL has no range-repartition syntax. But a plan need not have been built here. `datafusion-proto` encodes and decodes physical range partitioning and `datafusion-ffi` carries it in both directions, so `ExecutionPlan.from_bytes` can return a plan reporting it, as can an extension library's query planner. State that where each reader is: the reachability in the guide beside the rest of the scheme discussion, one sentence and a `:ref:` in the docstring. Also note that `hash_expressions` is `None` for `Range`, which leaves the ordering and split points reachable only through `repr`. No test: reaching `Range` from Python needs a plan built in Rust, and a Rust test here would never run in CI. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Both docstrings introducing `PhysicalPartitioning` distinguished it from `datafusion.expr.Partitioning` by calling that one "the partitioning `repartition_by_hash` asks for". No method takes one: `repartition` takes a count and `repartition_by_hash` takes expressions and a count. The type cannot be constructed from Python at all -- `Partitioning()` raises `TypeError` -- has no public members, and is only ever handed back by `Repartition.partitioning_scheme()`. Describe it that way instead: the logical partitioning a `Repartition` node records. The request-versus-result contrast the docstrings were reaching for is real and worth drawing, so keep it, but attach it to the node that holds the request rather than to a parameter that does not exist. The Rust comment also notes that the two enums differ, since the logical one has `DistributeBy` and no `UnknownPartitioning`. Pin the contrast with a test rather than only asserting it in prose. A `repartition_by_hash(num=8)` with nothing above it to consume the redistribution is dropped by the optimizer, so the request records Hash into 8 while the built plan reports `UnknownPartitioning(2)` -- disagreeing on both scheme and count, which is the reason the two types stay separate. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Drop the guide sentence claiming every config method mutates in place and returns itself. It is true, but every `with_*` docstring still promises "a new `SessionConfig` object", and a guide that contradicts the API docs it links to is worse than a guide that stays quiet. Correcting fifteen docstrings is its own change. `PhysicalPartitioning` gains `__eq__` and `__hash__`, computed structurally from scheme, partition count and hash expressions. Deliberately not delegated to `Partitioning`'s `PartialEq`, whose match lists no `UnknownPartitioning` arm and so falls through to `false`: two identical `UnknownPartitioning(2)` values are unequal there. That is defensible for deciding whether a partitioning satisfies a distribution requirement, but a non-reflexive `__eq__` would be a trap in Python, and `UnknownPartitioning` is what an ordinary file scan reports. `__hash__` comes along so the class stays usable in a set. `test_output_partitioning_reports_round_robin` asserted set equality over every scheme in the tree, pinning optimizer output that is not the property under test. Membership instead, and 50 rows rather than 2000 -- the round robin appears either way, so the larger file bought nothing. The PyO3 parameter is now `partition` to match the wrapper, so `ctx.execute(plan, partition=0)` works on the internal binding too, and the test covers the keyword form alongside the positional one. `from_bytes` had its false memory-table sentence removed but got nothing back, so it no longer said that a decoding session need share nothing with the encoder -- the property that makes it useful. Restored from the reader's side. Also cover the `OverflowError` a negative index raises, which was documented on `execute` but never exercised. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The two `SessionConfig.set` tests went into `test_plans.py` because they were committed alongside the `execute` bounds check, not because they have anything to do with plans. `test_context.py` is where `SessionConfig` construction is already covered, and it now also holds the three constructor tests for the other half of the same panic defect. Move them there, ahead of the constructor cases so the method they refer back to is read first, and fold the duplicated note about `PanicException` deriving from `BaseException` into the first of the five. `test_plans.py` keeps its `SessionConfig` import for the `with_target_partitions` calls in the partitioning tests. No assertion changed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
timsaucer
force-pushed
the
feat/plan-partitioning-and-errors
branch
from
September 10, 2026 16:39
39304d7 to
7123e7d
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Part of #1719. Closes nothing on its own.
Rationale for this change
While writing the the multi-library distributed example in #1719 and #1721, we discovered a few gaps in the existing code base that are beneficial to resolve. This PR addresses the bugs identified.
What changes are included in this PR?
SessionContext.executebounds-checks the partition indexSessionContext.executerenamed its second parameterSessionConfigno longer aborts on an unknown keyExecutionPlan.output_partitioningflushed outAre there any user-facing changes?
Yes. One breaking rename, the rest additive or strictly-better failures.
SessionContext.execute's second parameter is renamedpartitions→partition. Breaking for keyword callers, so this PR carries theapi changelabel anddocs/source/user-guide/upgrade-guides.mdhas a section with the before and after. It is the only entry that asks the reader to change anything, which is why it is the only one in the guide.