This is an automated email from the ASF dual-hosted git repository.
HyukjinKwon pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 6c4e7077218c [SPARK-46158][PYTHON][PS] Support axis columns in
DataFrame.xs
6c4e7077218c is described below
commit 6c4e7077218c474b360fc4fe8b92eb29bfaecbf1
Author: Nguyễn Mạnh hà <[email protected]>
AuthorDate: Tue Jun 23 07:17:19 2026 +0900
[SPARK-46158][PYTHON][PS] Support axis columns in DataFrame.xs
### What changes were proposed in this pull request?
This PR adds support for `DataFrame.xs` with `axis=1` / `axis="columns"` in
pandas API on Spark.
It supports selecting cross-sections from columns, including single-level
columns, MultiIndex columns, and level-based selection by level number or level
name.
### Why are the changes needed?
`DataFrame.xs` currently supports index-axis selection, but column-axis
selection raises `NotImplementedError`.
This improves compatibility with pandas and addresses SPARK-46158.
### Does this PR introduce _any_ user-facing change?
Yes.
Previously, `DataFrame.xs(..., axis=1)` raised `NotImplementedError`.
After this change, users can select column cross-sections, for example:
~~~python
psdf.xs("metrics", axis=1)
psdf.xs("metrics", axis="columns")
psdf.xs("num_legs", axis=1, level="feature")
~~~
### How was this patch tested?
Added positive and negative test cases in `FrameIndexingMixin.test_xs`
covering `axis=1`, `axis="columns"`, MultiIndex columns, level-based selection,
missing column keys, and invalid column levels.
Also ran:
~~~bash
python -m py_compile python/pyspark/pandas/frame.py
python/pyspark/pandas/tests/indexes/test_indexing.py
git diff --check
~~~
I attempted to run the targeted unittest locally, but the source checkout
has not built Spark jars yet.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: OpenAI Codex (GPT-5)
Closes #56284 from manhha2502/SPARK-46158-xs-axis-columns.
Authored-by: Nguyễn Mạnh hà <[email protected]>
Signed-off-by: Hyukjin Kwon <[email protected]>
(cherry picked from commit 7767678ad2094b30025cc71e0a1855f1456e5603)
Signed-off-by: Hyukjin Kwon <[email protected]>
---
python/pyspark/pandas/frame.py | 131 +++++++++++++++++++--
.../pyspark/pandas/tests/indexes/test_indexing.py | 45 ++++++-
2 files changed, 161 insertions(+), 15 deletions(-)
diff --git a/python/pyspark/pandas/frame.py b/python/pyspark/pandas/frame.py
index 4d366253c1d3..4b9ce84b3b7a 100644
--- a/python/pyspark/pandas/frame.py
+++ b/python/pyspark/pandas/frame.py
@@ -3584,8 +3584,9 @@ defaultdict(<class 'list'>, {'col..., 'col...})]
self._update_internal_frame(self.drop(columns=item)._internal)
return result
- # TODO(SPARK-46158): add axis parameter can work when '1' or 'columns'
- def xs(self, key: Name, axis: Axis = 0, level: Optional[int] = None) ->
DataFrameOrSeries:
+ def xs(
+ self, key: Name, axis: Axis = 0, level: Optional[Union[int, Name]] =
None
+ ) -> DataFrameOrSeries:
"""
Return cross-section from the DataFrame.
@@ -3596,9 +3597,8 @@ defaultdict(<class 'list'>, {'col..., 'col...})]
----------
key : label or tuple of label
Label contained in the index, or partially in a MultiIndex.
- axis : 0 or 'index', default 0
+ axis : {0 or 'index', 1 or 'columns'}, default 0
Axis to retrieve cross-section on.
- currently only support 0 or 'index'
level : object, defaults to first n levels (n=1 or len(key))
In case of a key partially contained in a MultiIndex, indicate
which levels are used. Levels can be referred by label or position.
@@ -3660,6 +3660,16 @@ defaultdict(<class 'list'>, {'col..., 'col...})]
num_legs num_wings
class locomotion
mammal walks 4 0
+
+ Get values at specified column
+
+ >>> df.xs('num_legs', axis=1) # doctest: +NORMALIZE_WHITESPACE
+ class animal locomotion
+ mammal cat walks 4
+ dog walks 4
+ bat flies 2
+ bird penguin walks 2
+ Name: num_legs, dtype: int64
"""
from pyspark.pandas.series import first_series
@@ -3670,8 +3680,85 @@ defaultdict(<class 'list'>, {'col..., 'col...})]
raise KeyError(key)
axis = validate_axis(axis)
- if axis != 0:
- raise NotImplementedError('axis should be either 0 or "index"
currently.')
+ if axis == 1:
+ column_level = self._internal.column_labels_level
+ if column_level == 1:
+ if level is not None:
+ raise TypeError("Index must be a MultiIndex")
+ if is_name_like_tuple(key):
+ raise KeyError(key)
+
+ column_key = cast(Label, key if is_name_like_tuple(key) else
(key,))
+ if len(column_key) > column_level:
+ raise KeyError(
+ "Key length ({}) exceeds index depth
({})".format(len(column_key), column_level)
+ )
+
+ if level is None:
+ level_idx = 0
+ elif isinstance(level, int):
+ level_idx = level
+ else:
+ level_name = cast(Label, level if is_name_like_tuple(level)
else (level,))
+ try:
+ level_idx =
self._internal.column_label_names.index(level_name)
+ except ValueError:
+ raise KeyError("Level {} not
found".format(name_like_string(level_name)))
+
+ original_level_idx = level_idx
+ if level_idx < 0:
+ level_idx += column_level
+ if level_idx < 0:
+ raise IndexError(
+ "Too many levels: Index has only {} levels, "
+ "{} is not a valid level number".format(column_level,
original_level_idx)
+ )
+ if level_idx >= column_level:
+ raise IndexError(
+ "Too many levels: Index has only {} levels, not {}".format(
+ column_level, level_idx + 1
+ )
+ )
+
+ drop_levels = range(level_idx, level_idx + len(column_key))
+ selected = [
+ (label, scol, field)
+ for label, scol, field in zip(
+ self._internal.column_labels,
+ self._internal.data_spark_columns,
+ self._internal.data_fields,
+ )
+ if tuple(label[i] for i in drop_levels) == column_key
+ ]
+ if not selected:
+ raise KeyError(key)
+
+ selected_column_labels, data_spark_columns, data_fields =
zip(*selected)
+ # Match pandas by dropping the selected column levels from the
result.
+ new_column_labels = [
+ tuple(label[i] for i in range(column_level) if i not in
drop_levels)
+ for label in selected_column_labels
+ ]
+ if len(new_column_labels[0]) == 0:
+ if len(selected_column_labels) == 1:
+ return self._psser_for(selected_column_labels[0])
+ else:
+ new_column_labels = list(selected_column_labels)
+ column_label_names = self._internal.column_label_names
+ else:
+ column_label_names = [
+ name
+ for i, name in enumerate(self._internal.column_label_names)
+ if i not in drop_levels
+ ]
+
+ internal = self._internal.copy(
+ column_labels=list(new_column_labels),
+ data_spark_columns=list(data_spark_columns),
+ data_fields=list(data_fields),
+ column_label_names=list(column_label_names),
+ )
+ return DataFrame(internal)
if not is_name_like_tuple(key):
key = (key,)
@@ -3682,10 +3769,25 @@ defaultdict(<class 'list'>, {'col..., 'col...})]
)
)
if level is None:
- level = 0
+ index_level = 0
+ elif isinstance(level, int):
+ index_level = level
+ else:
+ level_name = cast(Label, level if is_name_like_tuple(level) else
(level,))
+ if self._internal.index_names.count(level_name) > 1:
+ raise ValueError(
+ "The name {} occurs multiple times, use a level
number".format(
+ name_like_string(level_name)
+ )
+ )
+ try:
+ index_level = self._internal.index_names.index(level_name)
+ except ValueError:
+ raise KeyError("Level {} not
found".format(name_like_string(level_name)))
rows = [
- self._internal.index_spark_columns[lvl] == index for lvl, index in
enumerate(key, level)
+ self._internal.index_spark_columns[lvl] == index
+ for lvl, index in enumerate(key, index_level)
]
internal = self._internal.with_filter(reduce(lambda x, y: x & y, rows))
@@ -3700,11 +3802,16 @@ defaultdict(<class 'list'>, {'col..., 'col...})]
return first_series(DataFrame(pdf.transpose()))
else:
index_spark_columns = (
- internal.index_spark_columns[:level]
- + internal.index_spark_columns[level + len(key) :]
+ internal.index_spark_columns[:index_level]
+ + internal.index_spark_columns[index_level + len(key) :]
+ )
+ index_names = (
+ internal.index_names[:index_level] +
internal.index_names[index_level + len(key) :]
+ )
+ index_fields = (
+ internal.index_fields[:index_level]
+ + internal.index_fields[index_level + len(key) :]
)
- index_names = internal.index_names[:level] +
internal.index_names[level + len(key) :]
- index_fields = internal.index_fields[:level] +
internal.index_fields[level + len(key) :]
internal = internal.copy(
index_spark_columns=index_spark_columns,
diff --git a/python/pyspark/pandas/tests/indexes/test_indexing.py
b/python/pyspark/pandas/tests/indexes/test_indexing.py
index 78975a457d58..166a20e307c8 100644
--- a/python/pyspark/pandas/tests/indexes/test_indexing.py
+++ b/python/pyspark/pandas/tests/indexes/test_indexing.py
@@ -115,12 +115,51 @@ class FrameIndexingMixin:
pd.concat([pdf, pdf]).xs(("mammal", "dog", "walks")),
)
self.assert_eq(psdf.xs("cat", level=1), pdf.xs("cat", level=1))
+ self.assert_eq(psdf.xs("cat", level="animal"), pdf.xs("cat",
level="animal"))
self.assert_eq(psdf.xs("flies", level=2), pdf.xs("flies", level=2))
+ self.assert_eq(psdf.xs("mammal", level="class"), pdf.xs("mammal",
level="class"))
self.assert_eq(psdf.xs("mammal", level=-3), pdf.xs("mammal", level=-3))
- msg = 'axis should be either 0 or "index" currently.'
- with self.assertRaisesRegex(NotImplementedError, msg):
- psdf.xs("num_wings", axis=1)
+ self.assert_eq(psdf.xs("num_wings", axis=1), pdf.xs("num_wings",
axis=1))
+
+ columns = pd.MultiIndex.from_tuples(
+ [
+ ("metrics", "num_legs"),
+ ("metrics", "num_wings"),
+ ],
+ names=["group", "feature"],
+ )
+ pdf_with_columns = pdf.copy()
+ pdf_with_columns.columns = columns
+ psdf_with_columns = ps.from_pandas(pdf_with_columns)
+ self.assert_eq(
+ psdf_with_columns.xs("metrics", axis=1),
+ pdf_with_columns.xs("metrics", axis=1),
+ )
+ self.assert_eq(
+ psdf_with_columns.xs("metrics", axis="columns"),
+ pdf_with_columns.xs("metrics", axis="columns"),
+ )
+ self.assert_eq(
+ psdf_with_columns.xs(("metrics", "num_legs"), axis=1),
+ pdf_with_columns.xs(("metrics", "num_legs"), axis=1),
+ )
+ self.assert_eq(
+ psdf_with_columns.xs("num_legs", axis=1, level=1),
+ pdf_with_columns.xs("num_legs", axis=1, level=1),
+ )
+ self.assert_eq(
+ psdf_with_columns.xs("num_legs", axis=1, level="feature"),
+ pdf_with_columns.xs("num_legs", axis=1, level="feature"),
+ )
+ self.assert_eq(
+ psdf_with_columns.xs("num_legs", axis=1, level=-1),
+ pdf_with_columns.xs("num_legs", axis=1, level=-1),
+ )
+ self.assertRaises(KeyError, lambda: psdf_with_columns.xs("unknown",
axis=1))
+ self.assertRaises(KeyError, lambda: psdf_with_columns.xs(("metrics",
"unknown"), axis=1))
+ self.assertRaises(IndexError, lambda: psdf_with_columns.xs("metrics",
axis=1, level=2))
+
with self.assertRaises(KeyError):
psdf.xs(("mammal", "dog", "walk"))
msg = r"'Key length \(4\) exceeds index depth \(3\)'"
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
