From 5c6b20af59cc7e154591aa206f49eb47d00b8de6 Mon Sep 17 00:00:00 2001 From: ManvithPanyam <250704031+ManvithPanyam@users.noreply.github.com> Date: Sat, 22 Aug 2026 11:56:48 +0530 Subject: [PATCH 1/3] 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 <250704031+ManvithPanyam@users.noreply.github.com> --- sdks/python/apache_beam/dataframe/frames.py | 100 +++++++++++++++--- .../apache_beam/dataframe/frames_test.py | 29 ++++- 2 files changed, 111 insertions(+), 18 deletions(-) diff --git a/sdks/python/apache_beam/dataframe/frames.py b/sdks/python/apache_beam/dataframe/frames.py index 310791d2b58f..e8a62a7fb25e 100644 --- a/sdks/python/apache_beam/dataframe/frames.py +++ b/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()] + for col in proxy_frame.columns + } + dummy_df = pd.DataFrame(dummy_data, index=dummy_index) + xs_proxy = dummy_df.xs(k_val, **kwargs) + if isinstance(xs_proxy, pd.DataFrame) or isinstance(xs_proxy, + pd.Series): + xs_proxy = xs_proxy.iloc[:0] + else: + val = proxy_frame.dtype.type() + 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_ser = pd.Series([val], + index=dummy_index, + dtype=proxy_frame.dtype, + name=proxy_frame.name) + xs_proxy = dummy_ser.xs(k_val, **kwargs) + if isinstance(xs_proxy, pd.Series): + xs_proxy = xs_proxy.iloc[:0] else: - return result + xs_proxy = proxy_frame.dtype.type() - # key should be in this partition, call xs. Will raise KeyError if not - # present. - return frame.xs(key.item()) + def unwrap_xs(ser): + if ser.empty: + raise KeyError(k_val) + return ser.iloc[0] - 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())) + 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 7a03af6220b8..0c9a06bd27c7 100644 --- a/sdks/python/apache_beam/dataframe/frames_test.py +++ b/sdks/python/apache_beam/dataframe/frames_test.py @@ -331,6 +331,12 @@ def test_series_xs(self): 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) + def test_dataframe_xs(self): # Test cases reported in BEAM-13421 df = pd.DataFrame( @@ -342,10 +348,31 @@ def test_dataframe_xs(self): ]), 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'])) + def test_set_column(self): def new_column(df): df['NewCol'] = df['Speed'] From ac0525f6a32f1a80bd9aba8c8e6a536a2144f635 Mon Sep 17 00:00:00 2001 From: ManvithPanyam <250704031+ManvithPanyam@users.noreply.github.com> Date: Tue, 1 Sep 2026 10:06:46 +0530 Subject: [PATCH 2/3] 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 <250704031+ManvithPanyam@users.noreply.github.com> --- sdks/python/apache_beam/dataframe/frames.py | 37 ++++++++----------- .../apache_beam/dataframe/frames_test.py | 29 +++++++++++++++ 2 files changed, 45 insertions(+), 21 deletions(-) diff --git a/sdks/python/apache_beam/dataframe/frames.py b/sdks/python/apache_beam/dataframe/frames.py index e8a62a7fb25e..17bd04ba1338 100644 --- a/sdks/python/apache_beam/dataframe/frames.py +++ b/sdks/python/apache_beam/dataframe/frames.py @@ -1136,34 +1136,29 @@ def xs_partitioned_wrapped(frame, key): proxy_frame = reindexed._expr.proxy() k_val = key_series.iloc[0] if isinstance(proxy_frame, pd.DataFrame): + if not proxy_frame.index.is_unique: + proxy_frame = proxy_frame.copy() + proxy_frame.index = proxy_frame.index.drop_duplicates() 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()] - for col in proxy_frame.columns - } - dummy_df = pd.DataFrame(dummy_data, index=dummy_index) - xs_proxy = dummy_df.xs(k_val, **kwargs) - if isinstance(xs_proxy, pd.DataFrame) or isinstance(xs_proxy, - pd.Series): + 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: - val = proxy_frame.dtype.type() - 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_ser = pd.Series([val], - index=dummy_index, - dtype=proxy_frame.dtype, - name=proxy_frame.name) - xs_proxy = dummy_ser.xs(k_val, **kwargs) - if isinstance(xs_proxy, pd.Series): - xs_proxy = xs_proxy.iloc[:0] - else: + try: xs_proxy = proxy_frame.dtype.type() + except TypeError: + 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 not proxy_frame.index.is_unique: + proxy_frame = proxy_frame.copy() + proxy_frame.index = proxy_frame.index.drop_duplicates() + xs_proxy = proxy_frame.reindex(dummy_index).iloc[0] def unwrap_xs(ser): if ser.empty: diff --git a/sdks/python/apache_beam/dataframe/frames_test.py b/sdks/python/apache_beam/dataframe/frames_test.py index 0c9a06bd27c7..a2cb2e498005 100644 --- a/sdks/python/apache_beam/dataframe/frames_test.py +++ b/sdks/python/apache_beam/dataframe/frames_test.py @@ -337,6 +337,13 @@ def test_series_xs(self): 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( @@ -373,6 +380,28 @@ def test_dataframe_xs(self): 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_set_column(self): def new_column(df): df['NewCol'] = df['Speed'] From 5ff9b633e8c819030a098bec7e1c2e230c324792 Mon Sep 17 00:00:00 2001 From: ManvithPanyam <250704031+ManvithPanyam@users.noreply.github.com> Date: Fri, 4 Sep 2026 11:18:04 +0530 Subject: [PATCH 3/3] fix(dataframe): simplify xs() proxy construction and fix non-empty duplicate-proxy crash MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 <250704031+ManvithPanyam@users.noreply.github.com> --- sdks/python/apache_beam/dataframe/frames.py | 23 +++++++------------ .../apache_beam/dataframe/frames_test.py | 14 +++++++++++ 2 files changed, 22 insertions(+), 15 deletions(-) diff --git a/sdks/python/apache_beam/dataframe/frames.py b/sdks/python/apache_beam/dataframe/frames.py index 17bd04ba1338..452f519c58db 100644 --- a/sdks/python/apache_beam/dataframe/frames.py +++ b/sdks/python/apache_beam/dataframe/frames.py @@ -1133,16 +1133,16 @@ def xs_partitioned_wrapped(frame, key): requires_partition_by=partitionings.Index(list(range(key_size))), preserves_partition_by=partitionings.Singleton()) - proxy_frame = reindexed._expr.proxy() + 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): - if not proxy_frame.index.is_unique: - proxy_frame = proxy_frame.copy() - proxy_frame.index = proxy_frame.index.drop_duplicates() - 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)): @@ -1151,13 +1151,6 @@ def xs_partitioned_wrapped(frame, key): try: xs_proxy = proxy_frame.dtype.type() except TypeError: - 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 not proxy_frame.index.is_unique: - proxy_frame = proxy_frame.copy() - proxy_frame.index = proxy_frame.index.drop_duplicates() xs_proxy = proxy_frame.reindex(dummy_index).iloc[0] def unwrap_xs(ser): diff --git a/sdks/python/apache_beam/dataframe/frames_test.py b/sdks/python/apache_beam/dataframe/frames_test.py index a2cb2e498005..290d13adc405 100644 --- a/sdks/python/apache_beam/dataframe/frames_test.py +++ b/sdks/python/apache_beam/dataframe/frames_test.py @@ -402,6 +402,20 @@ def test_dataframe_xs(self): 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']