fix remaining tests

Author: jbrockmendel

pandas/core/arrays/arrow/array.py

@@ -396,11 +396,66 @@ def _cast_pointwise_result(self, values) -> ArrayLike:
         if len(values) == 0:
             # Retain our dtype
             return self[:0].copy()
-        arr = pa.array(values, from_pandas=True)
+
+        try:
+            arr = pa.array(values, from_pandas=True)
+        except (ValueError, TypeError):
+            # e.g. test_by_column_values_with_same_starting_value with nested
+            #  values, one entry of which is an ArrowStringArray
+            #  or test_agg_lambda_complex128_dtype_conversion for complex values
+            return super()._cast_pointwise_result(values)
+
+        if pa.types.is_duration(arr.type):
+            # workaround for https://github.com/apache/arrow/issues/40620
+            result = ArrowExtensionArray._from_sequence(values)
+            if pa.types.is_duration(self._pa_array.type):
+                result = result.astype(self.dtype)
+            elif pa.types.is_timestamp(self._pa_array.type):
+                # Try to retain original unit
+                new_dtype = ArrowDtype(pa.duration(self._pa_array.type.unit))
+                try:
+                    result = result.astype(new_dtype)
+                except ValueError:
+                    pass
+            elif pa.types.is_date64(self._pa_array.type):
+                # Try to match unit we get on non-pointwise op
+                dtype = ArrowDtype(pa.duration("ms"))
+                result = result.astype(dtype)
+            elif pa.types.is_date(self._pa_array.type):
+                # Try to match unit we get on non-pointwise op
+                dtype = ArrowDtype(pa.duration("s"))
+                result = result.astype(dtype)
+            return result
+
+        elif pa.types.is_date(arr.type) and pa.types.is_date(self._pa_array.type):
+            arr = arr.cast(self._pa_array.type)
+        elif pa.types.is_time(arr.type) and pa.types.is_time(self._pa_array.type):
+            arr = arr.cast(self._pa_array.type)
+        elif pa.types.is_decimal(arr.type) and pa.types.is_decimal(self._pa_array.type):
+            arr = arr.cast(self._pa_array.type)
+        elif pa.types.is_integer(arr.type) and pa.types.is_integer(self._pa_array.type):
+            try:
+                arr = arr.cast(self._pa_array.type)
+            except pa.lib.ArrowInvalid:
+                # e.g. test_combine_add if we can't cast
+                pass
+        elif pa.types.is_floating(arr.type) and pa.types.is_floating(
+            self._pa_array.type
+        ):
+            try:
+                arr = arr.cast(self._pa_array.type)
+            except pa.lib.ArrowInvalid:
+                # e.g. test_combine_add if we can't cast
+                pass
+
         if isinstance(self.dtype, StringDtype):
             if pa.types.is_string(arr.type) or pa.types.is_large_string(arr.type):
                 # ArrowStringArrayNumpySemantics
-                return type(self)(arr)
+                return type(self)(arr).astype(self.dtype)
+            if self.dtype.na_value is np.nan:
+                # ArrowEA has different semantics, so we return numpy-based
+                #  result instead
+                return super()._cast_pointwise_result(values)
             return ArrowExtensionArray(arr)
         return type(self)(arr)
 

pandas/core/arrays/masked.py

@@ -26,6 +26,7 @@
 from pandas.util._decorators import doc
 
 from pandas.core.dtypes.base import ExtensionDtype
