ENH: support reductions for pyarrow temporal types (#50998)

* ENH: support reductions for pyarrow temporal types

* unit check

* lint fixup

Author: jbrockmendel

pandas/compat/__init__.py

@@ -30,6 +30,7 @@
     pa_version_under7p0,
     pa_version_under8p0,
     pa_version_under9p0,
+    pa_version_under11p0,
 )
 
 
@@ -159,6 +160,7 @@ def get_lzma_file() -> type[pandas.compat.compressors.LZMAFile]:
     "pa_version_under7p0",
     "pa_version_under8p0",
     "pa_version_under9p0",
+    "pa_version_under11p0",
     "IS64",
     "PY39",
     "PY310",

pandas/compat/pyarrow.py

@@ -13,8 +13,10 @@
     pa_version_under8p0 = _palv < Version("8.0.0")
     pa_version_under9p0 = _palv < Version("9.0.0")
     pa_version_under10p0 = _palv < Version("10.0.0")
+    pa_version_under11p0 = _palv < Version("11.0.0")
 except ImportError:
     pa_version_under7p0 = True
     pa_version_under8p0 = True
     pa_version_under9p0 = True
     pa_version_under10p0 = True
+    pa_version_under11p0 = True

pandas/core/arrays/arrow/array.py

@@ -31,6 +31,7 @@
     pa_version_under7p0,
     pa_version_under8p0,
     pa_version_under9p0,
+    pa_version_under11p0,
 )
 from pandas.util._decorators import doc
 from pandas.util._validators import validate_fillna_kwargs
@@ -134,6 +135,16 @@ def floordiv_compat(
 ArrowExtensionArrayT = TypeVar("ArrowExtensionArrayT", bound="ArrowExtensionArray")
 
 
+def get_unit_from_pa_dtype(pa_dtype):
+    # https://github.com/pandas-dev/pandas/pull/50998#discussion_r1100344804
+    if pa_version_under11p0:
+        unit = str(pa_dtype).split("[", 1)[-1][:-1]
+        if unit not in ["s", "ms", "us", "ns"]:
+            raise ValueError(pa_dtype)
+        return unit
+    return pa_dtype.unit
+
+
 def to_pyarrow_type(
     dtype: ArrowDtype | pa.DataType | Dtype | None,
 ) -> pa.DataType | None:
@@ -1043,6 +1054,13 @@ def _reduce(self, name: str, *, skipna: bool = True, **kwargs):
         elif name in ["min", "max", "sum"] and pa.types.is_duration(pa_type):
             data_to_reduce = self._data.cast(pa.int64())
 
+        elif name in ["median", "mean", "std", "sem"] and pa.types.is_temporal(pa_type):
+            nbits = pa_type.bit_width
+            if nbits == 32:
+                data_to_reduce = self._data.cast(pa.int32())
+            else:
+                data_to_reduce = self._data.cast(pa.int64())
+
         if name == "sem":
 
             def pyarrow_meth(data, skip_nulls, **kwargs):
@@ -1080,6 +1098,22 @@ def pyarrow_meth(data, skip_nulls, **kwargs):
 
         if name in ["min", "max", "sum"] and pa.types.is_duration(pa_type):
             result = result.cast(pa_type)
+        if name in ["median", "mean"] and pa.types.is_temporal(pa_type):
+            result = result.cast(pa_type)
+        if name in ["std", "sem"] and pa.types.is_temporal(pa_type):
+            result = result.cast(pa.int64())
+            if pa.types.is_duration(pa_type):
+                result = result.cast(pa_type)
+            elif pa.types.is_time(pa_type):
+                unit = get_unit_from_pa_dtype(pa_type)
+                result = result.cast(pa.duration(unit))
+            elif pa.types.is_date(pa_type):
+                # go with closest available unit, i.e. "s"
+                result = result.cast(pa.duration("s"))
+            else:
+                # i.e. timestamp
+                result = result.cast(pa.duration(pa_type.unit))
+
         return result.as_py()
 
     def __setitem__(self, key, value) -> None:

pandas/tests/extension/test_arrow.py

@@ -513,13 +513,6 @@ def test_reduce_series(self, data, all_numeric_reductions, skipna, request):
         elif all_numeric_reductions == "sem" and pa_version_under8p0:
             request.node.add_marker(xfail_mark)
 
-        elif all_numeric_reductions in [
-            "mean",
-            "median",
-            "std",
-            "sem",
-        ] and pa.types.is_temporal(pa_dtype):
-            request.node.add_marker(xfail_mark)
         elif pa.types.is_boolean(pa_dtype) and all_numeric_reductions in {
             "sem",
             "std",