From 97799be869c1f2d68e32cc95babb832821d2027f Mon Sep 17 00:00:00 2001
From: Michael Ekstrand <md@ekstrandom.net>
Date: Mon, 4 Mar 2024 21:02:53 -0500
Subject: [PATCH] benchmark the KVP accumulator

---
 envs/lenskit-py3.10-ci.yaml   |  1 +
 envs/lenskit-py3.10-dev.yaml  |  1 +
 envs/lenskit-py3.10-test.yaml |  1 +
 envs/lenskit-py3.11-ci.yaml   |  1 +
 envs/lenskit-py3.11-dev.yaml  |  1 +
 envs/lenskit-py3.11-test.yaml |  1 +
 pyproject.toml                |  1 +
 tests/test_util_accum.py      | 33 ++++++++++++++++++++++++++++-----
 tests/test_util_kvp.py        | 19 +++++++++++++++++++
 9 files changed, 54 insertions(+), 5 deletions(-)

diff --git a/envs/lenskit-py3.10-ci.yaml b/envs/lenskit-py3.10-ci.yaml
index 93e405898..4621f7358 100644
--- a/envs/lenskit-py3.10-ci.yaml
+++ b/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.*
diff --git a/envs/lenskit-py3.10-dev.yaml b/envs/lenskit-py3.10-dev.yaml
index c63a3afa1..62d765293 100644
--- a/envs/lenskit-py3.10-dev.yaml
+++ b/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.*
diff --git a/envs/lenskit-py3.10-test.yaml b/envs/lenskit-py3.10-test.yaml
index f13b01127..ed91d7372 100644
--- a/envs/lenskit-py3.10-test.yaml
+++ b/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.*
diff --git a/envs/lenskit-py3.11-ci.yaml b/envs/lenskit-py3.11-ci.yaml
index 22b1d3f99..bce3e94bc 100644
--- a/envs/lenskit-py3.11-ci.yaml
+++ b/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.*
diff --git a/envs/lenskit-py3.11-dev.yaml b/envs/lenskit-py3.11-dev.yaml
index 7b20b76db..d8372e941 100644
--- a/envs/lenskit-py3.11-dev.yaml
+++ b/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.*
diff --git a/envs/lenskit-py3.11-test.yaml b/envs/lenskit-py3.11-test.yaml
index 0b34ed77f..8b3d05b65 100644
--- a/envs/lenskit-py3.11-test.yaml
+++ b/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.*
diff --git a/pyproject.toml b/pyproject.toml
index 8f117a73b..f04e61e65 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -54,6 +54,7 @@ test = [
   "pytest-doctestplus >= 0.9",
   "pytest-cov >= 2.12",
   "coverage >= 5",
+  "pytest-benchmark ==4.*",
   "cython ==3.*",
   "hypothesis >= 6",
 ]
diff --git a/tests/test_util_accum.py b/tests/test_util_accum.py
index 8e6cd4f35..055aef33e 100644
--- a/tests/test_util_accum.py
+++ b/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)
diff --git a/tests/test_util_kvp.py b/tests/test_util_kvp.py
index ee28fe04d..4bc8d644d 100644
--- a/tests/test_util_kvp.py
+++ b/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)