Skip to content

fix(dataframe): handle pandas dimensionality reduction in .xs() for single-item matches - #39851

Open
ManvithPanyam wants to merge 1 commit into
apache:masterfrom
ManvithPanyam:fix-beam-28559-xs
Open

fix(dataframe): handle pandas dimensionality reduction in .xs() for single-item matches#39851
ManvithPanyam wants to merge 1 commit into
apache:masterfrom
ManvithPanyam:fix-beam-28559-xs

Conversation

@ManvithPanyam

Copy link
Copy Markdown
Contributor

What does this PR do?

Fixes .xs() on DeferredDataFrame/DeferredSeries when a key matches
exactly one row and all index levels are selected.

The bug: pandas reduces dimensionality on single-item .xs() matches
(DataFrameSeries, Series → scalar), but Beam's implementation
assumed 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.concat either raised TypeError (scalar concat) or silently produced
a corrupted schema (NaN columns from a shape mismatch).

The fix: when key_size >= nlevels, matching partitions are routed
through 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 this
module. Tests exercising duplicate-match keys use check_proxy=False
accordingly, with the reasoning documented inline.

Fixes

Fixes #28559

Tests

Added regression coverage in frames_test.py for all reported failure
modes: single-level index single match, MultiIndex 0-levels-remaining
single match (both unique and duplicate-key datasets), and Series
single-item .xs(). Full frames_test.py suite: 452 passed, 19 skipped,
zero regressions.

…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>
@github-actions

Copy link
Copy Markdown
Contributor

Assigning reviewers:

R: @jrmccluskey for label python.

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

  • stop reviewer notifications - opt out of the automated review tooling
  • remind me after tests pass - tag the comment author after tests pass
  • waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

@ManvithPanyam

Copy link
Copy Markdown
Contributor Author

@tvalentyn — opened a fix for this. Root cause: Beam's .xs() assumed
static output shape across partitions, but pandas collapses dimensionality
on single-row matches, breaking cross-partition concat.

One thing flagged for review: the proxy can't distinguish single-match
vs. duplicate-match keys at graph-construction time (no row data
available yet), so it always predicts the dimensionality-reduced type.
Documented this as a known limitation in the code — same class of
tradeoff as other data-dependent-shape ops in this module. Open to
feedback if there's a cleaner way to handle it.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug]: Some methods on DeferredSeries and DeferredDataFrame don't work right when returning single items

1 participant