From 003ea31af87928ac019ff029f58d6050ce524c54 Mon Sep 17 00:00:00 2001 From: Brock Date: Sun, 12 Apr 2026 08:08:17 -0700 Subject: [PATCH 1/4] PERF: short-circuit array_equivalent and equals comparisons Add short-circuiting Cython functions for array equality checks: - array_equivalent_float: single-pass NaN-aware float comparison, replacing the 4-temporary-array expression ((left == right) | (isnan(left) & isnan(right))).all() - array_equivalent_bytes: memcmp-based comparison for int/bool/datetime arrays, replacing np.array_equal - has_nans/all_nans: short-circuiting replacements for np.isnan(arr).any()/all() For non-contiguous inputs, falls back to the original numpy expressions with no regression. closes #32339 Co-Authored-By: Claude Opus 4.6 (1M context) --- pandas/_libs/lib.pyx | 98 +++++++++++++++++++++++++++++ pandas/core/arrays/categorical.py | 2 +- pandas/core/arrays/datetimelike.py | 2 +- pandas/core/arrays/masked.py | 2 +- pandas/core/arrays/period.py | 2 +- pandas/core/dtypes/missing.py | 23 ++++++- pandas/core/indexes/datetimelike.py | 4 +- pandas/core/indexes/multi.py | 4 +- 8 files changed, 126 insertions(+), 11 deletions(-) diff --git a/pandas/_libs/lib.pyx b/pandas/_libs/lib.pyx index 8df146041944e..7417fdaeb0d02 100644 --- a/pandas/_libs/lib.pyx +++ b/pandas/_libs/lib.pyx @@ -37,6 +37,7 @@ from cython cimport ( Py_ssize_t, floating, ) +from libc.string cimport memcmp from pandas._config import using_string_dtype @@ -497,6 +498,103 @@ def has_infs(const floating[:] arr) -> bool: return ret +@cython.wraparound(False) +@cython.boundscheck(False) +def has_nans(const floating[:] arr) -> bool: + cdef: + Py_ssize_t i, n = len(arr) + Py_ssize_t n4 = n & ~3 # round down to multiple of 4 + bint found = False + + with nogil: + for i in range(0, n4, 4): + if ( + (arr[i] != arr[i]) + | (arr[i + 1] != arr[i + 1]) + | (arr[i + 2] != arr[i + 2]) + | (arr[i + 3] != arr[i + 3]) + ): + found = True + break + if not found: + for i in range(n4, n): + if arr[i] != arr[i]: + found = True + break + return found + + +@cython.wraparound(False) +@cython.boundscheck(False) +def all_nans(const floating[:] arr) -> bool: + cdef: + Py_ssize_t i, n = len(arr) + Py_ssize_t n4 = n & ~3 + bint found_non_nan = False + + with nogil: + for i in range(0, n4, 4): + if ( + (arr[i] == arr[i]) + | (arr[i + 1] == arr[i + 1]) + | (arr[i + 2] == arr[i + 2]) + | (arr[i + 3] == arr[i + 3]) + ): + found_non_nan = True + break + if not found_non_nan: + for i in range(n4, n): + if arr[i] == arr[i]: + found_non_nan = True + break + return not found_non_nan + + +@cython.wraparound(False) +@cython.boundscheck(False) +def array_equivalent_float(const floating[:] left, + const floating[:] right) -> bool: + cdef: + Py_ssize_t i, n = len(left) + floating lval, rval + bint mismatch = False + + with nogil: + for i in range(n): + lval = left[i] + rval = right[i] + if lval != rval: + if not (lval != lval and rval != rval): + mismatch = True + break + return not mismatch + + +def array_equivalent_bytes(left, right) -> bool: + cdef: + Py_ssize_t nbytes + int ndim, idx + ndarray left_arr, right_arr + + left_arr = np.asarray(left) + right_arr = np.asarray(right) + + ndim = cnp.PyArray_NDIM(left_arr) + if ndim != cnp.PyArray_NDIM(right_arr): + return False + for idx in range(ndim): + if cnp.PyArray_DIM(left_arr, idx) != cnp.PyArray_DIM(right_arr, idx): + return False + if not (cnp.PyArray_IS_C_CONTIGUOUS(left_arr) + and cnp.PyArray_IS_C_CONTIGUOUS(right_arr)): + return np.array_equal(left_arr, right_arr) + nbytes = cnp.PyArray_NBYTES(left_arr) + if nbytes == 0: + return True + return memcmp(cnp.PyArray_DATA(left_arr), cnp.PyArray_DATA(right_arr), + nbytes) == 0 + + @cython.boundscheck(False) @cython.wraparound(False) def has_only_ints_or_nan(const floating[:] arr) -> bool: diff --git a/pandas/core/arrays/categorical.py b/pandas/core/arrays/categorical.py index 83fe0965ec123..86e8571d4a00f 100644 --- a/pandas/core/arrays/categorical.py +++ b/pandas/core/arrays/categorical.py @@ -2626,7 +2626,7 @@ def equals(self, other: object) -> bool: return False elif self._categories_match_up_to_permutation(other): other = self._encode_with_my_categories(other) - return np.array_equal(self._codes, other._codes) + return lib.array_equivalent_bytes(self._codes, other._codes) return False def _accumulate(self, name: str, skipna: bool = True, **kwargs) -> Self: diff --git a/pandas/core/arrays/datetimelike.py b/pandas/core/arrays/datetimelike.py index da1ada2dffdd7..27ccac5bbf976 100644 --- a/pandas/core/arrays/datetimelike.py +++ b/pandas/core/arrays/datetimelike.py @@ -1946,7 +1946,7 @@ def _validate_frequency(cls, index, freq: BaseOffset, **kwargs) -> None: unit=index.unit, **kwargs, ) - if not np.array_equal(index.asi8, on_freq.asi8): + if not lib.array_equivalent_bytes(index.asi8, on_freq.asi8): raise ValueError except ValueError as err: if "non-fixed" in str(err): diff --git a/pandas/core/arrays/masked.py b/pandas/core/arrays/masked.py index dc4f55ebf086b..89d92c66fff27 100644 --- a/pandas/core/arrays/masked.py +++ b/pandas/core/arrays/masked.py @@ -1583,7 +1583,7 @@ def equals(self, other) -> bool: # GH#44382 if e.g. self[1] is np.nan and other[1] is pd.NA, we are NOT # equal. - if not np.array_equal(self._mask, other._mask): + if not lib.array_equivalent_bytes(self._mask, other._mask): return False left = self._data[~self._mask] diff --git a/pandas/core/arrays/period.py b/pandas/core/arrays/period.py index c9e0f47bbe852..ba7b6bf56f526 100644 --- a/pandas/core/arrays/period.py +++ b/pandas/core/arrays/period.py @@ -295,7 +295,7 @@ def _from_sequence( arrdata = np.asarray(scalars) if arrdata.dtype.kind == "f" and len(arrdata) > 0: - if not np.isnan(arrdata).all(): + if not lib.all_nans(arrdata): raise TypeError( "PeriodArray does not allow floating point in construction" ) diff --git a/pandas/core/dtypes/missing.py b/pandas/core/dtypes/missing.py index cb9d97d90da87..90143800cfb7b 100644 --- a/pandas/core/dtypes/missing.py +++ b/pandas/core/dtypes/missing.py @@ -447,7 +447,7 @@ def array_equivalent( # TODO: fastpath for pandas' StringDtype return _array_equivalent_object(left, right, strict_nan) else: - return np.array_equal(left, right) + return lib.array_equivalent_bytes(left, right) # Slow path when we allow comparing different dtypes. # Object arrays can contain None, NaN and NaT. @@ -477,15 +477,32 @@ def array_equivalent( ) and left.dtype != right.dtype: return False + if left.dtype == right.dtype: + return lib.array_equivalent_bytes(left, right) return np.array_equal(left, right) def _array_equivalent_float(left: np.ndarray, right: np.ndarray) -> bool: - return bool(((left == right) | (np.isnan(left) & np.isnan(right))).all()) + if left.dtype.kind == "c": + if not (left.flags.c_contiguous and right.flags.c_contiguous): + return bool(((left == right) | (np.isnan(left) & np.isnan(right))).all()) + # View complex as float pairs (complex128 -> float64, complex64 -> float32) + left = left.view(left.real.dtype) + right = right.view(right.real.dtype) + if left.ndim > 1: + if left.flags.f_contiguous and right.flags.f_contiguous: + # .T is a C-contiguous view of an F-contiguous array + left = left.T + right = right.T + if not (left.flags.c_contiguous and right.flags.c_contiguous): + return bool(((left == right) | (np.isnan(left) & np.isnan(right))).all()) + left = left.ravel() + right = right.ravel() + return lib.array_equivalent_float(left, right) def _array_equivalent_datetimelike(left: np.ndarray, right: np.ndarray) -> bool: - return np.array_equal(left.view("i8"), right.view("i8")) + return lib.array_equivalent_bytes(left.view("i8"), right.view("i8")) def _array_equivalent_object( diff --git a/pandas/core/indexes/datetimelike.py b/pandas/core/indexes/datetimelike.py index a5e001620e48a..976b65be40448 100644 --- a/pandas/core/indexes/datetimelike.py +++ b/pandas/core/indexes/datetimelike.py @@ -370,7 +370,7 @@ def equals(self, other: Any) -> bool: if type(self) != type(other): return False elif self.dtype == other.dtype: - return np.array_equal(self.asi8, other.asi8) + return lib.array_equivalent_bytes(self.asi8, other.asi8) elif (self.dtype.kind == "M" and self.tz == other.tz) or self.dtype.kind == "m": # type: ignore[attr-defined] # different units, otherwise matching try: @@ -379,7 +379,7 @@ def equals(self, other: Any) -> bool: except (OutOfBoundsDatetime, OutOfBoundsTimedelta): return False else: - return np.array_equal(left.view("i8"), right.view("i8")) + return lib.array_equivalent_bytes(left.view("i8"), right.view("i8")) return False def __contains__(self, key: Any) -> bool: diff --git a/pandas/core/indexes/multi.py b/pandas/core/indexes/multi.py index a50763a5efe50..0916d56b7cae4 100644 --- a/pandas/core/indexes/multi.py +++ b/pandas/core/indexes/multi.py @@ -4360,14 +4360,14 @@ def equals(self, other: object) -> bool: other_codes = other.codes[i] self_mask = self_codes == -1 other_mask = other_codes == -1 - if not np.array_equal(self_mask, other_mask): + if not lib.array_equivalent_bytes(self_mask, other_mask): return False self_level = self.levels[i] other_level = other.levels[i] new_codes = recode_for_categories( other_codes, other_level, self_level, copy=False ) - if not np.array_equal(self_codes, new_codes): + if not lib.array_equivalent_bytes(self_codes, new_codes): return False if not self_level[:0].equals(other_level[:0]): # e.g. Int64 != int64 From 2c604a1efc4d6fb2b20c228220c396ee0e629fe1 Mon Sep 17 00:00:00 2001 From: Brock Date: Sun, 12 Apr 2026 10:10:31 -0700 Subject: [PATCH 2/4] TYP: add type stubs for new Cython functions, fix pyright errors Co-Authored-By: Claude Opus 4.6 (1M context) --- pandas/_libs/lib.pyi | 10 ++++++++++ pandas/core/dtypes/missing.py | 5 +++-- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/pandas/_libs/lib.pyi b/pandas/_libs/lib.pyi index e50b301c34868..7b2bbb49e020e 100644 --- a/pandas/_libs/lib.pyi +++ b/pandas/_libs/lib.pyi @@ -221,6 +221,16 @@ def array_equivalent_object( right: npt.NDArray[np.object_], ) -> bool: ... def has_infs(arr: np.ndarray) -> bool: ... # const floating[:] +def has_nans(arr: np.ndarray) -> bool: ... # const floating[:] +def all_nans(arr: np.ndarray) -> bool: ... # const floating[:] +def array_equivalent_float( + left: np.ndarray, + right: np.ndarray, +) -> bool: ... # const floating[:] +def array_equivalent_bytes( + left: np.ndarray, + right: np.ndarray, +) -> bool: ... def has_only_ints_or_nan(arr: np.ndarray) -> bool: ... # const floating[:] def get_reverse_indexer( indexer: np.ndarray, # const intp_t[:] diff --git a/pandas/core/dtypes/missing.py b/pandas/core/dtypes/missing.py index 90143800cfb7b..b1a29637827cd 100644 --- a/pandas/core/dtypes/missing.py +++ b/pandas/core/dtypes/missing.py @@ -487,8 +487,9 @@ def _array_equivalent_float(left: np.ndarray, right: np.ndarray) -> bool: if not (left.flags.c_contiguous and right.flags.c_contiguous): return bool(((left == right) | (np.isnan(left) & np.isnan(right))).all()) # View complex as float pairs (complex128 -> float64, complex64 -> float32) - left = left.view(left.real.dtype) - right = right.view(right.real.dtype) + float_dtype = np.finfo(left.dtype).dtype + left = left.view(float_dtype) + right = right.view(float_dtype) if left.ndim > 1: if left.flags.f_contiguous and right.flags.f_contiguous: # .T is a C-contiguous view of an F-contiguous array From 5da15fef77f5179947d76babdbf0aa2fb4349438 Mon Sep 17 00:00:00 2001 From: Brock Date: Sun, 12 Apr 2026 12:35:26 -0700 Subject: [PATCH 3/4] BUG: skip memcmp fast path for structured dtypes with object fields Co-Authored-By: Claude Opus 4.6 (1M context) --- pandas/core/dtypes/missing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pandas/core/dtypes/missing.py b/pandas/core/dtypes/missing.py index b1a29637827cd..e827564051f99 100644 --- a/pandas/core/dtypes/missing.py +++ b/pandas/core/dtypes/missing.py @@ -477,7 +477,7 @@ def array_equivalent( ) and left.dtype != right.dtype: return False - if left.dtype == right.dtype: + if left.dtype == right.dtype and left.dtype.kind != "V": return lib.array_equivalent_bytes(left, right) return np.array_equal(left, right) From ebd1b21ce89b8da2249a708fc57834aa018374bb Mon Sep 17 00:00:00 2001 From: Brock Date: Sat, 18 Apr 2026 19:19:40 -0700 Subject: [PATCH 4/4] DOC: add docstrings to new array_equivalent / NaN-scan helpers Co-Authored-By: Claude Opus 4.7 (1M context) --- pandas/_libs/lib.pyx | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/pandas/_libs/lib.pyx b/pandas/_libs/lib.pyx index 7417fdaeb0d02..2ae18d1fb7a2f 100644 --- a/pandas/_libs/lib.pyx +++ b/pandas/_libs/lib.pyx @@ -501,6 +501,9 @@ def has_infs(const floating[:] arr) -> bool: @cython.wraparound(False) @cython.boundscheck(False) def has_nans(const floating[:] arr) -> bool: + """ + Faster equivalent to ``np.isnan(arr).any()``; exits on the first NaN found. + """ cdef: Py_ssize_t i, n = len(arr) Py_ssize_t n4 = n & ~3 # round down to multiple of 4 @@ -527,6 +530,9 @@ def has_nans(const floating[:] arr) -> bool: @cython.wraparound(False) @cython.boundscheck(False) def all_nans(const floating[:] arr) -> bool: + """ + Faster equivalent to ``np.isnan(arr).all()``; exits on the first non-NaN found. + """ cdef: Py_ssize_t i, n = len(arr) Py_ssize_t n4 = n & ~3 @@ -554,6 +560,10 @@ def all_nans(const floating[:] arr) -> bool: @cython.boundscheck(False) def array_equivalent_float(const floating[:] left, const floating[:] right) -> bool: + """ + Faster equivalent to ``((left == right) | (isnan(left) & isnan(right))).all()``; + exits on the first mismatch. Caller is responsible for checking shapes match. + """ cdef: Py_ssize_t i, n = len(left) floating lval, rval @@ -571,6 +581,12 @@ def array_equivalent_float(const floating[:] left, def array_equivalent_bytes(left, right) -> bool: + """ + Faster equivalent to ``np.array_equal(left, right)`` via ``memcmp`` on + C-contiguous inputs. Not safe for dtypes where distinct bit patterns can + represent the same value (e.g. floats with -0.0/+0.0 or NaN) or for arrays + that contain object pointers. + """ cdef: Py_ssize_t nbytes int ndim, idx