This is an automated email from the ASF dual-hosted git repository.
uros-b 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 b5eb682ab081 [SPARK-46161][PS] Add axis=1 support for DataFrame.diff
b5eb682ab081 is described below
commit b5eb682ab081e26d238f8a81bad2a26c4353cdfc
Author: Devin Petersohn <[email protected]>
AuthorDate: Wed Aug 12 11:50:06 2026 -0700
[SPARK-46161][PS] Add axis=1 support for DataFrame.diff
### What changes were proposed in this pull request?
Implement column-wise differencing for `DataFrame.diff(axis=1)`.
### Why are the changes needed?
It is a missing pandas API parameter.
### Does this PR introduce _any_ user-facing change?
Yes, `DataFrame.diff(axis=1)` now works.
### How was this patch tested?
CI
### Was this patch authored or co-authored using generative AI tooling?
Co-authored-by: Claude Opus 4
Closes #54516 from devin-petersohn/devin/diff-axis-columns.
Authored-by: Devin Petersohn <[email protected]>
Signed-off-by: Holden Karau <[email protected]>
(cherry picked from commit a9689fb515d2b1f83efa705b2c1eb63c82330ba1)
Signed-off-by: Uros Bojanic <[email protected]>
---
python/pyspark/pandas/frame.py | 40 +++++++++++++++++-----
.../pandas/tests/computation/test_compute.py | 8 +++--
2 files changed, 37 insertions(+), 11 deletions(-)
diff --git a/python/pyspark/pandas/frame.py b/python/pyspark/pandas/frame.py
index d196a85070bd..fdc7b6f7f49d 100644
--- a/python/pyspark/pandas/frame.py
+++ b/python/pyspark/pandas/frame.py
@@ -4946,7 +4946,6 @@ defaultdict(<class 'list'>, {'col..., 'col...})]
lambda psser: psser._shift(periods, fill_value),
should_resolve=True
)
- # TODO(SPARK-46161): axis should support 1 or 'columns' either at this
moment
def diff(self, periods: int = 1, axis: Axis = 0) -> "DataFrame":
"""
First discrete difference of element.
@@ -4954,7 +4953,7 @@ defaultdict(<class 'list'>, {'col..., 'col...})]
Calculates the difference of a DataFrame element compared with another
element in the
DataFrame (default is the element in the same column of the previous
row).
- .. note:: the current implementation of diff uses Spark's Window
without
+ .. note:: When ``axis=0``, the current implementation of diff uses
Spark's Window without
specifying partition specification. This leads to moving all data
into
a single partition in a single machine and could cause serious
performance degradation. Avoid this method with very large
datasets.
@@ -4963,8 +4962,8 @@ defaultdict(<class 'list'>, {'col..., 'col...})]
----------
periods : int, default 1
Periods to shift for calculating difference, accepts negative
values.
- axis : int, default 0 or 'index'
- Can only be set to 0 now.
+ axis : {0 or 'index', 1 or 'columns'}, default 0
+ Take difference over rows (0) or columns (1).
Returns
-------
@@ -5014,12 +5013,37 @@ defaultdict(<class 'list'>, {'col..., 'col...})]
3 -1.0 -2.0 -9.0
4 -1.0 -3.0 -11.0
5 NaN NaN NaN
+
+ Difference with previous column
+
+ >>> df.diff(axis=1)
+ a b c
+ 0 NaN 0 0
+ 1 NaN -1 3
+ 2 NaN -1 7
+ 3 NaN -1 13
+ 4 NaN 0 20
+ 5 NaN 2 28
"""
axis = validate_axis(axis)
- if axis != 0:
- raise NotImplementedError('axis should be either 0 or "index"
currently.')
-
- return self._apply_series_op(lambda psser: psser._diff(periods),
should_resolve=True)
+ if axis == 0:
+ return self._apply_series_op(lambda psser: psser._diff(periods),
should_resolve=True)
+ else:
+ column_labels = self._internal.column_labels
+ data_col_names = self._internal.data_spark_column_names
+ new_columns: list[PySparkColumn] = []
+ for i, label in enumerate(column_labels):
+ prev_idx = i - periods
+ if 0 <= prev_idx < len(column_labels):
+ prev_label = column_labels[prev_idx]
+ cur_col = self._internal.spark_column_for(label)
+ prev_col = self._internal.spark_column_for(prev_label)
+ new_columns.append(cur_col - prev_col)
+ else:
+ col_type = self._internal.spark_type_for(label)
+
new_columns.append(F.lit(None).cast(col_type).alias(data_col_names[i]))
+ internal = self._internal.with_new_columns(new_columns)
+ return DataFrame(internal)
def nunique(
self,
diff --git a/python/pyspark/pandas/tests/computation/test_compute.py
b/python/pyspark/pandas/tests/computation/test_compute.py
index 4538509cd4d3..924046fe256b 100644
--- a/python/pyspark/pandas/tests/computation/test_compute.py
+++ b/python/pyspark/pandas/tests/computation/test_compute.py
@@ -198,9 +198,11 @@ class FrameComputeMixin:
msg = "should be an int"
with self.assertRaisesRegex(TypeError, msg):
psdf.diff(1.5)
- msg = 'axis should be either 0 or "index" currently.'
- with self.assertRaisesRegex(NotImplementedError, msg):
- psdf.diff(axis=1)
+
+ # axis=1: difference across columns
+ self.assert_eq(pdf.diff(axis=1), psdf.diff(axis=1))
+ self.assert_eq(pdf.diff(periods=2, axis=1), psdf.diff(periods=2,
axis=1))
+ self.assert_eq(pdf.diff(periods=-1, axis=1), psdf.diff(periods=-1,
axis=1))
# multi-index columns
columns = pd.MultiIndex.from_tuples([("x", "Col1"), ("x", "Col2"),
("y", "Col3")])
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
