Skip to content

Commit 59d1c5b

Browse files
cmaloneyclaude
andcommitted
gh-159716: Faster integer array and memoryview searching
Speed up `in`, `.count()`, `.index()` and `array.array.remove()` for `array.array` and `memoryview` objects with integer formats by comparing the bytes of the items instead of creating an object for each item. The search is implemented using the buffer protocol to get a contiguous buffer, with a fallback to iteration and rich comparison. The buffer is held for the whole search, so a comparison which resizes the array or releases the memoryview now raises BufferError rather than allowing the resize while iterating. Add a `sq_contains` slot to memoryview so `in` uses the faster search when possible. The buffer protocol adds 7 ns of overhead per call on my M3 MacBook relative to passing native pointers. The overhead makes searches which stop at the first item or have fewer than two items slower. Longer cases get substantially faster; full results are below. `in`, `memoryview.index()` and `.count()` with the 'w' format (Py_UCS4) now work, because searching shares array's support for 'w'; they used to raise NotImplementedError. Other memoryview operations still don't support 'w'. `memoryview.index()` with an unsupported format now raises NotImplementedError even when there are no items to search, as `.count()` and `in` already did; it used to raise ValueError. pyperf, release build (-O3) without PGO or LTO, macOS arm64: | Benchmark | main | patch | |---------------------------------|:-------:|:-----------------------:| | array i[8] in (first) | 9.41 ns | 16.6 ns: 1.76x slower | | array i[8] in (last) | 57.3 ns | 19.0 ns: 3.01x faster | | array i[8] index (first) | 13.9 ns | 21.5 ns: 1.55x slower | | array i[8] count | 69.9 ns | 20.7 ns: 3.37x faster | | array d[8] index (first) | 20.4 ns | 26.3 ns: 1.29x slower | | array d[8] count | 94.4 ns | 96.8 ns: 1.03x slower | | memoryview i[8] index (first) | 14.4 ns | 22.6 ns: 1.57x slower | | memoryview i[8] count | 88.1 ns | 22.4 ns: 3.94x faster | | memoryview d[8] count | 125 ns | 96.7 ns: 1.29x faster | | array i[1k] count | 10.5 us | 136 ns: 77.22x faster | | memoryview i[1k:2] count | 10.0 us | 8.38 us: 1.19x faster | | memoryview d[1k] count | 10.8 us | 10.1 us: 1.07x faster | | array B[100k] in (miss) | 616 us | 1.77 us: 348.22x faster | | array i[100k] index (last) | 1.16 ms | 9.38 us: 123.36x faster | | array q[100k] count | 1.30 ms | 7.46 us: 173.79x faster | | memoryview B[100k] count | 805 us | 13.0 us: 61.90x faster | | memoryview i[100k] index (last) | 1.21 ms | 9.26 us: 131.19x faster | | Geometric mean | (ref) | 5.28x faster | Benchmark hidden because not significant (2): array d[1k] count, array w[1k] count Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 parent 8d11920 commit 59d1c5b

9 files changed

Lines changed: 633 additions & 334 deletions

File tree

‎Include/internal/pycore_pybuffer.h‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,23 @@ PyAPI_FUNC(int) _PyBuffer_ReleaseInInterpreter(
1515
PyAPI_FUNC(int) _PyBuffer_ReleaseInInterpreterAndRawFree(
1616
PyInterpreterState *interp, Py_buffer *view);
1717

18+
/* Search the items of the one-dimensional buffer of obj for value, comparing
19+
their bytes when that is equivalent to comparing them as objects.
20+
21+
_PyBuffer_FindItem() returns the index of the first item in [start, stop)
22+
equal to value, -1 if there is none, or -2 with an exception set on error.
23+
Negative start and stop count from the end, and are clipped like slice
24+
bounds.
25+
26+
_PyBuffer_CountItem() returns the number of items equal to value as an
27+
int, or NULL with an exception set on error.
28+
29+
Implemented in memoryobject.c. */
30+
PyAPI_FUNC(Py_ssize_t) _PyBuffer_FindItem(
31+
PyObject *obj, PyObject *value,
32+
Py_ssize_t start, Py_ssize_t stop);
33+
PyAPI_FUNC(PyObject *) _PyBuffer_CountItem(PyObject *obj, PyObject *value);
34+
1835
#ifdef __cplusplus
1936
}
2037
#endif

