Conversation
ngoldbaum
added
01 - Enhancement
component: numpy.strings
| if na_is_nan: | ||
| selection = np.isnan(ar) | ||
| else: | ||
| selection = np.equal(ar, typed_na) |
There was a problem hiding this comment.
@ikrommyd says this causes an issue:
In [8]: import numpy as np
In [9]: from numpy.dtypes import StringDType
In [10]: dt = StringDType(na_object=None)
In [11]: cands = np.array([""] + [f"s{i}" for i in range(30)], dtype="U5")
In [12]: np.isin(np.array([""], dtype=dt), cands)
Out[12]: array([ False])
|
We chatted a little bit about this at a triage review meeting and agreed that it's OK to have a bit of complexity if it buys a really significant speedup like this. |
| if na_is_nan: | ||
| selection = np.isnan(ar) | ||
| else: | ||
| selection = np.equal(ar, typed_na) |
There was a problem hiding this comment.
np.equal(NA, "") returns True which corrupts the results when you are searching if the empty string is in an array. The snippet above returns True on main.
| return ret[:len(ar1)] | ||
| mask = np.ones(len(ar1), dtype=bool) | ||
| for a in ar2: | ||
| mask &= (ar1 != a) |
There was a problem hiding this comment.
there is a problem with invert=True here depending on whether arr2 is small or large
dt = StringDType(na_object=np.nan)
v = np.array(["a", np.nan], dtype=dt)
print(np.isin(v, np.array(["a"], dtype="U1"), invert=True))
# [False False] ← small ar2: old loop
print(np.isin(v, np.array(["a"] + [f"s{i}" for i in range(40)], dtype="U3"), invert=True))
# [False True] ← large ar2: new path|
|
||
| # Unequal descriptors only need a compatibility check when it | ||
| # unlocks sorting; small arrays use the scalar fallback below. | ||
| if ( |
There was a problem hiding this comment.
Again another dependence on the length of arr2 is whether pandas NA sentinels raise.
values = np.array(["a", "b"] * 200, dtype=StringDType(na_object=pd.NA))
cands = np.array(["a", pd.NA], dtype=StringDType(na_object=pd.NA, coerce=False))
np.isin(values, cands) # Error
np.isin(values, np.concatenate([cands, np.array([f"s{i}" for i in range(40)], dtype=cands.dtype)])) # works fine| non_null_ar2, ar2_selection = _filter_stringdtype_nulls( | ||
| ar2, typed_na, na_is_nan | ||
| ) | ||
| non_null_result = _isin_sorting( |
There was a problem hiding this comment.
The decision to use sorting is made before _filter_stringdtype_nulls is called. But filtering can reduce the length. I think it's more efficient for speedup to decide after filtering which method to use.
| for value in ar2: | ||
| # A NaN-like sentinel never matches. Skipping it also avoids | ||
| # scalar object comparison for objects such as pandas.NA. | ||
| if na_is_nan and value is na_object: |
There was a problem hiding this comment.
na_object comes from ar1 here and na_objects are not equal from different arrays.
dt1 = StringDType(na_object=np.nan)
dt2 = StringDType(na_object=float("nan"))
print(dt1 == dt2) # True
print(dt1.na_object is dt2.na_object) #False
|
|
||
| import numpy as np | ||
| from numpy._core.tests._natype import pd_NA | ||
| from numpy._core.tests.test_stringdtype import string_list # noqa: F401 |
There was a problem hiding this comment.
I'm getting linting failures from spin lint locally about redefinition of string_list. Don't know why they're not in CI
| PyObject * | ||
| stringdtype_compatible_na_py( |
There was a problem hiding this comment.
Add NPY_NO_EXPORT to match the header file
| {NULL, NULL, 0, NULL}, | ||
| }; | ||
|
|
||
| PyObject * |
There was a problem hiding this comment.
Add NPY_NO_EXPORT to match the header file
There was a problem hiding this comment.
My AI model informs me that both new exports in this file are reachable from Python already:
np.result_type(StringDType(na_object=None), StringDType(na_object=np.nan)) # TypeError iff incompatible
np.result_type(StringDType(), StringDType(na_object=None)).na_object # effective NAand that
try: np.result_type(...) except TypeErrorreplaces both functions and the getattr/_NoValue chain at _arraysetops_impl.py:855-863
|
|
||
| param_names = ["size", "na_kind"] | ||
| params = [ | ||
| [10, 30_000], |
There was a problem hiding this comment.
size 10 is under the cutoff no? 10 < 10*10**0.145. So it measures the old scalar loop.
|
Closing this in favor of #32693. |
2 participants
PR summary
Fixes #32161. Opening as a draft for now while I test it.
AI Disclosure
I used an AI to iterate on this PR.