You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@beam.apache.org by "ManvithPanyam (via GitHub)" <gi...@apache.org> on 2026/08/22 15:40:32 UTC
[PR] fix(dataframe): handle pandas dimensionality reduction in .xs() for single-item matches [beam]
ManvithPanyam opened a new pull request, #39851:
URL: https://github.com/apache/beam/pull/39851
### 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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] fix(dataframe): handle pandas dimensionality reduction in .xs() for single-item matches [beam]
Posted by "tvalentyn (via GitHub)" <gi...@apache.org>.
tvalentyn commented on code in PR #39851:
URL: https://github.com/apache/beam/pull/39851#discussion_r3876287164
##########
sdks/python/apache_beam/dataframe/frames.py:
##########
@@ -1091,27 +1091,93 @@ def xs(self, key, axis, level, **kwargs):
reindexed = self.reorder_levels(
level + [i for i in range(self.index.nlevels) if i not in level])
- def xs_partitioned(frame, key):
- if not len(key):
- # key is not in this partition, return empty dataframe
- result = frame.iloc[:0]
- if key_size < frame.index.nlevels:
+ if key_size < reindexed.index.nlevels:
+
+ def xs_partitioned(frame, key):
+ if not len(key):
+ # key is not in this partition, return empty dataframe/series
+ result = frame.iloc[:0]
return result.droplevel(list(range(key_size)))
+ return frame.xs(key.item(), **kwargs)
+
+ return frame_base.DeferredFrame.wrap(
+ expressions.ComputedExpression(
+ 'xs',
+ xs_partitioned, [reindexed._expr, key_expr],
+ requires_partition_by=partitionings.Index(list(range(key_size))),
+ preserves_partition_by=partitionings.Singleton()))
+ else:
+ # When all index levels are matched (key_size >= nlevels), pandas .xs()
+ # return type is data-dependent:
+ # - Single match: reduces dimensionality (DataFrame -> Series, Series -> scalar)
+ # - Duplicate matches: preserves container type (DataFrame -> DataFrame, Series -> Series)
+ # Because proxy schemas are 0-row templates evaluated at graph construction time
+ # without knowledge of dataset contents or key frequencies, the proxy always assumes
+ # a single match (dimensionality-reduced type). At runtime, the Singleton unwrap stage
+ # correctly produces whichever type pandas returns. Tests with multi-matching keys
+ # therefore specify check_proxy=False.
+ def xs_partitioned_wrapped(frame, key):
+ if not len(key):
+ return pd.Series([], dtype=object)
+ k = key.item()
+ try:
+ res = frame.xs(k, **kwargs)
+ return pd.Series([res], dtype=object)
+ except KeyError:
+ return pd.Series([], dtype=object)
+
+ intermediate = expressions.ComputedExpression(
+ 'xs_partitioned_wrapped',
+ xs_partitioned_wrapped, [reindexed._expr, key_expr],
+ proxy=pd.Series([], dtype=object),
+ requires_partition_by=partitionings.Index(list(range(key_size))),
+ preserves_partition_by=partitionings.Singleton())
+
+ proxy_frame = reindexed._expr.proxy()
+ k_val = key_series.iloc[0]
+ if isinstance(proxy_frame, pd.DataFrame):
+ 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_data = {
+ col: [proxy_frame[col].dtype.type()]
Review Comment:
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]
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] fix(dataframe): handle pandas dimensionality reduction in .xs() for single-item matches [beam]
Posted by "tvalentyn (via GitHub)" <gi...@apache.org>.
tvalentyn commented on code in PR #39851:
URL: https://github.com/apache/beam/pull/39851#discussion_r3917446617
##########
sdks/python/apache_beam/dataframe/frames.py:
##########
@@ -1091,27 +1091,88 @@ def xs(self, key, axis, level, **kwargs):
reindexed = self.reorder_levels(
level + [i for i in range(self.index.nlevels) if i not in level])
- def xs_partitioned(frame, key):
- if not len(key):
- # key is not in this partition, return empty dataframe
- result = frame.iloc[:0]
- if key_size < frame.index.nlevels:
+ if key_size < reindexed.index.nlevels:
+
+ def xs_partitioned(frame, key):
+ if not len(key):
+ # key is not in this partition, return empty dataframe/series
+ result = frame.iloc[:0]
return result.droplevel(list(range(key_size)))
- else:
- return result
+ return frame.xs(key.item(), **kwargs)
- # key should be in this partition, call xs. Will raise KeyError if not
- # present.
- return frame.xs(key.item())
+ return frame_base.DeferredFrame.wrap(
+ expressions.ComputedExpression(
+ 'xs',
+ xs_partitioned, [reindexed._expr, key_expr],
+ requires_partition_by=partitionings.Index(list(range(key_size))),
+ preserves_partition_by=partitionings.Singleton()))
+ else:
+ # When all index levels are matched (key_size >= nlevels), pandas .xs()
+ # return type is data-dependent:
+ # - Single match: reduces dimensionality (DataFrame -> Series, Series -> scalar)
+ # - Duplicate matches: preserves container type (DataFrame -> DataFrame, Series -> Series)
+ # Because proxy schemas are 0-row templates evaluated at graph construction time
+ # without knowledge of dataset contents or key frequencies, the proxy always assumes
+ # a single match (dimensionality-reduced type). At runtime, the Singleton unwrap stage
+ # correctly produces whichever type pandas returns. Tests with multi-matching keys
+ # therefore specify check_proxy=False.
+ def xs_partitioned_wrapped(frame, key):
+ if not len(key):
+ return pd.Series([], dtype=object)
+ k = key.item()
+ try:
+ res = frame.xs(k, **kwargs)
+ return pd.Series([res], dtype=object)
+ except KeyError:
+ return pd.Series([], dtype=object)
+
+ intermediate = expressions.ComputedExpression(
+ 'xs_partitioned_wrapped',
+ xs_partitioned_wrapped, [reindexed._expr, key_expr],
+ proxy=pd.Series([], dtype=object),
+ requires_partition_by=partitionings.Index(list(range(key_size))),
+ preserves_partition_by=partitionings.Singleton())
- return frame_base.DeferredFrame.wrap(
- expressions.ComputedExpression(
- 'xs',
- xs_partitioned,
- [reindexed._expr, key_expr],
- requires_partition_by=partitionings.Index(list(range(key_size))),
- # Drops index levels, so partitioning is not preserved
- preserves_partition_by=partitionings.Singleton()))
+ proxy_frame = reindexed._expr.proxy()
Review Comment:
from AI review:
```
Calling proxy_frame.index.drop_duplicates() returns a deduplicated Index object which is shorter than the proxy_frame itself. Assigning it back to proxy_frame.index will raise a ValueError: Length mismatch from pandas.
You can simply slice proxy_frame.iloc[:0] unconditionally. An empty DataFrame/Series is guaranteed to have a uniquely-valued index, meaning .reindex() will always succeed in creating the required NaN dummy row, and it perfectly preserves the names and dtypes. This also eliminates the need to duplicate the logic for pd.DataFrame and pd.Series.
You can simplify the entire proxy generation block (lines 1136 to 1161 in your diff) to this:
proxy_frame = proxy_frame.iloc[:0]
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))
if isinstance(proxy_frame, pd.DataFrame):
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]
else:
try:
xs_proxy = proxy_frame.dtype.type()
except TypeError:
xs_proxy = proxy_frame.reindex(dummy_index).iloc[0]
```
What are your thoughts on this?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] fix(dataframe): handle pandas dimensionality reduction in .xs() for single-item matches [beam]
Posted by "tvalentyn (via GitHub)" <gi...@apache.org>.
tvalentyn merged PR #39851:
URL: https://github.com/apache/beam/pull/39851
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] fix(dataframe): handle pandas dimensionality reduction in .xs() for single-item matches [beam]
Posted by "ManvithPanyam (via GitHub)" <gi...@apache.org>.
ManvithPanyam commented on PR #39851:
URL: https://github.com/apache/beam/pull/39851#issuecomment-5452493914
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] fix(dataframe): handle pandas dimensionality reduction in .xs() for single-item matches [beam]
Posted by "tvalentyn (via GitHub)" <gi...@apache.org>.
tvalentyn commented on PR #39851:
URL: https://github.com/apache/beam/pull/39851#issuecomment-5544328256
thanks!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] fix(dataframe): handle pandas dimensionality reduction in .xs() for single-item matches [beam]
Posted by "ManvithPanyam (via GitHub)" <gi...@apache.org>.
ManvithPanyam commented on PR #39851:
URL: https://github.com/apache/beam/pull/39851#issuecomment-5540237036
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] fix(dataframe): handle pandas dimensionality reduction in .xs() for single-item matches [beam]
Posted by "tvalentyn (via GitHub)" <gi...@apache.org>.
tvalentyn commented on PR #39851:
URL: https://github.com/apache/beam/pull/39851#issuecomment-5485954907
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!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] fix(dataframe): handle pandas dimensionality reduction in .xs() for single-item matches [beam]
Posted by "ManvithPanyam (via GitHub)" <gi...@apache.org>.
ManvithPanyam commented on PR #39851:
URL: https://github.com/apache/beam/pull/39851#issuecomment-5381740288
@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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] fix(dataframe): handle pandas dimensionality reduction in .xs() for single-item matches [beam]
Posted by "ManvithPanyam (via GitHub)" <gi...@apache.org>.
ManvithPanyam commented on PR #39851:
URL: https://github.com/apache/beam/pull/39851#issuecomment-5496621151
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] fix(dataframe): handle pandas dimensionality reduction in .xs() for single-item matches [beam]
Posted by "tvalentyn (via GitHub)" <gi...@apache.org>.
tvalentyn commented on PR #39851:
URL: https://github.com/apache/beam/pull/39851#issuecomment-5445886720
thanks for the contribution. are you a Beam Dataframes user?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] fix(dataframe): handle pandas dimensionality reduction in .xs() for single-item matches [beam]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #39851:
URL: https://github.com/apache/beam/pull/39851#issuecomment-5381474645
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).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org