fix(dataframe): handle pandas dimensionality reduction in .xs() for single-item matches - #39851
fix(dataframe): handle pandas dimensionality reduction in .xs() for single-item matches#39851ManvithPanyam wants to merge 1 commit into
Conversation
…ingle-item matches Beam's .xs() implementation assumed static output shape (DataFrame/Series) across partitions, but pandas reduces dimensionality (DataFrame->Series, Series->scalar) when a key matches exactly one row and all index levels are selected. This caused TypeError/shape-mismatch failures during cross-partition concat. Fixes the key_size >= nlevels path to route matching partitions through a singleton unwrap stage that mirrors pandas' actual runtime behavior, while documenting the inherent proxy-time ambiguity for duplicate-match cases (proxy assumes single-match dimensionality; runtime produces whichever type pandas actually returns). Fixes apache#28559 Signed-off-by: ManvithPanyam <250704031+ManvithPanyam@users.noreply.github.com>
|
Assigning reviewers: R: @jrmccluskey for label python. Note: If you would like to opt out of this review, comment Available commands:
The PR bot will only process comments in the main thread (not review comments). |
19b0c08 to
5c6b20a
Compare
|
@tvalentyn — opened a fix for this. Root cause: Beam's One thing flagged for review: the proxy can't distinguish single-match |
What does this PR do?
Fixes
.xs()onDeferredDataFrame/DeferredSerieswhen a key matchesexactly one row and all index levels are selected.
The bug: pandas reduces dimensionality on single-item
.xs()matches(
DataFrame→Series,Series→ scalar), but Beam's implementationassumed a static output shape across partitions. Non-matching partitions
return an empty container of the original type, so when the matching
partition returned a dimensionality-reduced result, cross-partition
pd.concateither raisedTypeError(scalar concat) or silently produceda corrupted schema (NaN columns from a shape mismatch).
The fix: when
key_size >= nlevels, matching partitions are routedthrough a wrapped singleton stage that mirrors pandas' actual runtime
output, then unwrapped to the real return type — instead of assuming the
proxy shape holds at execution time.
Known limitation (documented in code): the proxy schema (computed
at graph-construction time from a 0-row template) can't know whether a
key will match 1 row or several at runtime, so it always assumes the
dimensionality-reduced type. If a key has duplicate matches, pandas
returns the non-reduced container instead — this is fundamentally
undecidable at proxy time, same class of limitation as
sort_values(),describe(), and other data-dependent-shape operations already in thismodule. Tests exercising duplicate-match keys use
check_proxy=Falseaccordingly, with the reasoning documented inline.
Fixes
Fixes #28559
Tests
Added regression coverage in
frames_test.pyfor all reported failuremodes: single-level index single match, MultiIndex 0-levels-remaining
single match (both unique and duplicate-key datasets), and Series
single-item
.xs(). Fullframes_test.pysuite: 452 passed, 19 skipped,zero regressions.