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

[ManvithPanyam]

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
(DataFrame → Series, 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]

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]

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

[tvalentyn]

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]

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

[ManvithPanyam]

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.

[tvalentyn]

thanks, please don't squash reviewed and unreviewed code in the same commits as it hides the diff between versions. would you be able to rebase so we have both commits? thanks!

[ManvithPanyam]

Done — split into two commits, rebased/pushed. First is the original as you reviewed it, second is just the reindex()-based fix and everything else from this thread.

[ManvithPanyam]

Verified — the reindex()-based approach itself is correct and I've adopted it, simplifying quite a bit (unified dummy_index construction, removed the duplicate per-branch dedup/copy logic).

One nuance: kept the is_unique/drop_duplicates() dedup check, just hoisted once at the top instead of duplicated per-branch. Traced it — pandas caches an IndexEngine on an index once something inspects it (e.g. .is_unique, .loc, .get_loc), and that cache can survive slicing to iloc[:0], leaving is_unique stale as False even on an empty result. This path is reachable in practice (e.g. the existing test harness calls .xs() on the full arg before slicing to an empty proxy), so removing the check entirely would reintroduce a reindex() failure in that case.

Also found and fixed a real bug while testing this: a non-empty user-supplied proxy with duplicate index labels (via to_dataframe(pcoll, proxy=df_with_dups)) crashed the old code with a length-mismatch error, since drop_duplicates() shrinks the index but not the frame before reassignment. Pre-slicing to iloc[:0] before the dedup check (as in your suggestion) fixes this too. Added a regression test for it.

453 passed, 19 skipped, zero regressions.

Note: 4 CI checks are failing on unrelated Bigtable IO tests — a fresh
google-cloud-bigtable 2.44.0 (published Sept 3, 22:30 UTC) broke
DirectRow._get_mutations() handling. Confirmed master's own scheduled
runs are failing identically before and after this PR's CI run, so
it's unrelated to this change.

[tvalentyn]

thanks!