This is an automated email from the ASF dual-hosted git repository.
tvalentyn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 0be6fca73c6 fix(dataframe): handle pandas dimensionality reduction in
.xs() for single-item matches (#39851)
0be6fca73c6 is described below
commit 0be6fca73c6d4e6a9bb4f18f9bbf6046e90cc7a9
Author: Manvith Panyam <[email protected]>
AuthorDate: Fri Sep 4 23:13:17 2026 +0530
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
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 #28559
Signed-off-by: ManvithPanyam
<[email protected]>
* fix(dataframe): use reindex() for xs() proxy construction to support
extension dtypes
Addresses review feedback: dtype.type() crashes on Categorical and
timezone-aware datetime dtypes since they aren't callable as zero-arg
scalar constructors. Switched the DataFrame branch to build the dummy
proxy via reindex() + xs(), letting pandas handle type construction
internally. Kept dtype.type() as the primary path for the Series
(single-match scalar) branch since reindex().iloc[0] alone silently
upcasts plain numeric types (e.g. int64 -> float64) by introducing NaN;
falls back to reindex().iloc[0] only on TypeError for extension types.
Also handles proxy indexes with duplicate labels, which reindex()
otherwise rejects.
Added regression tests for Categorical, tz-aware datetime, and nullable
Int64 columns.
Signed-off-by: ManvithPanyam
<[email protected]>
* fix(dataframe): simplify xs() proxy construction and fix non-empty
duplicate-proxy crash
Further simplifies the reindex()-based proxy generation from the
previous commit: unifies dummy_index construction (was duplicated
per-branch), and hoists the is_unique/drop_duplicates() dedup check to
run once instead of twice.
The dedup check is retained, not removed — traced that pandas can cache
an IndexEngine on an index once inspected (e.g. via .is_unique, .loc,
.get_loc), and that cache can survive slicing to iloc[:0], leaving
is_unique stale as False on an otherwise-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 fixes a real bug found while testing: a non-empty user-supplied
proxy with duplicate index labels (via
to_dataframe(pcoll, proxy=df_with_dup_index)) crashed the old code with
'ValueError: Length mismatch', since drop_duplicates() shrinks the
index but not the frame before reassignment. Pre-slicing to iloc[:0]
before the dedup check resolves this.
Added regression test test_dataframe_xs_non_empty_duplicate_proxy
covering both DataFrame and Series non-empty duplicate-proxy cases.
Signed-off-by: ManvithPanyam
<[email protected]>
---------
Signed-off-by: ManvithPanyam
<[email protected]>
---
sdks/python/apache_beam/dataframe/frames.py | 90 +++++++++++++++++++-----
sdks/python/apache_beam/dataframe/frames_test.py | 72 ++++++++++++++++++-
2 files changed, 143 insertions(+), 19 deletions(-)
diff --git a/sdks/python/apache_beam/dataframe/frames.py
b/sdks/python/apache_beam/dataframe/frames.py
index 310791d2b58..452f519c58d 100644
--- a/sdks/python/apache_beam/dataframe/frames.py
+++ b/sdks/python/apache_beam/dataframe/frames.py
@@ -1091,27 +1091,81 @@ class
DeferredDataFrameOrSeries(frame_base.DeferredFrame):
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().iloc[:0]
+ if not proxy_frame.index.is_unique:
+ proxy_frame.index = proxy_frame.index.drop_duplicates()
+ k_val = key_series.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]
+
+ def unwrap_xs(ser):
+ if ser.empty:
+ raise KeyError(k_val)
+ return ser.iloc[0]
+
+ with expressions.allow_non_parallel_operations(True):
+ return frame_base.DeferredFrame.wrap(
+ expressions.ComputedExpression(
+ 'xs',
+ unwrap_xs, [intermediate],
+ proxy=xs_proxy,
+ requires_partition_by=partitionings.Singleton(),
+ preserves_partition_by=partitionings.Singleton()))
@property
def dtype(self):
diff --git a/sdks/python/apache_beam/dataframe/frames_test.py
b/sdks/python/apache_beam/dataframe/frames_test.py
index 7a03af6220b..290d13adc40 100644
--- a/sdks/python/apache_beam/dataframe/frames_test.py
+++ b/sdks/python/apache_beam/dataframe/frames_test.py
@@ -331,6 +331,19 @@ class DeferredFrameTest(_AbstractFrameTest):
lambda df: df.num_legs.xs(('bird', 'walks'), level=[0, 'locomotion']),
df)
+ # Test cases reported in BEAM-28559
+ df_single_index = df.reset_index().set_index('class')
+ self._run_test(
+ lambda df: df.num_legs.xs('mammal'), df_single_index,
check_proxy=False)
+ self._run_test(lambda df: df.num_legs.xs('bird'), df_single_index)
+
+ # Categorical Series single match
+ s_cat = pd.Series(
+ pd.Categorical(['a', 'b', 'c']),
+ index=['r1', 'r2', 'r3'],
+ name='cat_col')
+ self._run_test(lambda s: s.xs('r1'), s_cat, check_proxy=False)
+
def test_dataframe_xs(self):
# Test cases reported in BEAM-13421
df = pd.DataFrame(
@@ -342,10 +355,67 @@ class DeferredFrameTest(_AbstractFrameTest):
]),
columns=['provider', 'time', 'value'])
- self._run_test(lambda df: df.xs('state'), df.set_index(['provider']))
+ self._run_test(
+ lambda df: df.xs('state'),
+ df.set_index(['provider']),
+ check_proxy=False)
self._run_test(
lambda df: df.xs('state'), df.set_index(['provider', 'time']))
+ # Test cases reported in BEAM-28559
+ self._run_test(lambda df: df.xs('county'), df.set_index(['provider']))
+ self._run_test(
+ lambda df: df.xs(('state', 'day1')),
+ df.set_index(['provider', 'time']),
+ check_proxy=False)
+
+ df_unique = pd.DataFrame(
+ np.array([
+ ['state', 'day1', 12],
+ ['state', 'day2', 14],
+ ['county', 'day1', 9],
+ ]),
+ columns=['provider', 'time', 'value'])
+ self._run_test(
+ lambda df: df.xs(('state', 'day2')),
+ df_unique.set_index(['provider', 'time']))
+
+ # Categorical and extension dtype tests
+ df_cat = pd.DataFrame({
+ 'cat': pd.Categorical(['a', 'b', 'c']), 'val': [1, 2, 3]
+ },
+ index=['r1', 'r2', 'r3'])
+ self._run_test(lambda df: df.xs('r1'), df_cat)
+
+ df_dt_tz = pd.DataFrame({
+ 'dt': pd.Series([
+ pd.Timestamp('2023-01-01', tz='UTC'),
+ pd.Timestamp('2023-01-02', tz='UTC')
+ ],
+ dtype='datetime64[ns, UTC]'),
+ 'val': [1, 2]
+ },
+ index=['r1', 'r2'])
+ self._run_test(lambda df: df.xs('r1'), df_dt_tz)
+
+ df_null_int = pd.DataFrame({'num': pd.Series([1, 2, None], dtype='Int64')},
+ index=['r1', 'r2', 'r3'])
+ self._run_test(lambda df: df.xs('r1'), df_null_int)
+
+ def test_dataframe_xs_non_empty_duplicate_proxy(self):
+ df_dups = pd.DataFrame({'a': [1, 2]}, index=['x', 'x'])
+ p = beam.Pipeline()
+ deferred = to_dataframe(p | beam.Create([{'a': 1}]), proxy=df_dups)
+ res = deferred.xs('x')
+ self.assertIsInstance(res, frames.DeferredSeries)
+ self.assertTrue(res._expr.proxy().empty)
+
+ s_dups = pd.Series(pd.Categorical(['a', 'b']), index=['x', 'x'], name='s')
+ deferred_s = to_dataframe(
+ p | 'CreateSeries' >> beam.Create(['a']), proxy=s_dups)
+ res_s = deferred_s.xs('x')
+ self.assertIsInstance(res_s, frame_base.DeferredBase)
+
def test_set_column(self):
def new_column(df):
df['NewCol'] = df['Speed']
