Skip to content

ENH: allow StringDType to use the sort paths in isin - #32693

Open
ngoldbaum wants to merge 4 commits into
numpy:mainfrom
ngoldbaum:speedup-isin
Open

ngoldbaum wants to merge 4 commits into
numpy:mainfrom
ngoldbaum:speedup-isin

Conversation

@ngoldbaum

@ngoldbaum ngoldbaum commented Sep 18, 2026

Copy link
Copy Markdown
Member

PR summary

Fixes #32161. Supersedes #32217.

This follows up on recent improvements to StringDType (in particular #32563 and #32464) to enable substantial performance improvements for isin. Also adds more tests for isin in the StringDType tests.

Tests and benchmarks ran on an M1 MacBook Air using Python 3.14. The pandas.NA case uses pandas 3.0.5.

Missing sentinel Before (main) After Speedup
No sentinel 5.8702 s 6.923 ms 848×
None 12.6482 s 5.752 ms 2,199×
np.nan 3.8978 s 5.636 ms 692×
"NA" 8.3059 s 6.129 ms 1,355×
object() 12.7385 s 5.876 ms 2,168×
pandas.NA Raises a casting error 5.919 ms

For pandas.NA, the patch also fixes an error and correctly reports missing entries as non-matches. All returned results were checked against the expected membership mask.

Each baseline was measured once; patched timings are the best of three runs. Input creation is excluded from the timings. Reproduce with repeats=1 on the baseline and repeats=3 on the patched build:

import timeit
import numpy as np
import pandas as pd

repeats = 3
sentinels = (None, np.nan, "NA", object(), pd.NA)
dtypes = [np.dtypes.StringDType()]
dtypes += [np.dtypes.StringDType(na_object=na) for na in sentinels]
for dtype in dtypes:
    values = ["abcd" * 5, "defg" * 5, "abhoj" * 6]
    if hasattr(dtype, "na_object"):
        values[-1] = dtype.na_object
    a = np.array(values * 10_000, dtype=dtype)
    print(dtype, min(timeit.repeat(lambda: np.isin(a, a),
                                  number=1, repeat=repeats)))

AI Disclosure

I used an AI to iterate on this.

@ikrommyd

ikrommyd commented Sep 18, 2026

Copy link
Copy Markdown
Member

Just curious, how does this stack up against the PR you closed? Cause the speedups you're showing here are amazing.

@ngoldbaum

Copy link
Copy Markdown
Member Author

I'm not sure. Probably better because we're relying on C code rather than re-implementing rules for StringDType operations in Python.

@ngoldbaum

Copy link
Copy Markdown
Member Author

I just quickly measured and they're actually about the same. The improvement is mostly algorithmic: the sorting path is O(nlogn) and the slow object path is O(N^2). But thank you for asking me to do this exercise because I see some marginal performance improvements.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

PERF: np.isin is very slow for StringDType

2 participants