benchmark the KVP accumulator

envs/lenskit-py3.10-ci.yaml

@@ -27,6 +27,7 @@ dependencies:
   - numpy>=1.23
   - pandas<3,>=1.5
   - pyproject2conda~=0.11
+  - pytest-benchmark==4.*
   - pytest-cov>=2.12
   - pytest-doctestplus>=0.9
   - pytest==7.*

envs/lenskit-py3.10-dev.yaml

@@ -32,6 +32,7 @@ dependencies:
   - numpy>=1.23
   - pandas<3,>=1.5
   - pyproject2conda~=0.11
+  - pytest-benchmark==4.*
   - pytest-cov>=2.12
   - pytest-doctestplus>=0.9
   - pytest==7.*

envs/lenskit-py3.10-test.yaml

@@ -20,6 +20,7 @@ dependencies:
   - numba<0.59,>=0.56
   - numpy>=1.23
   - pandas<3,>=1.5
+  - pytest-benchmark==4.*
   - pytest-cov>=2.12
   - pytest-doctestplus>=0.9
   - pytest==7.*

envs/lenskit-py3.11-ci.yaml

@@ -27,6 +27,7 @@ dependencies:
   - numpy>=1.23
   - pandas<3,>=1.5
   - pyproject2conda~=0.11
+  - pytest-benchmark==4.*
   - pytest-cov>=2.12
   - pytest-doctestplus>=0.9
   - pytest==7.*

envs/lenskit-py3.11-dev.yaml

@@ -32,6 +32,7 @@ dependencies:
   - numpy>=1.23
   - pandas<3,>=1.5
   - pyproject2conda~=0.11
+  - pytest-benchmark==4.*
   - pytest-cov>=2.12
   - pytest-doctestplus>=0.9
   - pytest==7.*

envs/lenskit-py3.11-test.yaml

@@ -20,6 +20,7 @@ dependencies:
   - numba<0.59,>=0.56
   - numpy>=1.23
   - pandas<3,>=1.5
+  - pytest-benchmark==4.*
   - pytest-cov>=2.12
   - pytest-doctestplus>=0.9
   - pytest==7.*

pyproject.toml

@@ -54,6 +54,7 @@ test = [
   "pytest-doctestplus >= 0.9",
   "pytest-cov >= 2.12",
   "coverage >= 5",
+  "pytest-benchmark ==4.*",
   "cython ==3.*",
   "hypothesis >= 6",
 ]

tests/test_util_accum.py

@@ -5,13 +5,14 @@
 # SPDX-License-Identifier: MIT
 
 import numpy as np
+from numba import njit
 
-from lenskit.util.accum import kvp_minheap_insert, kvp_minheap_sort
-
-
-from hypothesis import given, assume, settings
-import hypothesis.strategies as st
 import hypothesis.extra.numpy as nph
+import hypothesis.strategies as st
+from hypothesis import assume, given, settings
+from pytest import mark
+
+from lenskit.util.accum import kvp_minheap_insert, kvp_minheap_sort
 
 
 def test_kvp_add_to_empty():
@@ -185,3 +186,25 @@ def test_kvp_sort(values):
     assert vs[-1] == np.min(ovs)
     assert all(ks == oks[ord])
     assert all(vs == ovs[ord])
+
+
+@mark.benchmark(group="KVPSort")
+def test_kvp_sort_numba(rng, benchmark):
+    N = 10000
+    K = 500
+    in_keys = np.arange(N)
+    in_vals = rng.uniform(size=N)
+
+    def op():
+        ks = np.zeros(K, np.int32)
+        vs = np.zeros(K, np.float64)
+        ep = 0
+        for i in range(N):
+            ep = kvp_minheap_insert(0, ep, K, in_keys[i], in_vals[i], ks, vs)
+
+        kvp_minheap_sort(0, ep, ks, vs)
+
+    # dry run to compile
+    op()
+
+    benchmark(op)

tests/test_util_kvp.py

@@ -9,6 +9,7 @@
 import hypothesis.extra.numpy as nph
 import hypothesis.strategies as st
 from hypothesis import assume, given, settings
+from pytest import mark
 
 from lenskit.util.kvp import KVPHeap
 
@@ -195,3 +196,21 @@ def test_kvp_sort(values):
     assert vs[-1] == np.min(ovs)
     assert all(ks == oks[ord])
     assert all(vs == ovs[ord])
+
+
+@mark.benchmark(group="KVPSort")
+def test_kvp_sort_cython(rng, benchmark):
+    N = 10000
+    K = 500
+    in_keys = np.arange(N)
+    in_vals = rng.uniform(size=N)
+
+    def op():
+        ks = np.zeros(K, np.int32)
+        vs = np.zeros(K, np.float64)
+        kvp = KVPHeap(0, 0, K, ks, vs)
+        for i in range(N):
+            kvp.insert(in_keys[i], in_vals[i])
+        kvp.sort()
+
+    benchmark(op)