‎Lib/test/seq_tests.py‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
"""
44

55
import unittest
6+
import struct
67
import sys
78
import pickle
89
from test import support
@@ -442,3 +443,78 @@ def test_pickle(self):
442443
def test_free_after_iterating(self):
443444
support.check_free_after_iterating(self, iter, self.type2test)
444445
support.check_free_after_iterating(self, reversed, self.type2test)
446+
447+
448+
class NativeItemsSearchTest:
449+
"""Tests of count(), index() and "in" for 1-d buffer protocol objects.
450+
451+
Tests use "needle" for the value being searched for and "haystack" for
452+
the object being searched.
453+
"""
454+
455+
# All the native single item formats of memoryview.
456+
search_formats = ('b', 'B', 'h', 'H', 'i', 'I', 'l', 'L', 'q', 'Q',
457+
'n', 'N', 'P', 'e', 'f', 'd', 'Zf', 'Zd', '?', 'c')
458+
459+
def make_haystacks(self, fmt, raw):
460+
"""Make type-specific haystacks interpreting *raw* as *fmt*.
461+
462+
Subclasses make one haystack per layout under test (for example a
463+
contiguous, strided or unaligned memoryview). Yield (name, haystack)
464+
pairs, where name identifies the haystack in subtests.
465+
"""
466+
raise NotImplementedError
467+
468+
def run_search_cases(self, check):
469+
"""Call check(haystack, haystack_objects, needle) for every needle
470+
of every haystack, with a subtest for each haystack.
471+
472+
haystack_objects is the list of the items of haystack as Python
473+
objects, to check results against.
474+
"""
475+
class AlwaysEqual(int):
476+
def __eq__(self, other):
477+
return True
478+
__hash__ = int.__hash__
479+
480+
nan = float('nan')
481+
# Every byte value, then the doubles 0.0, -0.0, NaN and 1.0, all
482+
# twice so that values are found more than once.
483+
raw = (bytes(range(256)) + struct.pack('4d', 0.0, -0.0, nan, 1)) * 2
484+
# Needles which may not be items of the format: of other types, out
485+
# of range, or equal to items with different bytes.
486+
other_needles = [0, 1, -1, 2**64, 0.0, -0.0, nan, True, False, 1.0,
487+
b'a', AlwaysEqual(5)]
488+
for fmt in self.search_formats:
489+
for name, haystack in self.make_haystacks(fmt, raw):
490+
with self.subTest(fmt=fmt, haystack=name):
491+
haystack_objects = haystack.tolist()
492+
for needle in haystack_objects + other_needles:
493+
check(haystack, haystack_objects, needle)
494+
495+
def test_search_formats(self):
496+
self.run_search_cases(self.check_search)
497+
498+
def check_search(self, haystack, haystack_objects, needle):
499+
"""Validate optimized search matches python object search.
500+
501+
Compares search results between optimized format search and python
502+
rich object comparison.
503+
"""
504+
msg = f'needle={needle!r}'
505+
n = len(haystack_objects)
506+
matches = [i for i, x in enumerate(haystack_objects) if x == needle]
507+
self.assertEqual(haystack.count(needle), len(matches), msg)
508+
self.assertEqual(needle in haystack, bool(matches), msg)
509+
# index() bounds which are negative or past either end are clipped
510+
# like slice bounds.
511+
bounds = [(0, n), (5, -5), (-3, n), (2, n + 100), (-n - 5, n // 2)]
512+
for start, stop in bounds:
513+
r = range(n)[start:stop]
514+
expected = [i for i in matches if i in r]
515+
if expected:
516+
self.assertEqual(haystack.index(needle, start, stop),
517+
expected[0], msg)
518+
else:
519+
with self.assertRaises(ValueError, msg=msg):
520+
haystack.index(needle, start, stop)

‎Lib/test/test_array.py‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import collections.abc
66
import unittest
77
from test import support
8+
from test.seq_tests import NativeItemsSearchTest
89
from test.support import import_helper, script_helper
910
from test.support import os_helper
1011
from test.support import _2G
@@ -130,6 +131,22 @@ def test_typecodes_documented(self):
130131
self.assertGreaterEqual(array.array(typecode).itemsize,
131132
minimum_size)
132133

134+
def test_search_mutated_during_compare(self):
135+
# The buffer is held while searching, so the array cannot be resized
136+
# by a comparison.
137+
for search in (array.array.index, array.array.count,
138+
array.array.remove, array.array.__contains__):
139+
for mutate in (lambda a: a.append(9.0), lambda a: a.clear()):
140+
with self.subTest(search=search.__name__):
141+
a = array.array('d', [1.0, 2.0, 3.0])
142+
class Mutate:
143+
def __eq__(self, other):
144+
mutate(a)
145+
return False
146+
self.assertRaises(BufferError, search, a, Mutate())
147+
self.assertEqual(a.tolist(), [1.0, 2.0, 3.0])
148+
a.append(4.0) # The buffer was released.
149+
133150

134151
# Machine format codes.
135152
#
@@ -167,6 +184,45 @@ def test_typecodes_documented(self):
167184

168185
MACHINE_FORMAT_CODE_MAX = 27
169186

187+
class SearchTest(NativeItemsSearchTest, unittest.TestCase):
188+
# 'w' is tested by test_search_unicode. The raw is not valid UCS4 so
189+
# conversion to a list fails.
190+
search_formats = tuple(fmt for fmt in array.typecodes if fmt != 'w')
191+
192+
def make_haystacks(self, typecode, raw):
193+
yield 'array', array.array(typecode, raw)
194+
195+
def test_search_remove(self):
196+
self.run_search_cases(self.check_remove)
197+
198+
def check_remove(self, a, haystack_objects, needle):
199+
"""Check a.remove(needle) against removing the first item of
200+
haystack_objects equal to needle."""
201+
matches = [i for i, x in enumerate(haystack_objects) if x == needle]
202+
b = array.array(a.typecode, a)
203+
if matches:
204+
b.remove(needle)
205+
after = array.array(a.typecode, a)
206+
del after[matches[0]]
207+
self.assertEqual(b.tobytes(), after.tobytes(), repr(needle))
208+
else:
209+
with self.assertRaises(ValueError, msg=repr(needle)):
210+
b.remove(needle)
211+
212+
def test_search_unicode(self):
213+
a = array.array('w', 'abc\U0010ffff\0ab' * 2)
214+
haystack_objects = a.tolist()
215+
for needle in ['a', 'b', '\U0010ffff', '\0', 'z', 'ab', '', b'a', 97]:
216+
self.check_search(a, haystack_objects, needle)
217+
self.check_remove(a, haystack_objects, needle)
218+
# memoryview does not support 'w' items, but can search them.
219+
m = memoryview(a)
220+
self.assertEqual(m.count('a'), a.count('a'))
221+
self.assertEqual(m.index('\0', 1), a.index('\0', 1))
222+
self.assertIn('\U0010ffff', m)
223+
self.assertNotIn('z', m)
224+
self.assertRaises(NotImplementedError, m.__getitem__, 0)
225+
170226

171227
class ArrayReconstructorTest(unittest.TestCase):
172228

‎Lib/test/test_buffer.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2963,6 +2963,16 @@ def assert_out_of_bounds_error(self, dim):
29632963
self.assertEqual(str(cm.exception),
29642964
"index out of bounds on dimension %d" % (dim,))
29652965

2966+
def test_memoryview_search_unsupported_format(self):
2967+
nd = ndarray([1, 2], shape=[2], format='<i')
2968+
m = memoryview(nd)
2969+
# Ensure empty and non-empty behave the same.
2970+
for mv in m, m[:0]:
2971+
self.assertRaises(NotImplementedError, mv.count, 1)
2972+
self.assertRaises(NotImplementedError, mv.index, 1)
2973+
self.assertRaises(NotImplementedError, lambda: 1 in mv)
2974+
self.assertRaises(NotImplementedError, m.index, 1, 2)
2975+
29662976
def test_memoryview_index(self):
29672977

29682978
# ndim = 0
@@ -4631,6 +4641,17 @@ def __release_buffer__(self, buffer):
46314641
self.assertEqual(mv.tobytes(), b"hello")
46324642
self.assertFalse(wr.held)
46334643

4644+
def test_search_restricted_memoryview(self):
4645+
# The memoryview passed to __release_buffer__() refuses to export new
4646+
# views, but can still be searched.
4647+
results = []
4648+
class Searcher(bytearray):
4649+
def __release_buffer__(self, mv):
4650+
results.append((97 in mv, mv.count(98), mv.index(99)))
4651+
4652+
memoryview(Searcher(b"abc")).release()
4653+
self.assertEqual(results, [(True, 1, 2)])
4654+
46344655
def test_buffer_flags(self):
46354656
class PossiblyMutable:
46364657
def __init__(self, data, mutable) -> None:

‎Lib/test/test_memoryview.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
from itertools import product
1919
from test import support
20+
from test.seq_tests import NativeItemsSearchTest
2021
from test.support import import_helper, threading_helper
2122

2223

@@ -949,6 +950,26 @@ def test_overflows_in_floats(self):
949950
complex_view[0] = 1e300j
950951
self.assertEqual(complex_view[0], 123+321j)
951952

953+
def test_search_released_during_compare(self):
954+
# The buffer is held while searching, so a comparison cannot release
955+
# the memoryview. Leaving the with block fails if it is still held.
956+
for name in 'index', 'count', '__contains__':
957+
with memoryview(array.array('d', [1.0, 2.0, 3.0])) as m:
958+
class Release:
959+
def __eq__(self, other):
960+
m.release()
961+
return False
962+
with self.assertRaises(BufferError, msg=name):
963+
getattr(m, name)(Release())
964+
965+
966+
class SearchTest(NativeItemsSearchTest, unittest.TestCase):
967+
def make_haystacks(self, fmt, raw):
968+
mv = memoryview(raw).cast(fmt)
969+
yield 'contiguous', mv
970+
yield 'strided', mv[::3]
971+
yield 'unaligned', memoryview(b'\0' + raw)[1:].cast(fmt)
972+
952973

953974
@threading_helper.requires_working_threading()
954975
@support.requires_resource("cpu")
Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
Speed up ``in``, :meth:`memoryview.count` and :meth:`memoryview.index` for
2+
:class:`memoryview` objects with integer formats. Attempts to resize during
3+
these operations may now raise :exc:`BufferError`. Searching is 60-130x faster
4+
for 100,000 items. :meth:`memoryview.index` now raises
5+
:exc:`NotImplementedError` for an unsupported format even when there are no
6+
items to search, like :meth:`~memoryview.count` and ``in``.
Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
Speed up ``in``, :meth:`array.array.count`, :meth:`array.array.index` and
2+
:meth:`array.array.remove` for :class:`array.array` objects with integer
3+
typecodes. Attempts to resize during these operations may now raise
4+
:exc:`BufferError`. Searching is 120-350x faster for 100,000 items.

0 commit comments

Comments
 (0)