Skip to content

Commit cd2a5fa

Browse files
jbrockmendelclaude
andcommitted
PERF: optimize _libs.algos for SIMD auto-vectorization and GIL release
- diff_2d: use raw C pointers for inner loops so the C compiler can auto-vectorize the subtraction (~20-25% faster for DataFrame.diff) - take_2d_axis0/axis1: use raw C pointers for inner row copy/conversion loops enabling auto-vectorization (~10% faster for cross-type reindex) - nancorr: add no_nans fast path skipping per-element mask checks - nancorr_spearman: add @cython.cdivision(True) - rank_sorted_1d: add @cython.cdivision(True) to remove zero-division checks from the hot ranking loop - rank_1d: skip unnecessary putmask when check_mask is False - pad/backfill/pad_inplace/pad_2d_inplace: release GIL for numeric types - get_fill_indexer: release GIL - unique_deltas: use C array instead of Python list to avoid GIL interaction per unique delta - is_monotonic: remove redundant assignments before break Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 4b69e1c commit cd2a5fa

3 files changed

Lines changed: 198 additions & 56 deletions

File tree

doc/source/whatsnew/v3.1.0.rst

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,8 @@ Performance improvements
136136
- Performance improvement in :func:`read_sas` for compressed SAS7BDAT files by reusing the decompression buffer instead of allocating per row (:issue:`47339`)
137137
- Performance improvement in :func:`read_sas` when decoding strings (:issue:`47339`)
138138
- Performance improvement in :func:`util.hash_pandas_object` for PyArrow-backed string and binary types by using PyArrow's ``dictionary_encode`` instead of converting to NumPy for factorization (:issue:`48964`)
139+
- Performance improvement in :meth:`DataFrame.corr` and :meth:`DataFrame.cov` when data contains no NaN values (:issue:`64857`)
140+
- Performance improvement in :meth:`DataFrame.diff` by using raw C pointers to enable auto-vectorization of the inner loop (:issue:`64857`)
139141
- Performance improvement in :meth:`DataFrame.from_records` when passing a 2D :class:`numpy.ndarray` (:issue:`22025`)
140142
- Performance improvement in :meth:`DataFrame.insert` when the number of blocks is small (:issue:`57641`)
141143
- Performance improvement in :meth:`DataFrame.loc` with non-unique masked index (:issue:`56759`)
@@ -151,6 +153,7 @@ Performance improvements
151153
- Performance improvement in datetime/timedelta unit conversion (e.g. ``datetime64[s]`` to ``datetime64[ns]``) (:issue:`35025`)
152154
- Performance improvement in indexing a :class:`DataFrame` with a :class:`CategoricalIndex` of :class:`Interval` categories (:issue:`61928`)
153155
- Performance improvement in indexing a :class:`MultiIndex` with a list-like indexer (:issue:`55786`)
156+
- Performance improvement in indexing operations that use ``take_2d`` (e.g. :meth:`DataFrame.reindex`) by using raw C pointers to enable auto-vectorization of the inner row copy (:issue:`64857`)
154157
- Performance improvement in partial-string indexing on a monotonic decreasing :class:`DatetimeIndex` or :class:`PeriodIndex` (:issue:`64811`)
155158
- Performance improvement in plotting :class:`DatetimeIndex` with multiplied frequencies (e.g. ``"1000ms"``, ``"100s"``) (:issue:`50355`)
156159
- Performance improvement in reading zip-compressed files (e.g. :func:`read_pickle`, :func:`read_csv`) on Python < 3.12 (:issue:`59279`)

pandas/_libs/algos.pyx

