Merge remote-tracking branch 'upstream/main' into sroline_issue_60923

Author: Roline Stapny Saldanha

doc/source/whatsnew/v2.3.2.rst

@@ -25,7 +25,7 @@ Bug fixes
 - Fix :meth:`~DataFrame.to_json` with ``orient="table"`` to correctly use the
   "string" type in the JSON Table Schema for :class:`StringDtype` columns
   (:issue:`61889`)
-
+- Boolean operations (``|``, ``&``, ``^``) with bool-dtype objects on the left and :class:`StringDtype` objects on the right now cast the string to bool, with a deprecation warning (:issue:`60234`)
 
 .. ---------------------------------------------------------------------------
 .. _whatsnew_232.contributors:

doc/source/whatsnew/v3.0.0.rst

@@ -854,6 +854,7 @@ Groupby/resample/rolling
 - Bug in :meth:`DataFrame.ewm` and :meth:`Series.ewm` when passed ``times`` and aggregation functions other than mean (:issue:`51695`)
 - Bug in :meth:`DataFrame.resample` and :meth:`Series.resample` were not keeping the index name when the index had :class:`ArrowDtype` timestamp dtype (:issue:`61222`)
 - Bug in :meth:`DataFrame.resample` changing index type to :class:`MultiIndex` when the dataframe is empty and using an upsample method (:issue:`55572`)
+- Bug in :meth:`DataFrameGroupBy.agg` and :meth:`SeriesGroupBy.agg` that was returning numpy dtype values when input values are pyarrow dtype values, instead of returning pyarrow dtype values. (:issue:`53030`)
 - Bug in :meth:`DataFrameGroupBy.agg` that raises ``AttributeError`` when there is dictionary input and duplicated columns, instead of returning a DataFrame with the aggregation of all duplicate columns. (:issue:`55041`)
 - Bug in :meth:`DataFrameGroupBy.agg` where applying a user-defined function to an empty DataFrame returned a Series instead of an empty DataFrame. (:issue:`61503`)
 - Bug in :meth:`DataFrameGroupBy.apply` and :meth:`SeriesGroupBy.apply` for empty data frame with ``group_keys=False`` still creating output index using group keys. (:issue:`60471`)

pandas/core/arrays/arrow/array.py

@@ -12,6 +12,7 @@
     overload,
 )
 import unicodedata
+import warnings
 
 import numpy as np
 
@@ -27,6 +28,7 @@
     pa_version_under13p0,
 )
 from pandas.util._decorators import doc
+from pandas.util._exceptions import find_stack_level
 
 from pandas.core.dtypes.cast import (
     can_hold_element,
@@ -852,6 +854,25 @@ def _logical_method(self, other, op) -> Self:
         # integer types. Otherwise these are boolean ops.
         if pa.types.is_integer(self._pa_array.type):
             return self._evaluate_op_method(other, op, ARROW_BIT_WISE_FUNCS)
+        elif (
+            (
+                pa.types.is_string(self._pa_array.type)
+                or pa.types.is_large_string(self._pa_array.type)
+            )
+            and op in (roperator.ror_, roperator.rand_, roperator.rxor)
+            and isinstance(other, np.ndarray)
+            and other.dtype == bool
+        ):
+            # GH#60234 backward compatibility for the move to StringDtype in 3.0
+            op_name = op.__name__[1:].strip("_")
+            warnings.warn(
+                f"'{op_name}' operations between boolean dtype and {self.dtype} are "
+                "deprecated and will raise in a future version. Explicitly "
+                "cast the strings to a boolean dtype before operating instead.",
+                FutureWarning,
+                stacklevel=find_stack_level(),
+            )
+            return op(other, self.astype(bool))
         else:
             return self._evaluate_op_method(other, op, ARROW_LOGICAL_FUNCS)
 

pandas/core/arrays/string_.py

@@ -52,6 +52,7 @@
     missing,
     nanops,
     ops,
+    roperator,
 )
 from pandas.core.algorithms import isin
 from pandas.core.array_algos import masked_reductions
@@ -390,6 +391,26 @@ class BaseStringArray(ExtensionArray):
 
     dtype: StringDtype
 
+    # TODO(4.0): Once the deprecation here is enforced, this method can be
+    #  removed and we use the parent class method instead.
+    def _logical_method(self, other, op):
+        if (
+            op in (roperator.ror_, roperator.rand_, roperator.rxor)
+            and isinstance(other, np.ndarray)
+            and other.dtype == bool
+        ):
+            # GH#60234 backward compatibility for the move to StringDtype in 3.0
+            op_name = op.__name__[1:].strip("_")
+            warnings.warn(
+                f"'{op_name}' operations between boolean dtype and {self.dtype} are "
+                "deprecated and will raise in a future version. Explicitly "
+                "cast the strings to a boolean dtype before operating instead.",
+                FutureWarning,
+                stacklevel=find_stack_level(),
+            )
+            return op(other, self.astype(bool))
+        return NotImplemented
+
     @doc(ExtensionArray.tolist)
     def tolist(self) -> list:
         if self.ndim > 1:

pandas/core/groupby/ops.py

@@ -44,13 +44,15 @@
     ensure_platform_int,
     ensure_uint64,
     is_1d_only_ea_dtype,
+    is_string_dtype,
 )
 from pandas.core.dtypes.missing import (
     isna,
     maybe_fill,
 )
 
 from pandas.core.arrays import Categorical
+from pandas.core.arrays.arrow.array import ArrowExtensionArray
 from pandas.core.frame import DataFrame
 from pandas.core.groupby import grouper
 from pandas.core.indexes.api import (
@@ -963,18 +965,26 @@ def agg_series(
         -------
         np.ndarray or ExtensionArray
         """
+        result = self._aggregate_series_pure_python(obj, func)
+        npvalues = lib.maybe_convert_objects(result, try_float=False)
+
+        if isinstance(obj._values, ArrowExtensionArray):
+            # When obj.dtype is a string, any object can be cast. Only do so if the
+            # UDF returned strings or NA values.
+            if not is_string_dtype(obj.dtype) or lib.is_string_array(
+                npvalues, skipna=True
+            ):
+                out = maybe_cast_pointwise_result(
+                    npvalues, obj.dtype, numeric_only=True, same_dtype=preserve_dtype
+                )
+            else:
+                out = npvalues
 
-        if not isinstance(obj._values, np.ndarray):
+        elif not isinstance(obj._values, np.ndarray):
             # we can preserve a little bit more aggressively with EA dtype
             #  because maybe_cast_pointwise_result will do a try/except
             #  with _from_sequence.  NB we are assuming here that _from_sequence
             #  is sufficiently strict that it casts appropriately.
-            preserve_dtype = True
-
-        result = self._aggregate_series_pure_python(obj, func)
-
-        npvalues = lib.maybe_convert_objects(result, try_float=False)
-        if preserve_dtype:
             out = maybe_cast_pointwise_result(npvalues, obj.dtype, numeric_only=True)
         else:
             out = npvalues

pandas/tests/groupby/aggregate/test_aggregate.py

@@ -10,6 +10,7 @@
 import pytest
 
 from pandas.errors import SpecificationError
+import pandas.util._test_decorators as td
 
 from pandas.core.dtypes.common import is_integer_dtype
 
@@ -23,6 +24,7 @@
     to_datetime,
 )
 import pandas._testing as tm
+from pandas.arrays import ArrowExtensionArray
 from pandas.core.groupby.grouper import Grouping
 
 
@@ -1809,6 +1811,102 @@ def test_groupby_aggregation_func_list_multi_index_duplicate_columns():
     tm.assert_frame_equal(result, expected)
 
 
+@td.skip_if_no("pyarrow")
+@pytest.mark.parametrize(
+    "input_dtype, output_dtype",
+    [
+        # With NumPy arrays, the results from the UDF would be e.g. np.float32 scalars
+        # which we can therefore preserve. However with PyArrow arrays, the results are
+        # Python scalars so we have no information about size or uint vs int.
+        ("float[pyarrow]", "double[pyarrow]"),
+        ("int64[pyarrow]", "int64[pyarrow]"),
+        ("uint64[pyarrow]", "int64[pyarrow]"),
+        ("bool[pyarrow]", "bool[pyarrow]"),
+    ],
+)
+def test_agg_lambda_pyarrow_dtype_conversion(input_dtype, output_dtype):
+    # GH#59601
+    # Test PyArrow dtype conversion back to PyArrow dtype
+    df = DataFrame(
+        {
+            "A": ["c1", "c2", "c3", "c1", "c2", "c3"],
+            "B": pd.array([100, 200, 255, 0, 199, 40392], dtype=input_dtype),
+        }
+    )
+    gb = df.groupby("A")
+    result = gb.agg(lambda x: x.min())
+
+    expected = DataFrame(
+        {"B": pd.array([0, 199, 255], dtype=output_dtype)},
+        index=Index(["c1", "c2", "c3"], name="A"),
+    )
+    tm.assert_frame_equal(result, expected)
+
+
+@td.skip_if_no("pyarrow")
+def test_agg_lambda_complex128_dtype_conversion():
+    # GH#59601
+    df = DataFrame(
+        {"A": ["c1", "c2", "c3"], "B": pd.array([100, 200, 255], "int64[pyarrow]")}
+    )
+    gb = df.groupby("A")
+    result = gb.agg(lambda x: complex(x.sum(), x.count()))
+
+    expected = DataFrame(
+        {
+            "B": pd.array(
+                [complex(100, 1), complex(200, 1), complex(255, 1)], dtype="complex128"
+            ),
+        },
+        index=Index(["c1", "c2", "c3"], name="A"),
+    )
+    tm.assert_frame_equal(result, expected)
+
+
+@td.skip_if_no("pyarrow")
+def test_agg_lambda_numpy_uint64_to_pyarrow_dtype_conversion():
+    # GH#59601
+    df = DataFrame(
+        {
+            "A": ["c1", "c2", "c3"],
+            "B": pd.array([100, 200, 255], dtype="uint64[pyarrow]"),
+        }
+    )
+    gb = df.groupby("A")
+    result = gb.agg(lambda x: np.uint64(x.sum()))
+
+    expected = DataFrame(
+        {
+            "B": pd.array([100, 200, 255], dtype="uint64[pyarrow]"),
+        },
+        index=Index(["c1", "c2", "c3"], name="A"),
+    )
+    tm.assert_frame_equal(result, expected)
+
+
+@td.skip_if_no("pyarrow")
+def test_agg_lambda_pyarrow_struct_to_object_dtype_conversion():
+    # GH#59601
+    import pyarrow as pa
+
+    df = DataFrame(
+        {
+            "A": ["c1", "c2", "c3"],
+            "B": pd.array([100, 200, 255], dtype="int64[pyarrow]"),
+        }
+    )
+    gb = df.groupby("A")
+    result = gb.agg(lambda x: {"number": 1})
+
+    arr = pa.array([{"number": 1}, {"number": 1}, {"number": 1}])
+    expected = DataFrame(
+        {"B": ArrowExtensionArray(arr)},
+        index=Index(["c1", "c2", "c3"], name="A"),
+    )
+
+    tm.assert_frame_equal(result, expected)
+
+
 def test_groupby_aggregate_empty_builtin_sum():
     df = DataFrame(columns=["Group", "Data"])
     result = df.groupby(["Group"], as_index=False)["Data"].agg("sum")

pandas/tests/groupby/test_groupby.py

@@ -2434,25 +2434,28 @@ def test_rolling_wrong_param_min_period():
 
 def test_by_column_values_with_same_starting_value(any_string_dtype):
     # GH29635
+    dtype = any_string_dtype
     df = DataFrame(
         {
             "Name": ["Thomas", "Thomas", "Thomas John"],
             "Credit": [1200, 1300, 900],
-            "Mood": Series(["sad", "happy", "happy"], dtype=any_string_dtype),
+            "Mood": Series(["sad", "happy", "happy"], dtype=dtype),
         }
     )
     aggregate_details = {"Mood": Series.mode, "Credit": "sum"}
 
     result = df.groupby(["Name"]).agg(aggregate_details)
-    expected_result = DataFrame(
+    expected = DataFrame(
         {
             "Mood": [["happy", "sad"], "happy"],
             "Credit": [2500, 900],
             "Name": ["Thomas", "Thomas John"],
-        }
+        },
     ).set_index("Name")
-
-    tm.assert_frame_equal(result, expected_result)
+    if getattr(dtype, "storage", None) == "pyarrow":
+        mood_values = pd.array(["happy", "sad"], dtype=dtype)
+        expected["Mood"] = [mood_values, "happy"]
+    tm.assert_frame_equal(result, expected)
 
 
 def test_groupby_none_in_first_mi_level():

pandas/tests/strings/test_strings.py

@@ -787,3 +787,27 @@ def test_decode_with_dtype_none():
         result = ser.str.decode("utf-8", dtype=None)
         expected = Series(["a", "b", "c"], dtype="str")
         tm.assert_series_equal(result, expected)
+
+
+def test_reversed_logical_ops(any_string_dtype):
+    # GH#60234
+    dtype = any_string_dtype
+    warn = None if dtype == object else FutureWarning
+    left = Series([True, False, False, True])
+    right = Series(["", "", "b", "c"], dtype=dtype)
+
+    msg = "operations between boolean dtype and"
+    with tm.assert_produces_warning(warn, match=msg):
+        result = left | right
+    expected = left | right.astype(bool)
+    tm.assert_series_equal(result, expected)
+
+    with tm.assert_produces_warning(warn, match=msg):
+        result = left & right
+    expected = left & right.astype(bool)
+    tm.assert_series_equal(result, expected)
+
+    with tm.assert_produces_warning(warn, match=msg):
+        result = left ^ right
+    expected = left ^ right.astype(bool)
+    tm.assert_series_equal(result, expected)