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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
