PERF: setting values via df.loc / df.iloc with pyarrow-backed columns

[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.

Performance improvement in ArrowExtensionArray.__setitem__ when key is a null slice. The rationale for adding a fast path here is that internally null slices get used in a variety of DataFrame setitem operations. Here is an example:

import pandas as pd
import numpy as np

arr = pd.array(np.arange(10**6), dtype="int64[pyarrow]")
df = pd.DataFrame({'a': arr, 'b': arr})

%timeit df.iloc[0, 0] = 0

# 483 ms ± 74.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)    <- main
# 3.4 ms ± 590 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)  <- PR

The example above does not use a null slice, however, internally a null slice is used in ExtensionBlock.set_inplace:

pandas/pandas/core/internals/blocks.py

Line 1646 in 0189674

self.values[:] = values

ASV added:

       before           after         ratio
     [7ef6a71c]       [4e4b36d4]
                      <arrow-ea-setitem-null-slice>
-         802±2μs       57.4±0.5μs     0.07  array.ArrowStringArray.time_setitem_null_slice(False)
-     3.12±0.04ms          194±2μs     0.06  array.ArrowStringArray.time_setitem_null_slice(True)

[mroeschke]

Do we need to directly check not pa_version_under6p0? If a user has data with a pyarrow data type I think it's implicitly assumed that pyarrow is installed?

[lukemanley]

that path hits pc.if_else which I believe was added in 6.0. pc.if_else was the fastest way I found to create an array of all the same value. There are slower options of course.

[mroeschke]

We recently bumped the min pyarrow version to 6.0 so I think this should be safe to remove

[lukemanley]

done

[mroeschke]

1. Is this robust to the case when the value being set is not the same type as the original array e.g. ser = pd.Series([1], dtype="int64"); ser[:] = "foo" should not work
2. Related to the above, was this case before in setitem?
3. Could you add tests for the null slice if/else branches you added?

[lukemanley]

1. Is this robust to the case when the value being set is not the same type as the original array e.g. ser = pd.Series([1], dtype="int64"); ser[:] = "foo" should not work
2. Related to the above, was this case before in setitem?

It is robust in the sense that it raises:

import pandas as pd
ser = pd.Series([1], dtype="int64[pyarrow]")
ser[:] = "foo"

# with this PR:
ArrowInvalid: Could not convert 'foo' with type str: tried to convert to int64

# Previously:
ArrowNotImplementedError: NumPy type not implemented: unrecognized type (19) in GetNumPyTypeName

However, different indexers are raising different errors:.

import pandas as pd
ser = pd.Series([1], dtype="int64[pyarrow]")
ser[0] = "foo"

# ArrowNotImplementedError: NumPy type not implemented: unrecognized type (19) in GetNumPyTypeName

I suspect that refactoring setitem to not go through numpy would help with consistency and might make it cleaner.

3. Could you add tests for the null slice if/else branches you added?

added

[mroeschke]

Okay glad to see that Arrow is erroring in this case. Could you add a test this invalid setting case?

Yeah I think generally ArrowInvalid: Could not convert 'foo' with type str: tried to convert to int64 is a better error message but the refactoring can be left for a separate PR

[lukemanley]

Okay glad to see that Arrow is erroring in this case. Could you add a test this invalid setting case?

Added

[mroeschke]

Thanks @lukemanley