Repository: pandas-dev/pandas
Pull request: BUG: ArrowExtensionArray.mode(dropna=False) not respecting NAs
- [x] closes #50982 (Replace xxxx with the GitHub issue number)
- [x] [Tests added and passed](https://pandas.pydata.org/pandas-docs/dev/development/contributing_codebase.html#writing-tests) if fixing a bug or adding a new feature
- [x] All [code checks passed](https://pandas.pydata.org/pandas-docs/dev/development/contributing_codebase.html#pre-commit).
- [ ] Added [type annotations](https://pandas.pydata.org/pandas-docs/dev/development/contributing_codebase.html#type-hints) to new arguments/methods/functions.
- [x] Added an entry in the latest `doc/source/whatsnew/vX.X.X.rst` file if fixing a bug or adding a new feature.


Problem statement:
BUG: .mode with pyarrow dtypes ignores dropna
```python
ser = pd.Series([1, 1, None, None], dtype="int64[pyarrow]")

ser.mode(dropna=True)
ser.mode(dropna=False)
```

With either dropna=True or dropna=False, we get back `Series([1], dtype="int64[pyarrow]")`.  I'd expect that with dropna=True, but expeced to get `Series([1, None], dtype="int64[pyarrow]")` with dropna=False.

Encountered while trying to fix the mode xfail for string dtype by implementing it in terms of value_counts, which does behave the way I expected.

Hints from issue discussion:

Commit message: BUG: ArrowExtensionArray.mode(dropna=False) not respecting NAs

Patch to review:
diff --git a/doc/source/whatsnew/v2.0.0.rst b/doc/source/whatsnew/v2.0.0.rst
index c28c9fdad1804..075a3f390d01a 100644
--- a/doc/source/whatsnew/v2.0.0.rst
+++ b/doc/source/whatsnew/v2.0.0.rst
@@ -1004,6 +1004,7 @@ Numeric
 - Bug in DataFrame reduction methods (e.g. :meth:`DataFrame.sum`) with object dtype, ``axis=1`` and ``numeric_only=False`` would not be coerced to float (:issue:`49551`)
 - Bug in :meth:`DataFrame.sem` and :meth:`Series.sem` where an erroneous ``TypeError`` would always raise when using data backed by an :class:`ArrowDtype` (:issue:`49759`)
 - Bug in :meth:`Series.__add__` casting to object for list and masked :class:`Series` (:issue:`22962`)
+- Bug in :meth:`~arrays.ArrowExtensionArray.mode` where ``dropna=False`` was not respected when there was ``NA`` values (:issue:`50982`)
 
 Conversion
 ^^^^^^^^^^
diff --git a/pandas/core/arrays/arrow/array.py b/pandas/core/arrays/arrow/array.py
index 0e70b3795bc85..98397fc422131 100644
--- a/pandas/core/arrays/arrow/array.py
+++ b/pandas/core/arrays/arrow/array.py
@@ -1269,7 +1269,6 @@ def _mode(self: ArrowExtensionArrayT, dropna: bool = True) -> ArrowExtensionArra
         ----------
         dropna : bool, default True
             Don't consider counts of NA values.
-            Not implemented by pyarrow.
 
         Returns
         -------
@@ -1292,15 +1291,21 @@ def _mode(self: ArrowExtensionArrayT, dropna: bool = True) -> ArrowExtensionArra
             data = self._data
 
         modes = pc.mode(data, pc.count_distinct(data).as_py())
-        values = modes.field(0)
         counts = modes.field(1)
         # counts sorted descending i.e counts[0] = max
+        if not dropna and self._data.null_count > counts[0].as_py():
+            return type(self)(pa.array([None], type=pa_type))
         mask = pc.equal(counts, counts[0])
-        most_common = values.filter(mask)
+        most_common = modes.field(0).filter(mask)
 
         if pa.types.is_temporal(pa_type):
             most_common = most_common.cast(pa_type)
 
+        if not dropna and self._data.null_count == counts[0].as_py():
+            most_common = pa.concat_arrays(
+                [most_common, pa.array([None], type=pa_type)]
+            )
+
         return type(self)(most_common)
 
     def _maybe_convert_setitem_value(self, value):
diff --git a/pandas/tests/extension/test_arrow.py b/pandas/tests/extension/test_arrow.py
index 9db49470edaf2..5f42c9f3307f7 100644
--- a/pandas/tests/extension/test_arrow.py
+++ b/pandas/tests/extension/test_arrow.py
@@ -1385,6 +1385,27 @@ def test_mode(data_for_grouping, dropna, take_idx, exp_idx, request):
     tm.assert_series_equal(result, expected)
 
 
+def test_mode_dropna_false_mode_na(data, request):
+    # GH 50982
+    pa_dtype = data.dtype.pyarrow_dtype
+    if pa.types.is_string(pa_dtype) or pa.types.is_binary(pa_dtype):
+        request.node.add_marker(
+            pytest.mark.xfail(
+                raises=pa.ArrowNotImplementedError,
+                reason=f"mode not supported by pyarrow for {pa_dtype}",
+            )
+        )
+    more_nans = pd.Series([None, None, data[0]], dtype=data.dtype)
+    result = more_nans.mode(dropna=False)
+    expected = pd.Series([None], dtype=data.dtype)
+    tm.assert_series_equal(result, expected)
+
+    equal_nans = pd.Series([None, data[0]], dtype=data.dtype)
+    result = equal_nans.mode(dropna=False)
+    expected = pd.Series([data[0], None], dtype=data.dtype)
+    tm.assert_series_equal(result, expected)
+
+
 def test_is_bool_dtype():
     # GH 22667
     data = ArrowExtensionArray(pa.array([True, False, True]))


Code review comments:
diff_hunk: @@ -1292,15 +1291,21 @@ def _mode(self: ArrowExtensionArrayT, dropna: bool = True) -> ArrowExtensionArra
             data = self._data
 
         modes = pc.mode(data, pc.count_distinct(data).as_py())
-        values = modes.field(0)
         counts = modes.field(1)
         # counts sorted descending i.e counts[0] = max
+        if not dropna and self._data.null_count > counts[0].as_py():
path: pandas/core/arrays/arrow/array.py
text: @user1:
i know you've wanted to keep as much of this "in pyarrow" as possible, but i find this hard to follow.  would it be that bad to just implement mode in terms of value_counts?

```
if not len(self):[...]
vcs = self.value_counts(dropna=dropna)
res_ser = vcs[vcs == vcs.max()].sort_index()
return res_ser.index._values
```

@author:
Locally with dispatching the mode implementation to value counts results in a noticeable slowdown, so I would prefer to keep using the `pc.mode` implementation

```
In [9]: %timeit arr._mode()
147 µs ± 1.69 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)


In [4]: %timeit arr._mode()
1.13 ms ± 10.1 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
```

@user1:
makes sense, thanks for taking a look

@user2:
Seeing this comment here, I find it surprising that there is such a big difference between both. Especially because you are doing a count_distinct to include all unique values in the result of `mode`, so essentially making it equivalent to a value_counts. So if there is a big difference, that seems to indicate some performance issue in the implementation in arrow. 

Now, trying with the following, I see different results:

```
In [8]: arr = pa.array(np.random.randint(0, 1000, 1_000_000))

In [12]: %timeit pc.mode(arr, pc.count_distinct(arr).as_py())
12.7 ms ± 496 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [13]: %timeit pc.value_counts(arr)
8.96 ms ± 170 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
```
(it might depend a lot on the exact characteristics of the data, though, number of uniques vs total number of values, etc)

@author:
This was a quick benchmark replacing the implementation to use `ArrowExtensionArray.value_counts`, not just comparing `pc.mode` vs `pc.value_counts` per se

```
% ipython
Python 3.8.15 | packaged by conda-forge | (default, Nov 22 2022, 08:53:40)
Type 'copyright', 'credits' or 'license' for more information
IPython 8.8.0 -- An enhanced Interactive Python. Type '?' for help.

In [1]: import pyarrow as pa

In [2]: data = list(range(100_000)) + [None] * 100_000

In [3]: arr = pd.arrays.ArrowExtensionArray(pa.array(data))

In [4]: %timeit arr._mode()
15.9 ms ± 57.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

% git diff
+        vcs = self.value_counts(dropna=dropna)
+        res_ser = vcs[vcs == vcs.max()].sort_index()
+        return res_ser.index._values
+        # pa_type = self._data.type
+        # if pa.types.is_temporal(pa_type):
+        #     nbits = pa_type.bit_width
+        #     if nbits == 32:
+        #         data = self._data.cast(pa.int32())
+        #     elif nbits == 64:
+        #         data = self._data.cast(pa.int64())
+        #     else:
+        #         raise NotImplementedError(pa_type)
+        # else:
+        #     data = self._data
+        #
+        # modes = pc.mode(data, pc.count_distinct(data).as_py())
+        # counts = modes.field(1)
+        # # counts sorted descending i.e counts[0] = max
+        # if not dropna and self._data.null_count > counts[0].as_py():
+        #     return type(self)(pa.array([None], type=pa_type))
+        # mask = pc.equal(counts, counts[0])
+        # most_common = modes.field(0).filter(mask)
+        #
+        # if pa.types.is_temporal(pa_type):
+        #     most_common = most_common.cast(pa_type)
+        #
+        # if not dropna and self._data.null_count == counts[0].as_py():
+        #     most_common = pa.concat_arrays(
+        #         [most_common, pa.array([None], type=pa_type)]
+        #     )
+        #
+        # return type(self)(most_common)

% ipython
Python 3.8.15 | packaged by conda-forge | (default, Nov 22 2022, 08:53:40)
Type 'copyright', 'credits' or 'license' for more information
IPython 8.8.0 -- An enhanced Interactive Python. Type '?' for help.

In [1]: import pyarrow as pa

In [2]: data = list(range(100_000)) + [None] * 100_000

In [3]: arr = pd.arrays.ArrowExtensionArray(pa.array(data))

In [4]: %timeit arr._mode()
54.7 ms ± 672 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
```

@user2:
Then I suppose the overhead is coming from the extra work the `ArrowExtensionArray.value_counts` is doing? 

Doing it with pyarrow compute directly here in `_mode` might be faster and simpler (compared to current `_mode`)? Something like: 

```
res = pc.value_counts(self.data)
most_common = res.field("values").filter(pc.equal(res.field("counts"), pc.max(res.field("counts"))))
```

(this is still faster than calling `pc.mode` on your example data)

@author:
Nice! Looks like `pc.value_counts` also has some benefits

1. Works for string and binary types
2. The result maintains the order of the original values.

So it's good to switch to using value_counts here