Lines changed: 106 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -136,25 +136,37 @@ cpdef ndarray[int64_t, ndim=1] unique_deltas(const int64_t[:] arr):
136136
An ordered ndarray[int64_t]
137137
"""
138138
cdef:
139-
Py_ssize_t i, n = len(arr)
139+
Py_ssize_t i, n = len(arr), num_uniques = 0
140140
int64_t val
141141
khiter_t k
142142
kh_int64_t *table
143143
int ret = 0
144-
list uniques = []
144+
int64_t *uniques = NULL
145145
ndarray[int64_t, ndim=1] result
146146

147147
table = kh_init_int64()
148148
kh_resize_int64(table, 10)
149+
150+
# n - 1 is the max possible number of unique deltas
151+
if n > 1:
152+
uniques = <int64_t*>malloc((n - 1) * sizeof(int64_t))
153+
if uniques is NULL:
154+
kh_destroy_int64(table)
155+
raise MemoryError()
156+
149157
for i in range(n - 1):
150158
val = arr[i + 1] - arr[i]
151159
k = kh_get_int64(table, val)
152160
if k == table.n_buckets:
153161
kh_put_int64(table, val, &ret)
154-
uniques.append(val)
162+
uniques[num_uniques] = val
163+
num_uniques += 1
155164
kh_destroy_int64(table)
156165

157-
result = np.array(uniques, dtype=np.int64)
166+
result = np.empty(num_uniques, dtype=np.int64)
167+
if num_uniques > 0:
168+
memcpy(cnp.PyArray_DATA(result), uniques, num_uniques * sizeof(int64_t))
169+
free(uniques)
158170
result.sort()
159171
return result
160172

@@ -347,6 +359,7 @@ def nancorr(const float64_t[:, :] mat, bint cov=False, minp=None):
347359
uint8_t[:, :] mask
348360
int64_t nobs = 0
349361
float64_t vx, vy, dx, dy, meanx, meany, divisor, ssqdmx, ssqdmy, covxy, val
362+
bint no_nans
350363

351364
N, K = (<object>mat).shape
352365
if minp is None:
@@ -356,25 +369,40 @@ def nancorr(const float64_t[:, :] mat, bint cov=False, minp=None):
356369

357370
result = np.empty((K, K), dtype=np.float64)
358371
mask = np.isfinite(mat).view(np.uint8)
372+
no_nans = np.asarray(mask).all()
359373

360374
with nogil:
361375
for xi in range(K):
362376
for yi in range(xi + 1):
363377
# Welford's method for the variance-calculation
364378
# https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance
365379
nobs = ssqdmx = ssqdmy = covxy = meanx = meany = 0
366-
for i in range(N):
367-
if mask[i, xi] and mask[i, yi]:
380+
381+
if no_nans:
382+
nobs = N
383+
for i in range(N):
368384
vx = mat[i, xi]
369385
vy = mat[i, yi]
370-
nobs += 1
371386
dx = vx - meanx
372387
dy = vy - meany
373-
meanx += 1. / nobs * dx
374-
meany += 1. / nobs * dy
388+
meanx += 1. / (i + 1) * dx
389+
meany += 1. / (i + 1) * dy
375390
ssqdmx += (vx - meanx) * dx
376391
ssqdmy += (vy - meany) * dy
377392
covxy += (vx - meanx) * dy
393+
else:
394+
for i in range(N):
395+
if mask[i, xi] and mask[i, yi]:
396+
vx = mat[i, xi]
397+
vy = mat[i, yi]
398+
nobs += 1
399+
dx = vx - meanx
400+
dy = vy - meany
401+
meanx += 1. / nobs * dx
402+
meany += 1. / nobs * dy
403+
ssqdmx += (vx - meanx) * dx
404+
ssqdmy += (vy - meany) * dy
405+
covxy += (vx - meanx) * dy
378406

379407
if nobs < minpv:
380408
result[xi, yi] = result[yi, xi] = NaN
@@ -843,8 +871,6 @@ def is_monotonic(const numeric_object_t[:] arr, bint timelike):
843871
is_monotonic_dec = 0
844872
break
845873
if not is_monotonic_inc and not is_monotonic_dec:
846-
is_monotonic_inc = 0
847-
is_monotonic_dec = 0
848874
break
849875
prev = cur
850876

@@ -1022,7 +1048,6 @@ def rank_1d(
10221048
# will flip the ordering to still end up with lowest rank.
10231049
# Symmetric logic applies to `na_option == 'bottom'`
10241050
nans_rank_highest = ascending ^ (na_option == "top")
1025-
nan_fill_val = get_rank_nan_fill_val(nans_rank_highest, <numeric_object_t>0)
10261051
if nans_rank_highest:
10271052
order = [masked_vals, mask]
10281053
else:
@@ -1031,7 +1056,9 @@ def rank_1d(
10311056
if check_labels:
10321057
order.append(labels)
10331058

1034-
np.putmask(masked_vals, mask, nan_fill_val)
1059+
if check_mask:
1060+
nan_fill_val = get_rank_nan_fill_val(nans_rank_highest, <numeric_object_t>0)
1061+
np.putmask(masked_vals, mask, nan_fill_val)
10351062
# putmask doesn't accept a memoryview, so we assign as a separate step
10361063
masked_vals_memview = masked_vals
10371064

@@ -1356,7 +1383,7 @@ ctypedef fused out_t:
13561383
@cython.wraparound(False)
13571384
def diff_2d(
13581385
const diff_t[:, :] arr,
1359-
ndarray[out_t, ndim=2] out,
1386+
out_t[:, :] out,
13601387
Py_ssize_t periods,
13611388
int axis,
13621389
bint datetimelike=False,
@@ -1365,6 +1392,15 @@ def diff_2d(
13651392
Py_ssize_t i, j, sx, sy, start, stop
13661393
bint f_contig = arr.is_f_contig()
13671394
diff_t left, right
1395+
# Raw pointer variables for auto-vectorization of inner loops.
1396+
# Memoryview indexing generates strided pointer arithmetic that
1397+
# prevents the C compiler from auto-vectorizing even when the
1398+
# data is contiguous. Using raw typed pointers gives the compiler
1399+
# unit-stride access patterns it can vectorize.
1400+
const diff_t *arr_row
1401+
const diff_t *arr_prev_row
1402+
out_t *out_row
1403+
bint rows_contiguous
13681404

13691405
# Disable for unsupported dtype combinations,
13701406
# see https://github.com/cython/cython/issues/2646
@@ -1381,6 +1417,13 @@ def diff_2d(
13811417
# We put this inside an indented else block to avoid cython build
13821418
# warnings about unreachable code
13831419
sx, sy = (<object>arr).shape
1420+
1421+
# Check whether both arrays have unit stride along axis 1 (columns),
1422+
# i.e. rows are contiguous. When true, inner loops over j can use
1423+
# raw C pointers enabling SIMD auto-vectorization.
1424+
rows_contiguous = (arr.strides[1] == <Py_ssize_t>sizeof(diff_t) and
1425+
out.strides[1] == <Py_ssize_t>sizeof(out_t))
1426+
13841427
with nogil:
13851428
if f_contig:
13861429
if axis == 0:
@@ -1421,33 +1464,64 @@ def diff_2d(
14211464
start, stop = periods, sx
14221465
else:
14231466
start, stop = 0, sx + periods
1424-
for i in range(start, stop):
1425-
for j in range(sy):
1426-
left = arr[i, j]
1427-
right = arr[i - periods, j]
1428-
if out_t is int64_t and datetimelike:
1429-
if left == NPY_NAT or right == NPY_NAT:
1430-
out[i, j] = NPY_NAT
1467+
if rows_contiguous:
1468+
for i in range(start, stop):
1469+
arr_row = &arr[i, 0]
1470+
arr_prev_row = &arr[i - periods, 0]
1471+
out_row = &out[i, 0]
1472+
for j in range(sy):
1473+
if out_t is int64_t and datetimelike:
1474+
left = arr_row[j]
1475+
right = arr_prev_row[j]
1476+
if left == NPY_NAT or right == NPY_NAT:
1477+
out_row[j] = NPY_NAT
1478+
else:
1479+
out_row[j] = left - right
1480+
else:
1481+
out_row[j] = arr_row[j] - arr_prev_row[j]
1482+
else:
1483+
for i in range(start, stop):
1484+
for j in range(sy):
1485+
left = arr[i, j]
1486+
right = arr[i - periods, j]
1487+
if out_t is int64_t and datetimelike:
1488+
if left == NPY_NAT or right == NPY_NAT:
1489+
out[i, j] = NPY_NAT
1490+
else:
1491+
out[i, j] = left - right
14311492
else:
14321493
out[i, j] = left - right
1433-
else:
1434-
out[i, j] = left - right
14351494
else:
14361495
if periods >= 0:
14371496
start, stop = periods, sy
14381497
else:
14391498
start, stop = 0, sy + periods
1440-
for i in range(sx):
1441-
for j in range(start, stop):
1442-
left = arr[i, j]
1443-
right = arr[i, j - periods]
1444-
if out_t is int64_t and datetimelike:
1445-
if left == NPY_NAT or right == NPY_NAT:
1446-
out[i, j] = NPY_NAT
1499+
if rows_contiguous:
1500+
for i in range(sx):
1501+
arr_row = &arr[i, 0]
1502+
out_row = &out[i, 0]
1503+
for j in range(start, stop):
1504+
if out_t is int64_t and datetimelike:
1505+
left = arr_row[j]
1506+
right = arr_row[j - periods]
1507+
if left == NPY_NAT or right == NPY_NAT:
1508+
out_row[j] = NPY_NAT
1509+
else:
1510+
out_row[j] = left - right
1511+
else:
1512+
out_row[j] = arr_row[j] - arr_row[j - periods]
1513+
else:
1514+
for i in range(sx):
1515+
for j in range(start, stop):
1516+
left = arr[i, j]
1517+
right = arr[i, j - periods]
1518+
if out_t is int64_t and datetimelike:
1519+
if left == NPY_NAT or right == NPY_NAT:
1520+
out[i, j] = NPY_NAT
1521+
else:
1522+
out[i, j] = left - right
14471523
else:
14481524
out[i, j] = left - right
1449-
else:
1450-
out[i, j] = left - right
14511525

14521526

14531527
# ----------------------------------------------------------------------

0 commit comments

Comments
 (0)