Skip to content

Commit

Permalink
ENH: pyarrow temporal dtypes support quantile in some cases (#50868)
Browse files Browse the repository at this point in the history
  • Loading branch information
jbrockmendel authored Jan 20, 2023
1 parent f69efb6 commit 6ecb52e
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 6 deletions.
18 changes: 17 additions & 1 deletion pandas/core/arrays/arrow/array.py
Original file line number Diff line number Diff line change
Expand Up @@ -1208,7 +1208,23 @@ def _quantile(
-------
same type as self
"""
result = pc.quantile(self._data, q=qs, interpolation=interpolation)
pa_dtype = self._data.type

data = self._data
if pa.types.is_temporal(pa_dtype) and interpolation in ["lower", "higher"]:
# https://github.com/apache/arrow/issues/33769 in these cases
# we can cast to ints and back
nbits = pa_dtype.bit_width
if nbits == 32:
data = data.cast(pa.int32())
else:
data = data.cast(pa.int64())

result = pc.quantile(data, q=qs, interpolation=interpolation)

if pa.types.is_temporal(pa_dtype) and interpolation in ["lower", "higher"]:
result = result.cast(pa_dtype)

return type(self)(result)

def _mode(self: ArrowExtensionArrayT, dropna: bool = True) -> ArrowExtensionArrayT:
Expand Down
14 changes: 9 additions & 5 deletions pandas/tests/extension/test_arrow.py
Original file line number Diff line number Diff line change
Expand Up @@ -1279,7 +1279,11 @@ def test_quantile(data, interpolation, quantile, request):
ser.quantile(q=quantile, interpolation=interpolation)
return

if not (pa.types.is_integer(pa_dtype) or pa.types.is_floating(pa_dtype)):
if pa.types.is_integer(pa_dtype) or pa.types.is_floating(pa_dtype):
pass
elif pa.types.is_temporal(data._data.type) and interpolation in ["lower", "higher"]:
pass
else:
request.node.add_marker(
pytest.mark.xfail(
raises=pa.ArrowNotImplementedError,
Expand All @@ -1293,10 +1297,10 @@ def test_quantile(data, interpolation, quantile, request):
assert result == data[0]
else:
# Just check the values
result = result.astype("float64[pyarrow]")
expected = pd.Series(
data.take([0, 0]).astype("float64[pyarrow]"), index=[0.5, 0.5]
)
expected = pd.Series(data.take([0, 0]), index=[0.5, 0.5])
if pa.types.is_integer(pa_dtype) or pa.types.is_floating(pa_dtype):
expected = expected.astype("float64[pyarrow]")
result = result.astype("float64[pyarrow]")
tm.assert_series_equal(result, expected)


Expand Down

0 comments on commit 6ecb52e

Please sign in to comment.