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.

@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.

isinstance(k_val, tuple) else pd.Index([k_val],
name=proxy_frame.index.name))
dummy_data = {
col: [proxy_frame[col].dtype.type()]

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

dtype.type() will crash on some Pandas types like Categorical Type.

will smth like this work?

        proxy_frame = reindexed._expr.proxy()
        
        dummy_index = (
            pd.MultiIndex.from_tuples([k_val], names=proxy_frame.index.names) if
            isinstance(k_val, tuple) else pd.Index([k_val],
                                                   name=proxy_frame.index.name))
        
        dummy_obj = proxy_frame.reindex(dummy_index)
        
        xs_proxy = dummy_obj.xs(k_val, **kwargs)
        if isinstance(xs_proxy, (pd.DataFrame, pd.Series)):
            xs_proxy = xs_proxy.iloc[:0] 

@tvalentyn

Copy link
Copy Markdown
Contributor

thanks for the contribution. are you a Beam Dataframes user?

…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>
@ManvithPanyam

Copy link
Copy Markdown
Contributor Author

Not really an active Beam DataFrames user day-to-day — found this while
working through the dataframe module after my last PR (#39581, the
empty-CSV restriction tracker fix). This was the next issue that stood
out while digging deeper into the module.

Thanks for the catch on dtype.type() — confirmed it breaks on
Categorical (TypeError: type.__new__() takes exactly 3 arguments) and
tz-aware datetime (TypeError: function missing required argument 'year').

Applied your reindex() approach for the DataFrame branch — works
cleanly and is simpler than the manual per-column construction. For the
Series branch (single-match → scalar case), reindex().iloc[0] alone
silently changes dtype on plain numeric types (e.g. int64 → float64,
since reindex introduces NaN), so I kept dtype.type() as the primary
path there and only fall back to reindex().iloc[0] on the TypeError
extension-dtype case.

Also had to handle proxy indexes with duplicate labels — reindex()
raises on non-unique index, so I drop duplicates on the proxy copy
before reindexing (doesn't affect real execution, this is proxy-only).

Added regression tests for Categorical, tz-aware datetime, and nullable
Int64 columns. Full suite: 452 passed, 19 skipped, zero regressions.

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

2 participants