[SPARK-56375][PANDAS/PS] Implement DataFrame.set_axis and Series.set_axis

[HenryBui777]

What changes were proposed in this pull request?

This PR implements DataFrame.set_axis() and Series.set_axis() for the Pandas API on Spark, which were previously unsupported (listed in missing/frame.py and missing/series.py).

Changes:

• python/pyspark/pandas/frame.py: Added DataFrame.set_axis(labels, axis=0) method.
• • axis=0 or axis='index': Reassigns the row index using the provided labels (delegates to set_index with a pd.Index)
• • axis=1 or axis='columns': Reassigns column labels using rename(columns=...)
• • Raises ValueError if the number of labels doesn't match the axis length
• • python/pyspark/pandas/series.py: Added Series.set_axis(labels, axis=0) method.
• • Only axis=0 is supported for Series (consistent with pandas)
• • Raises ValueError if labels length doesn't match the Series length
• • Raises ValueError if axis != 0
• • python/pyspark/pandas/missing/frame.py: Removed set_axis from _unsupported_function list
• • python/pyspark/pandas/missing/series.py: Removed set_axis from _unsupported_function list
    Supported behavior (matching pandas):

import pyspark.pandas as ps

# DataFrame - change row index
df = ps.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]})
df.set_axis(['x', 'y', 'z'])           # axis=0 (default)
df.set_axis(['x', 'y', 'z'], axis=0)
df.set_axis(['x', 'y', 'z'], axis='index')

# DataFrame - change column labels
df.set_axis(['I', 'II'], axis=1)
df.set_axis(['I', 'II'], axis='columns')

# Series
s = ps.Series([1, 2, 3])
s.set_axis(['a', 'b', 'c'])

Why are the changes needed?

DataFrame.set_axis() and Series.set_axis() are standard pandas APIs used extensively for reassigning index or column labels. They were completely unsupported in the Pandas API on Spark, causing PandasNotImplementedError even though the underlying functionality (index/column renaming) is already supported. This PR enables compatibility with pandas code that uses set_axis.

Reference: https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.set_axis.html

Does this PR introduce any user-facing change?

Yes. DataFrame.set_axis() and Series.set_axis() now work in the Pandas API on Spark instead of raising PandasNotImplementedError.

How was this patch tested?

Added new unit test files:

• python/pyspark/pandas/tests/frame/test_set_axis.py - covers:
• • test_set_axis_index: axis=0 with list, string 'index', and pd.Index labels
• • test_set_axis_columns: axis=1 with list, string 'columns', and pd.Index labels
• • test_set_axis_errors: ValueError on length mismatch and invalid axis
• • test_set_axis_numeric_index: numeric index labels
• • python/pyspark/pandas/tests/series/test_set_axis.py - covers:
• • test_set_axis_index: axis=0 with list, string 'index', and pd.Index labels
• • test_set_axis_errors: ValueError on length mismatch and invalid axis (axis=1)
• • test_set_axis_named: named Series

Was this patch authored or co-authored using generative AI tooling?

No.