From 59410e810a1e93f096103c79bad9fc02362f071a Mon Sep 17 00:00:00 2001
From: Ferdinand Schlatt <ferdi@seltz.ai>
Date: Wed, 23 Sep 2026 15:45:19 +0200
Subject: [PATCH] Avoid numba recompilation for varying document-id lengths
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

`Run`/`Qrels` store document ids as a fixed-width unicode numpy array whose
dtype (`<U{max_len}`) is derived from the longest id in the input. numba
encodes that width in the type signature of `create_and_sort` and the helpers
it calls, so each distinct max id length triggers a fresh JIT compilation
(seconds to minutes). Workloads that build many `Run`/`Qrels` objects with
differing id lengths — e.g. computing IR metrics per batch during model
validation — recompile the numba kernels on nearly every call, which can turn
a few-minute evaluation into tens of minutes.

Round the fixed width up to the next power of two so the number of distinct
dtypes is bounded by ~log2(max_len). The kernels then compile a handful of
times and are reused. The width only ever grows, so ids are never truncated
and metric values are unchanged.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 ranx/data_structures/qrels.py | 3 ++-
 ranx/data_structures/run.py   | 3 ++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/ranx/data_structures/qrels.py b/ranx/data_structures/qrels.py
index f0d24bf..4ad6ada 100644
--- a/ranx/data_structures/qrels.py
+++ b/ranx/data_structures/qrels.py
@@ -74,7 +74,8 @@ def __init__(self, qrels: Dict[str, Dict[str, int]] = None, name: str = None):
             # Doc IDs
             doc_ids = [list(doc.keys()) for doc in qrels.values()]
             max_len = max(len(y) for x in doc_ids for y in x)
-            dtype = f"<U{max_len}"
+            width = 1 << (max_len - 1).bit_length() if max_len > 1 else 1
+            dtype = f"<U{width}"
             doc_ids = TypedList([np.array(x, dtype=dtype) for x in doc_ids])
 
             # Scores
diff --git a/ranx/data_structures/run.py b/ranx/data_structures/run.py
index 6df514d..31141d5 100644
--- a/ranx/data_structures/run.py
+++ b/ranx/data_structures/run.py
@@ -59,7 +59,8 @@ def __init__(self, run: Dict[str, Dict[str, float]] = None, name: str = None):
             # Doc IDs
             doc_ids = [list(doc.keys()) for doc in run.values()]
             max_len = max(len(y) for x in doc_ids for y in x)
-            dtype = f"<U{max_len}"
+            width = 1 << (max_len - 1).bit_length() if max_len > 1 else 1
+            dtype = f"<U{width}"
             doc_ids = TypedList([np.array(x, dtype=dtype) for x in doc_ids])
 
             # Scores