PERF: Series.fillna for pyarrow-backed dtypes

[lukemanley]

• Tests added and passed if fixing a bug or adding a new feature
• All code checks passed.
• Added type annotations to new arguments/methods/functions.
• Added an entry in the latest doc/source/whatsnew/v2.0.0.rst file if fixing a bug or adding a new feature.

Perf improvement in Series.fillna for pyarrow-backed dtypes.

      before           after         ratio
    [86f18282]       [34319599]
    <main>           <arrow-ea-fillna>
-        98.6±1ms       18.7±0.4ms     0.19  series_methods.Fillna.time_fillna('string[pyarrow]', 'backfill')
-        99.2±1ms       17.6±0.2ms     0.18  series_methods.Fillna.time_fillna('string[pyarrow]', 'pad')
-       1.86±0.1s      3.86±0.03ms     0.00  series_methods.Fillna.time_fillna('int64[pyarrow]', 'backfill')
-      1.89±0.07s      3.01±0.03ms     0.00  series_methods.Fillna.time_fillna('int64[pyarrow]', 'pad')

SOME BENCHMARKS HAVE CHANGED SIGNIFICANTLY.
PERFORMANCE INCREASED.

[mroeschke]

Could you make sure these are being checked with tm.maybe_produces_warning in the test suite (should show up in the pyarrow=6 build)? Ideally we want to catch all occurrences of these in this test suite

[lukemanley]

added. thanks for the reviews

[mroeschke]

Looking good. Just a single comment

[mroeschke]

A few other places where this is showing up otherwise LGTM

pandas/tests/extension/test_string.py::TestMissing::test_fillna_no_op_returns_copy[True-pyarrow]
pandas/tests/extension/test_string.py::TestMissing::test_fillna_no_op_returns_copy[False-pyarrow]
  /home/runner/work/pandas/pandas/pandas/tests/extension/base/missing.py:93: PerformanceWarning: Falling back on a non-pyarrow code path which may decrease performance. Upgrade to pyarrow >=7 to possibly suppress this warning.
    result = data.fillna(method="backfill")

pandas/tests/extension/test_string.py::TestMissing::test_fillna_series_method[ffill-True-pyarrow]
pandas/tests/extension/test_string.py::TestMissing::test_fillna_series_method[ffill-False-pyarrow]
pandas/tests/extension/test_string.py::TestMissing::test_fillna_series_method[bfill-True-pyarrow]
pandas/tests/extension/test_string.py::TestMissing::test_fillna_series_method[bfill-False-pyarrow]
  /home/runner/work/pandas/pandas/pandas/tests/extension/base/missing.py:123: PerformanceWarning: Falling back on a non-pyarrow code path which may decrease performance. Upgrade to pyarrow >=7 to possibly suppress this warning.
    result = pd.Series(data_missing).fillna(method=fillna_method)

[lukemanley]

A few other places where this is showing up otherwise LGTM

fixed, thx

[mroeschke]

Thanks @lukemanley

[lukemanley]

cc @mroeschke, @jbrockmendel

It turns out Block.fillna does not dispatch to the underlying arrays in all cases including the simple case where method and limit are None. This highlights the impact of not dispatching for Arrow-backed arrays:

import pandas as pd
import pandas._testing as tm

N = 10**6
arr = pd.array(tm.rands_array(10, N), dtype="string[pyarrow]")
arr[::2] = None
ser = pd.Series(arr)

%timeit ser.fillna("foo")
60.5 ms ± 119 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)     # <- no dispatch

%timeit arr.fillna("foo")
4.89 ms ± 22.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)   # <- potential, if dispatched

It looks like the pattern might be to define ArrowBackedExtensionBlock(ExtensionBlock) with a fillna method that dispatches to ArrowExtensionArray.fillna. However, the Block.fillna method is marked final. Any suggestions? Would it be problematic to remove the final attribute?

[mroeschke]

I imagine ExtensionBlock should generally define and just dispatch to to the array's fillna since (at least for the foreseeable future) they will be 1D?

[jbrockmendel]

I imagine ExtensionBlock should generally define and just dispatch to to the array's fillna

That makes sense to me. Would want to make sure we're not hurting perf in non-arrow cases.