Skip to content

Commit

Permalink
[SPARK-40579][PS] GroupBy.first should skip NULLs
Browse files Browse the repository at this point in the history
### What changes were proposed in this pull request?
make `GroupBy.first` skip nulls

### Why are the changes needed?
to fix the behavior difference

```
In [1]:
   ...: import pandas as pd
   ...: import numpy as np
   ...: import pyspark.pandas as ps
   ...:
   ...: pdf = pd.DataFrame({"A": [1, 2, 1, 2],"B": [-1.5, np.nan, -3.2, 0.1],})
   ...: psdf = ps.from_pandas(pdf)
   ...:

In [2]: pdf.groupby("A").first()
Out[2]:
     B
A
1 -1.5
2  0.1

In [3]: psdf.groupby("A").first()

     B
A
1 -1.5
2  NaN
```

### Does this PR introduce _any_ user-facing change?
yes, updated `GroupBy.first` will skip NULLs

### How was this patch tested?
added UT

Closes apache#38017 from zhengruifeng/ps_first_skip_na.

Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Hyukjin Kwon <[email protected]>
  • Loading branch information
zhengruifeng authored and HyukjinKwon committed Sep 28, 2022
1 parent a96ac7e commit e932e0a
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 1 deletion.
3 changes: 2 additions & 1 deletion python/pyspark/pandas/groupby.py
Original file line number Diff line number Diff line change
Expand Up @@ -449,7 +449,8 @@ def first(self, numeric_only: Optional[bool] = False) -> FrameLike:
2 False 3
"""
return self._reduce_for_stat_function(
F.first, accepted_spark_types=(NumericType, BooleanType) if numeric_only else None
lambda col: F.first(col, ignorenulls=True),
accepted_spark_types=(NumericType, BooleanType) if numeric_only else None,
)

def last(self, numeric_only: Optional[bool] = False) -> FrameLike:
Expand Down
11 changes: 11 additions & 0 deletions python/pyspark/pandas/tests/test_groupby.py
Original file line number Diff line number Diff line change
Expand Up @@ -1419,6 +1419,17 @@ def test_first(self):
self._test_stat_func(lambda groupby_obj: groupby_obj.first(numeric_only=None))
self._test_stat_func(lambda groupby_obj: groupby_obj.first(numeric_only=True))

pdf = pd.DataFrame(
{
"A": [1, 2, 1, 2],
"B": [-1.5, np.nan, -3.2, 0.1],
}
)
psdf = ps.from_pandas(pdf)
self.assert_eq(
pdf.groupby("A").first().sort_index(), psdf.groupby("A").first().sort_index()
)

def test_last(self):
self._test_stat_func(lambda groupby_obj: groupby_obj.last())
self._test_stat_func(lambda groupby_obj: groupby_obj.last(numeric_only=None))
Expand Down

0 comments on commit e932e0a

Please sign in to comment.