+from pandas.core.dtypes.cast import maybe_downcast_to_dtype
 from pandas.core.dtypes.common import (
     is_bool,
     is_integer_dtype,
@@ -149,7 +150,15 @@ def _from_sequence(cls, scalars, *, dtype=None, copy: bool = False) -> Self:
 
     def _cast_pointwise_result(self, values) -> ArrayLike:
         values = np.asarray(values, dtype=object)
-        return lib.maybe_convert_objects(values, convert_to_nullable_dtype=True)
+        result = lib.maybe_convert_objects(values, convert_to_nullable_dtype=True)
+        lkind = self.dtype.kind
+        rkind = result.dtype.kind
+        if (lkind in "iu" and rkind in "iu") or (lkind == rkind == "f"):
+            new_data = maybe_downcast_to_dtype(
+                result._data, dtype=self.dtype.numpy_dtype
+            )
+            result = type(result)(new_data, result._mask)
+        return result
 
     @classmethod
     @doc(ExtensionArray._empty)

pandas/core/arrays/numpy_.py

@@ -14,7 +14,10 @@
 from pandas.compat.numpy import function as nv
 
 from pandas.core.dtypes.astype import astype_array
-from pandas.core.dtypes.cast import construct_1d_object_array_from_listlike
+from pandas.core.dtypes.cast import (
+    construct_1d_object_array_from_listlike,
+    maybe_downcast_to_dtype,
+)
 from pandas.core.dtypes.common import pandas_dtype
 from pandas.core.dtypes.dtypes import NumpyEADtype
 from pandas.core.dtypes.missing import isna
@@ -34,6 +37,7 @@
     from collections.abc import Callable
 
     from pandas._typing import (
+        ArrayLike,
         AxisInt,
         Dtype,
         FillnaOptions,
@@ -145,6 +149,24 @@ def _from_sequence(
             result = result.copy()
         return cls(result)
 
+    def _cast_pointwise_result(self, values) -> ArrayLike:
+        result = super()._cast_pointwise_result(values)
+        lkind = self.dtype.kind
+        rkind = result.dtype.kind
+        if (
+            (lkind in "iu" and rkind in "iu")
+            or (lkind == "f" and rkind == "f")
+            or (lkind == rkind == "c")
+        ):
+            result = maybe_downcast_to_dtype(result, self.dtype.numpy_dtype)
+        elif rkind == "M":
+            # Ensure potential subsequent .astype(object) doesn't incorrectly
+            #  convert Timestamps to ints
+            from pandas import array as pd_array
+
+            result = pd_array(result, copy=False)
+        return result
+
     # ------------------------------------------------------------------------
     # Data
 

pandas/core/arrays/sparse/array.py

@@ -609,7 +609,20 @@ def _from_factorized(cls, values, original) -> Self:
 
     def _cast_pointwise_result(self, values):
         result = super()._cast_pointwise_result(values)
-        return type(self)._from_sequence(result)
+        if result.dtype.kind == self.dtype.kind:
+            try:
+                # e.g. test_groupby_agg_extension
+                res = type(self)._from_sequence(result, dtype=self.dtype)
+                if ((res == result) | (isna(result) & res.isna())).all():
+                    # This does not hold for e.g.
+                    #  test_arith_frame_with_scalar[0-__truediv__]
+                    return res
+                return type(self)._from_sequence(result)
+            except (ValueError, TypeError):
+                return type(self)._from_sequence(result)
+        else:
+            # e.g. test_combine_le avoid casting bools to Sparse[float64, nan]
+            return type(self)._from_sequence(result)
 
     # ------------------------------------------------------------------------
     # Data

pandas/core/indexes/base.py

@@ -6397,18 +6397,20 @@ def map(self, mapper, na_action: Literal["ignore"] | None = None):
         if not new_values.size:
             # empty
             dtype = self.dtype
-            return Index(new_values, dtype=dtype, copy=False, name=self.name)
+        elif isinstance(new_values, Categorical):
+            # cast_pointwise_result is unnecessary
+            dtype = new_values.dtype
         else:
+            if isinstance(self, MultiIndex):
+                arr = self[:0].to_flat_index().array
+            else:
+                arr = self[:0].array
             # e.g. if we are floating and new_values is all ints, then we
             #  don't want to cast back to floating.  But if we are UInt64
             #  and new_values is all ints, we want to try.
-            if isinstance(self._values, np.ndarray):
-                return Index(new_values, dtype=dtype, copy=False, name=self.name)
-            else:
-                new_values = self._values._cast_pointwise_result(new_values)
-                return Index(
-                    new_values, dtype=new_values.dtype, copy=False, name=self.name
-                )
+            new_values = arr._cast_pointwise_result(new_values)
+            dtype = new_values.dtype
+        return Index(new_values, dtype=dtype, copy=False, name=self.name)
 
     # TODO: De-duplicate with map, xref GH#32349
     @final

pandas/tests/extension/base/methods.py

@@ -367,6 +367,18 @@ def test_combine_le(self, data_repeated):
         )
         tm.assert_series_equal(result, expected)
 
+    def _construct_for_combine_add(self, left, right):
+        if isinstance(right, type(left)):
+            return left._from_sequence(
+                [a + b for (a, b) in zip(list(left), list(right))],
+                dtype=left.dtype,
+            )
+        else:
+            return left._from_sequence(
+                [a + right for a in list(left)],
+                dtype=left.dtype,
+            )
+
     def test_combine_add(self, data_repeated):
         # GH 20825
         orig_data1, orig_data2 = data_repeated(2)
@@ -377,26 +389,22 @@ def test_combine_add(self, data_repeated):
         #  we will expect Series.combine to raise as well.
         try:
             with np.errstate(over="ignore"):
-                expected = pd.Series(
-                    orig_data1._from_sequence(
-                        [a + b for (a, b) in zip(list(orig_data1), list(orig_data2))]
-                    )
-                )
+                arr = self._construct_for_combine_add(orig_data1, orig_data2)
         except TypeError:
             # If the operation is not supported pointwise for our scalars,
             #  then Series.combine should also raise
             with pytest.raises(TypeError):
                 s1.combine(s2, lambda x1, x2: x1 + x2)
             return
+        expected = pd.Series(arr)
 
         result = s1.combine(s2, lambda x1, x2: x1 + x2)
         tm.assert_series_equal(result, expected)
 
         val = s1.iloc[0]
         result = s1.combine(val, lambda x1, x2: x1 + x2)
-        expected = pd.Series(
-            orig_data1._from_sequence([a + val for a in list(orig_data1)])
-        )
+        arr = self._construct_for_combine_add(orig_data1, val)
+        expected = pd.Series(arr)
         tm.assert_series_equal(result, expected)
 
     def test_combine_first(self, data):

pandas/tests/extension/decimal/test_decimal.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import decimal
-import operator
 
 import numpy as np
 import pytest
@@ -282,33 +281,10 @@ def _create_arithmetic_method(cls, op):
 DecimalArrayWithoutCoercion._add_arithmetic_ops()
 
 
-def test_combine_from_sequence_raises(monkeypatch):
-    # https://github.com/pandas-dev/pandas/issues/22850
-    cls = DecimalArrayWithoutFromSequence
-
-    def construct_array_type(self):
-        return DecimalArrayWithoutFromSequence
-
-    monkeypatch.setattr(DecimalDtype, "construct_array_type", construct_array_type)
-
-    arr = cls([decimal.Decimal("1.0"), decimal.Decimal("2.0")])
-    ser = pd.Series(arr)
-    result = ser.combine(ser, operator.add)
-
-    # note: object dtype
-    expected = pd.Series(
-        [decimal.Decimal("2.0"), decimal.Decimal("4.0")], dtype="object"
-    )
-    tm.assert_series_equal(result, expected)
-
-
-@pytest.mark.parametrize(
-    "class_", [DecimalArrayWithoutFromSequence, DecimalArrayWithoutCoercion]
-)
-def test_scalar_ops_from_sequence_raises(class_):
+def test_scalar_ops_from_sequence_raises():
     # op(EA, EA) should return an EA, or an ndarray if it's not possible
     # to return an EA with the return values.
-    arr = class_([decimal.Decimal("1.0"), decimal.Decimal("2.0")])
+    arr = DecimalArrayWithoutCoercion([decimal.Decimal("1.0"), decimal.Decimal("2.0")])
     result = arr + arr
     expected = np.array(
         [decimal.Decimal("2.0"), decimal.Decimal("4.0")], dtype="object"

pandas/tests/extension/json/array.py

@@ -90,6 +90,13 @@ def _from_sequence(cls, scalars, *, dtype=None, copy=False):
     def _from_factorized(cls, values, original):
         return cls([UserDict(x) for x in values if x != ()])
 
+    def _cast_pointwise_result(self, values):
+        result = super()._cast_pointwise_result(values)
+        try:
+            return type(self)._from_sequence(result, dtype=self.dtype)
+        except (ValueError, TypeError):
+            return result
+
     def __getitem__(self, item):
         if isinstance(item, tuple):
             item = unpack_tuple_and_ellipses(item)

pandas/tests/extension/test_arrow.py

@@ -46,6 +46,7 @@
 )
 from pandas.errors import Pandas4Warning
 
+from pandas.core.dtypes.common import pandas_dtype
 from pandas.core.dtypes.dtypes import (
     ArrowDtype,
     CategoricalDtypeType,
@@ -271,6 +272,26 @@ def data_for_twos(data):
 
 
 class TestArrowArray(base.ExtensionTests):
+    def _construct_for_combine_add(self, left, right):
+        dtype = left.dtype
+
+        # in a couple cases, addition is not dtype-preserving
+        if dtype == "bool[pyarrow]":
+            dtype = pandas_dtype("int64[pyarrow]")
+        elif dtype == "int8[pyarrow]" and isinstance(right, type(left)):
+            dtype = pandas_dtype("int64[pyarrow]")
+
+        if isinstance(right, type(left)):
+            return left._from_sequence(
+                [a + b for (a, b) in zip(list(left), list(right))],
+                dtype=dtype,
+            )
+        else:
+            return left._from_sequence(
+                [a + right for a in list(left)],
+                dtype=dtype,
+            )
+
     def test_compare_scalar(self, data, comparison_op):
         ser = pd.Series(data)
         self._compare_other(ser, data, comparison_op, data[0])
@@ -797,14 +818,24 @@ def _cast_pointwise_result(self, op_name: str, obj, other, pointwise_result):
         if op_name in ["eq", "ne", "lt", "le", "gt", "ge"]:
             return pointwise_result.astype("boolean[pyarrow]")
 
+        original_dtype = tm.get_dtype(expected)
+
         was_frame = False
         if isinstance(expected, pd.DataFrame):
             was_frame = True
             expected_data = expected.iloc[:, 0]
-            original_dtype = obj.iloc[:, 0].dtype
         else:
             expected_data = expected
-            original_dtype = obj.dtype
+
+        # the pointwise method will have retained our original dtype, while
+        #  the op(ser, other) version will have cast to 64bit
+        if type(other) is int and op_name not in ["__floordiv__"]:
+            if original_dtype.kind == "f":
+                return expected.astype("float64[pyarrow]")
+            else:
+                return expected.astype("int64[pyarrow]")
+        elif type(other) is float:
+            return expected.astype("float64[pyarrow]")
 
         orig_pa_type = original_dtype.pyarrow_dtype
         if not was_frame and isinstance(other, pd.Series):
@@ -836,29 +867,7 @@ def _cast_pointwise_result(self, op_name: str, obj, other, pointwise_result):
 
         pa_expected = pa.array(expected_data._values)
 
-        if pa.types.is_duration(pa_expected.type):
-            if pa.types.is_date(orig_pa_type):
-                if pa.types.is_date64(orig_pa_type):
-                    # TODO: why is this different vs date32?
-                    unit = "ms"
-                else:
-                    unit = "s"
-            else:
-                # pyarrow sees sequence of datetime/timedelta objects and defaults
-                #  to "us" but the non-pointwise op retains unit
-                # timestamp or duration
-                unit = orig_pa_type.unit
-                if type(other) in [datetime, timedelta] and unit in ["s", "ms"]:
-                    # pydatetime/pytimedelta objects have microsecond reso, so we
-                    #  take the higher reso of the original and microsecond. Note
-                    #  this matches what we would do with DatetimeArray/TimedeltaArray
-                    unit = "us"
-
-            pa_expected = pa_expected.cast(f"duration[{unit}]")
-
-        elif pa.types.is_decimal(pa_expected.type) and pa.types.is_decimal(
-            orig_pa_type
-        ):
+        if pa.types.is_decimal(pa_expected.type) and pa.types.is_decimal(orig_pa_type):
             # decimal precision can resize in the result type depending on data
             # just compare the float values
             alt = getattr(obj, op_name)(other)