PERF: ArrowExtensionArray.searchsorted by lukemanley · Pull Request #50447 · pandas-dev/pandas (original) (raw)

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service andprivacy statement. We’ll occasionally send you account related emails.

Already on GitHub?Sign in to your account

Conversation7 Commits11 Checks0 Files changed

Conversation

This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.Learn more about bidirectional Unicode characters

[ Show hidden characters]({{ revealButtonHref }})

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 ArrowExtensionArray.searchsorted:

import pandas as pd
import numpy as np

arr = pd.array(np.arange(10**6), dtype="int64[pyarrow]")

%timeit arr.searchsorted(500)

# 367 ms ± 2.52 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)         <- main
# 21.3 µs ± 1.18 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)  <- PR

sorter: NumpySorter = None,
) -> npt.NDArray[np.intp] \| np.intp:
if self._hasna:
raise ValueError(

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this the same error the base class would raise? Are there tests for this?

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No, the base class raises:

TypeError: boolean value of NA is ambiguous

The ArrowExtensionArray.searchsorted method added here is pretty much a copy of BaseMaskedArray.searchsorted. I will add a test. (I'll add one for BaseMaskedArray as well - I don't see it being tested there)

"which is impossible with NAs present."
)
with pytest.raises(ValueError, match=err_msg):
print(arr_with_na.dtype)

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you remove this print?

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

oops, removed. thx

@@ -438,6 +438,18 @@ def test_searchsorted(self, data_for_sorting, as_series):
sorter = np.array([1, 2, 0])
assert data_for_sorting.searchsorted(a, sorter=sorter) == 0

# arr containing na value

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think since this is applicable to pandas specific masked arrays + arrow array, this test should live in those specific test files. Technically these base extension tests should be passable for any EA developer who might not define searchsorted to error in this way

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Makes sense, I've moved the tests.

2 participants