Skip to content

Stop OneOfExec's schema from tracking the best candidate - #56

Merged
wyatt-herkamp merged 2 commits into
branch-54from
nullable-column-bug-fix
Aug 15, 2026
Merged

wyatt-herkamp merged 2 commits into
branch-54from
nullable-column-bug-fix

Conversation

@wyatt-herkamp

@wyatt-herkamp wyatt-herkamp commented Aug 14, 2026 •

Copy link
Copy Markdown

Rationale for this change

OneOfExec holds several logically equivalent candidate plans and selects the cheapest one
using a user-supplied cost function. Candidates are only required to agree on field names and
data types — schemas_equal_ignoring_nullability deliberately permits them to differ in
nullability, 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:

  1. properties() returned self.candidates[self.best].properties(), so the node's declared
    schema was whichever candidate happened to be cheapest at that moment.
  2. best is recomputed on every with_new_children (it goes through try_new), and optimizer
    rules call with_new_children whenever they rewrite a child.

So the declared schema could change in the middle of a single optimizer rule. EnforceSorting
is the clearest way to hit this: it adds a SortExec above whichever candidate branches don't
already satisfy the node's required_input_ordering, then rebuilds the node. Any cost function
that is sensitive to sorts — which is the whole reason to declare required_input_ordering here,
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 the
plan's schema (OptimizationInvariantChecker / is_allowed_field_change in
datafusion/core/src/physical_planner.rs). Nullability may be tightened but never loosened.
So when best moved from a non-nullable candidate to a nullable one, planning failed with:

PhysicalOptimizer rule 'EnforceSorting' failed. Schema mismatch.
Expected original schema: ... Field { "event_date": Date32 }
got new schema:           ... Field { "event_date": nullable Date32 }

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?

  • OneOfExec now owns its declared PlanProperties instead of borrowing the current winner's.
    Ordering, partitioning, emission type and boundedness still come from best; only the schema
    is decoupled.
  • New widened_schema: field i is nullable if any candidate's field i is 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_change compares it too.
  • New declared_properties builds those properties, using
    EquivalenceProperties::with_new_schema to rewrite equivalence classes and orderings onto the
    widened schema. That helper requires only matching field counts and data types, so a
    nullability-only difference is exactly what it accepts.
  • execute relabels batches onto the declared schema when the selected candidate's schema is
    tighter, so the stream stays consistent with what the node advertises. Columns are moved by
    Arc, so this is not a data copy. In practice PruneCandidates replaces this node before
    execution; this just removes the inconsistency.

Widening is the only safe direction, and it is what makes a stable schema possible:

  • while the OneOfExec is in the plan its schema is fixed, so no rule can observe a change no
    matter how often best is recomputed;
  • when PruneCandidates later collapses the node to a single candidate, the schema either stays
    identical or tightens — both permitted.

Pinning to candidates[0] instead (mirroring the logical OneOf::schema(), which always exposes
the 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 changes
    its mind across calls, asserting the schema survives a with_new_children that moves best.
    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
EnforceSorting rule over a mixed-nullability candidate set.

The full suite passes (73 tests), with cargo fmt and cargo clippy --all-targets clean.

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. A
consumer reading nullability off this node will see nullable: true where it previously saw a
non-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.

@wyatt-herkamp
wyatt-herkamp marked this pull request as ready for review August 14, 2026 16:50
candidates: &[Arc<dyn ExecutionPlan>],
best: usize,
) -> Result<PlanProperties> {
let properties = candidates[best].properties().as_ref().clone();

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)?;

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@zhuqi-lucas

Copy link
Copy Markdown

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 ReverseOrder trips identically, and so would any schema_check rule that rebuilds this node. Something like "Stop OneOfExec's schema from tracking the best candidate" would read better in git log. Worth keeping the error string in the body so anyone searching for it still lands here.

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.

@wyatt-herkamp wyatt-herkamp changed the title Fix Schema Mismatch on Enforce Sorting Stop OneOfExec's schema from tracking the best candidate Aug 15, 2026
@wyatt-herkamp
wyatt-herkamp merged commit d369494 into branch-54 Aug 15, 2026
8 checks passed
@wyatt-herkamp
wyatt-herkamp deleted the nullable-column-bug-fix branch August 15, 2026 11:33
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants