Repository navigation
Stop OneOfExec's schema from tracking the best candidate - #56
Conversation
| candidates: &[Arc<dyn ExecutionPlan>], | ||
| best: usize, | ||
| ) -> Result<PlanProperties> { | ||
| let properties = candidates[best].properties().as_ref().clone(); |
There was a problem hiding this comment.
Small perf note. Both this deep clone and widened_schema below run unconditionally, so every OneOfExec pays them on every with_new_children, even when no candidate widens anything, which is every table except this one family. The early return sits below both.
The clone is the bigger half since PlanProperties carries EquivalenceProperties, and widened_schema also allocates a fields Vec, a metadata HashMap clone and a Schema. Both scale with column count.
properties() already hands back an &Arc<PlanProperties>, so a no-allocation scan first would make the common path a plain Arc::clone:
let properties = if needs_widening(&candidates) {
Arc::new(declared_properties(&candidates, best)?)
} else {
Arc::clone(candidates[best].properties())
};where needs_widening just checks whether any non-nullable field in candidate 0 is nullable in another candidate. Not much on 19 columns, but this sits in the optimizer path on every rebuild.
| context: Arc<TaskContext>, | ||
| ) -> Result<SendableRecordBatchStream> { | ||
| self.candidates[self.best].execute(partition, context) | ||
| let stream = self.candidates[self.best].execute(partition, context)?; |
There was a problem hiding this comment.
Both new tests assert on schema() only, so nothing covers the relabeling below. If someone later simplifies this branch away, the tests still pass and the breakage surfaces as a declared-vs-actual schema mismatch on the wire, which is unpleasant to trace.
Worth executing in a test: build the node so best is the NOT NULL candidate, collect a batch, assert batch.schema() equals the widened declared schema. divergent_candidates() already sets that up.
|
Two more, both non-blocking. The title describes the symptom and scopes it narrower than the fix actually is. EnforceSorting is just the rule that surfaced it first; per your own doc comment This also looks worth sending to datafusion-contrib. Nothing in it is specific to our setup: partition columns are non-nullable in DataFusion's own ListingTable model too, so any deployment that ends up with a mixed candidate set hits the same flip. #37 was the first half of the same problem, making mixed candidate sets legal at planning time without making the physical schema stable, so the two together are the complete fix and might be easier to land upstream as a pair. |
Rationale for this change
OneOfExecholds several logically equivalent candidate plans and selects the cheapest oneusing a user-supplied cost function. Candidates are only required to agree on field names and
data types —
schemas_equal_ignoring_nullabilitydeliberately permits them to differ innullability, because different table providers legitimately expose an equivalent column with
different nullability. A typical case: a column that is a hive partition column on a base table
(non-nullable) is an ordinary file column on a materialized view defined over that table
(nullable).
Two things then combine badly:
properties()returnedself.candidates[self.best].properties(), so the node's declaredschema was whichever candidate happened to be cheapest at that moment.
bestis recomputed on everywith_new_children(it goes throughtry_new), and optimizerrules call
with_new_childrenwhenever they rewrite a child.So the declared schema could change in the middle of a single optimizer rule.
EnforceSortingis the clearest way to hit this: it adds a
SortExecabove whichever candidate branches don'talready satisfy the node's
required_input_ordering, then rebuilds the node. Any cost functionthat is sensitive to sorts — which is the whole reason to declare
required_input_orderinghere,per the comment on that field — can then rank a different candidate first, and the node's schema
follows.
DataFusion checks that a physical optimizer rule declaring
schema_check()does not change theplan's schema (
OptimizationInvariantChecker/is_allowed_field_changeindatafusion/core/src/physical_planner.rs). Nullability may be tightened but never loosened.So when
bestmoved from a non-nullable candidate to a nullable one, planning failed with:This is intermittent, and it blames whichever rule happened to trigger the re-pick rather than
the node responsible. It only surfaces when two candidates are close enough in cost that
inserting a sort on one branch flips the ranking, so it comes and goes with statistics and data
volume.
What changes are included in this PR?
OneOfExecnow owns its declaredPlanPropertiesinstead of borrowing the current winner's.Ordering, partitioning, emission type and boundedness still come from
best; only the schemais decoupled.
widened_schema: fieldiis nullable if any candidate's fieldiis nullable. Names,data types, field metadata and schema metadata come from the first candidate, which the
existing candidate validation has already established the others match. Field metadata is
preserved deliberately —
is_allowed_field_changecompares it too.declared_propertiesbuilds those properties, usingEquivalenceProperties::with_new_schemato rewrite equivalence classes and orderings onto thewidened schema. That helper requires only matching field counts and data types, so a
nullability-only difference is exactly what it accepts.
executerelabels batches onto the declared schema when the selected candidate's schema istighter, so the stream stays consistent with what the node advertises. Columns are moved by
Arc, so this is not a data copy. In practicePruneCandidatesreplaces this node beforeexecution; this just removes the inconsistency.
Widening is the only safe direction, and it is what makes a stable schema possible:
OneOfExecis in the plan its schema is fixed, so no rule can observe a change nomatter how often
bestis recomputed;PruneCandidateslater collapses the node to a single candidate, the schema either staysidentical or tightens — both permitted.
Pinning to
candidates[0]instead (mirroring the logicalOneOf::schema(), which always exposesthe base branch) would be equally stable but wrong: collapsing to a nullable candidate would then
be a loosening, moving the same failure into
PruneCandidates.Are these changes tested?
Yes, two new unit tests in
tests_nullability:one_of_exec_declares_widened_nullability— with one non-nullable and one nullable candidate,the node advertises the nullable field even while the non-nullable candidate is selected, and
nothing else about the schema is disturbed.
one_of_exec_schema_is_stable_when_best_changes— a cost function that deliberately changesits mind across calls, asserting the schema survives a
with_new_childrenthat movesbest.It also asserts the winner really did change, so it cannot pass vacuously.
Verified against the original reproduction as well: the failure reproduces on the parent commit
and is resolved on this branch, using a separate end-to-end test that drives the real
EnforceSortingrule over a mixed-nullability candidate set.The full suite passes (73 tests), with
cargo fmtandcargo clippy --all-targetsclean.Are there any user-facing changes?
No public API changes.
One behavioral change: for a candidate set whose members disagree on nullability,
OneOfExec::schema()now reports the widest nullability rather than the selected candidate's. Aconsumer reading nullability off this node will see
nullable: truewhere it previously saw anon-nullable field, depending on which candidate was selected. This is the direction consumers
must already tolerate, since the selected candidate could change between planning runs.
Candidate sets whose schemas already agree on nullability are unaffected — the widened schema is
identical to the existing one, and the properties are reused as-is.