ENH: Implement cum* methods for PyArrow strings #60633
Conversation
rhshadrach
added
Enhancement
Strings
The only other way I think that would be reasonable to implement this in would be to use nanoarrow, but that's a larger implementation. I think this is fine for now - just not very performant but that can always be improved later |
pandas/conftest.py
Outdated
| @@ -1317,6 +1317,22 @@ def nullable_string_dtype(request): | |||
| return request.param | |||
|
|
|||
|
|
|||
| @pytest.fixture( | |||
| params=[ | |||
| pytest.param("str[pyarrow]", marks=td.skip_if_no("pyarrow")), | |||
There was a problem hiding this comment.
I was going to comment: I don't think this can work. Although it is then strange the tests are passing :) But it seems this was not doing what I think you expected it was doing -> #60661
I would use the same approach of creating the dtype through StringDtype(..) like in some of the fixtures above
pandas/core/arrays/arrow/array.py
Outdated
| nulls = pc.is_null(pa_array) | ||
| idx = pc.index(nulls, True).as_py() | ||
| tail = pa.nulls(len(pa_array) - idx, type=pa_array.type) | ||
| pa_array = pa_array[:idx].combine_chunks() |
There was a problem hiding this comment.
Is the combine chunks needed here? (I would expect that the conversion to numpy (when calling the numpy func) will do this automatically (and potentially more efficiently))
| (["x", "z", "y"], "cumsum", False, ["x", "xz", "xzy"]), | ||
| (["x", pd.NA, "y"], "cumsum", True, ["x", "x", "xy"]), | ||
| (["x", pd.NA, "y"], "cumsum", False, ["x", pd.NA, pd.NA]), | ||
| ([pd.NA, "x", "y"], "cumsum", True, ["", "x", "xy"]), |
There was a problem hiding this comment.
It seems that for numerical data, we actually (somewhat inconsistently?) propagate leading NAs:
In [7]: pd.Series([np.nan, 0.5, 2.5]).cumsum()
Out[7]:
0 NaN
1 0.5
2 3.0
dtype: float64
(i.e. the result doesn't have 0.0 for the first element)
There was a problem hiding this comment.
Actually not related to "leading" NAs. It seems what is happening is that missing values are ignored to calculate the cumulative result, but then are propagated to the result elementwise. This is also shown in the docstring example of cumsum, so this seems intentional.
There was a problem hiding this comment.
Thanks for catching this. Agreed we should match this behavior. I do find it odd, but that's (possibly) for another day!
doc/source/whatsnew/vX.X.X.rstfile if fixing a bug or adding a new feature.Implements
cumsum,cummin, andcummaxfor PyArrow-backed strings. I don't see a way to do this without passing through NumPy - if there are other ideas for approaching happy to give those a shot.cc @WillAyd @jorisvandenbossche