Make Numba optional with graceful fallback to pure Python #75

ranx/__init__.py

@@ -7,12 +7,21 @@
     "plot",
     "Qrels",
     "Run",
+    "use_numba",
+    "set_numba_enabled",
 ]
 
-from numba import config
-
+from .config import set_numba_enabled, use_numba
 from .data_structures import Qrels, Run
 from .meta import compare, evaluate, fuse, normalize, optimize_fusion, plot
 
-# Set numba threading layer to workqueue
-config.THREADING_LAYER = "workqueue"
+# Conditional Numba configuration
+if use_numba():
+    try:
+        from numba import config
+
+        # Set numba threading layer to workqueue
+        config.THREADING_LAYER = "workqueue"
+    except ImportError:
+        # Numba not available, silently continue
+        pass

ranx/config.py

@@ -0,0 +1,46 @@
+"""Configuration system for ranx library."""
+
+import os
+from typing import Optional
+
+# Global configuration state
+_USE_NUMBA: Optional[bool] = None
+
+
+def use_numba() -> bool:
+    """
+    Check if Numba should be used for performance optimizations.
+
+    Returns True by default, but can be disabled via:
+    1. Environment variable: RANX_USE_NUMBA=false
+    2. Programmatically: set_numba_enabled(False)
+
+    Returns:
+        bool: True if Numba should be used, False otherwise
+    """
+    global _USE_NUMBA
+    if _USE_NUMBA is None:
+        env_value = os.environ.get("RANX_USE_NUMBA", "true").lower()
+        _USE_NUMBA = env_value not in ("false", "0", "no", "off")
+    return _USE_NUMBA
+
+
+def set_numba_enabled(enabled: bool) -> None:
+    """
+    Programmatically enable or disable Numba usage.
+
+    Args:
+        enabled: True to enable Numba, False to disable
+    """
+    global _USE_NUMBA
+    _USE_NUMBA = enabled
+
+
+def reset_numba_config() -> None:
+    """
+    Reset Numba configuration to default (reads from environment again).
+
+    Primarily used for testing purposes.
+    """
+    global _USE_NUMBA
+    _USE_NUMBA = None

ranx/data_structures/generic.py

@@ -1,24 +1,29 @@
-from numba import njit, types
-from numba.typed import Dict as TypedDict
-from numba.typed import List as TypedList
+from ..decorators import create_typed_dict, create_typed_list
+
+# Handle Numba-specific imports conditionally
+try:
+    from numba import types
+
+    NUMBA_AVAILABLE = True
+except ImportError:
+    NUMBA_AVAILABLE = False
 
 
-@njit(cache=True)
 def create_empty_results_dict():
-    return TypedDict.empty(
-        key_type=types.unicode_type,
-        value_type=types.float64,
+    return create_typed_dict(
+        key_type=types.unicode_type if NUMBA_AVAILABLE else None,
+        value_type=types.float64 if NUMBA_AVAILABLE else None,
     )
 
 
-@njit(cache=True)
 def create_empty_results_dict_list(length):
-    return TypedList([create_empty_results_dict() for _ in range(length)])
+    return create_typed_list(
+        initial_list=[create_empty_results_dict() for _ in range(length)]
+    )
 
 
-@njit(cache=True)
 def convert_results_dict_list_to_run(q_ids, results_dict_list):
-    combined_run = TypedDict()
+    combined_run = create_typed_dict()
 
     for i, q_id in enumerate(q_ids):
         combined_run[q_id] = results_dict_list[i]

ranx/data_structures/qrels.py

@@ -7,9 +7,21 @@
 import numpy as np
 import orjson
 import pandas as pd
-from numba import njit, prange, types
-from numba.typed import Dict as TypedDict
-from numba.typed import List as TypedList
+
+try:
+    from numba import njit, prange, types
+
+    NUMBA_AVAILABLE = True
+except ImportError:
+    NUMBA_AVAILABLE = False
+
+
+try:
+    from numba.typed import Dict as TypedDict
+    from numba.typed import List as TypedList
+except ImportError:
+    TypedDict = dict
+    TypedList = list
 
 from .common import (
     add_and_sort,

ranx/data_structures/run.py

@@ -5,9 +5,17 @@
 
 import numpy as np
 import pandas as pd
-from numba import types
-from numba.typed import Dict as TypedDict
-from numba.typed import List as TypedList
+
+try:
+    from numba import types
+    from numba.typed import Dict as TypedDict
+    from numba.typed import List as TypedList
+
+    NUMBA_AVAILABLE = True
+except ImportError:
+    NUMBA_AVAILABLE = False
+    TypedDict = dict
+    TypedList = list
 
 from ..io import download, load_json, load_lz4, save_json, save_lz4
 from .common import (

ranx/decorators.py

@@ -0,0 +1,135 @@
+"""Conditional decorators for optional Numba support."""
+
+from .config import use_numba
+
+# Check if Numba is available
+try:
+    from numba import jit, njit
+    from numba.typed import Dict as TypedDict
+    from numba.typed import List as TypedList
+
+    NUMBA_AVAILABLE = True
+except ImportError:
+    NUMBA_AVAILABLE = False
+    # Create dummy objects for when Numba is not available
+    TypedDict = dict
+    TypedList = list
+
+
+def maybe_njit(*args, **kwargs):
+    """
+    Conditional njit decorator that falls back to identity function when Numba is disabled.
+
+    This decorator will apply Numba's njit compilation when:
+    1. Numba is available (installed)
+    2. Numba is enabled via configuration
+
+    Otherwise, it returns the function unchanged (pure Python execution).
+
+    Args:
+        *args, **kwargs: Same arguments as numba.njit
+
+    Returns:
+        Function decorator that conditionally applies Numba compilation
+    """
+
+    def decorator(func):
+        if NUMBA_AVAILABLE and use_numba():
+            return njit(*args, **kwargs)(func)
+        else:
+            return func
+
+    # Handle the case where maybe_njit is used without arguments: @maybe_njit
+    if len(args) == 1 and callable(args[0]) and not kwargs:
+        func = args[0]
+        if NUMBA_AVAILABLE and use_numba():
+            return njit()(func)
+        else:
+            return func
+
+    return decorator
+
+
+def maybe_jit(*args, **kwargs):
+    """
+    Conditional jit decorator that falls back to identity function when Numba is disabled.
+
+    Similar to maybe_njit but uses jit instead of njit.
+
+    Args:
+        *args, **kwargs: Same arguments as numba.jit
+
+    Returns:
+        Function decorator that conditionally applies Numba compilation
+    """
+
+    def decorator(func):
+        if NUMBA_AVAILABLE and use_numba():
+            return jit(*args, **kwargs)(func)
+        else:
+            return func
+
+    # Handle the case where maybe_jit is used without arguments: @maybe_jit
+    if len(args) == 1 and callable(args[0]) and not kwargs:
+        func = args[0]
+        if NUMBA_AVAILABLE and use_numba():
+            return jit()(func)
+        else:
+            return func
+
+    return decorator
+
+
+# Note: prange requires separate implementations because Numba needs compile-time
+# knowledge of whether to use prange or range. See metrics files for examples.
+
+
+def create_typed_dict(key_type=None, value_type=None, initial_dict=None):
+    """
+    Create a typed dictionary that falls back to regular dict when Numba is disabled.
+
+    Args:
+        key_type: Numba type for keys (ignored when Numba disabled)
+        value_type: Numba type for values (ignored when Numba disabled)
+        initial_dict: Initial dictionary to populate
+
+    Returns:
+        numba.typed.Dict when Numba enabled, regular dict otherwise
+    """
+    if NUMBA_AVAILABLE and use_numba():
+        if initial_dict:
+            typed_dict = TypedDict()
+            for k, v in initial_dict.items():
+                typed_dict[k] = v
+            return typed_dict
+        elif key_type is not None and value_type is not None:
+            return TypedDict.empty(key_type, value_type)
+        else:
+            return TypedDict()
+    else:
+        return dict(initial_dict) if initial_dict else {}
+
+
+def create_typed_list(item_type=None, initial_list=None):
+    """
+    Create a typed list that falls back to regular list when Numba is disabled.
+
+    Args:
+        item_type: Numba type for items (ignored when Numba disabled)
+        initial_list: Initial list to populate
+
+    Returns:
+        numba.typed.List when Numba enabled, regular list otherwise
+    """
+    if NUMBA_AVAILABLE and use_numba():
+        if initial_list:
+            typed_list = TypedList()
+            for item in initial_list:
+                typed_list.append(item)
+            return typed_list
+        elif item_type is not None:
+            return TypedList.empty_list(item_type)
+        else:
+            return TypedList()
+    else:
+        return list(initial_list) if initial_list else []

ranx/metrics/average_precision.py

@@ -1,14 +1,13 @@
 from typing import Union
 
-import numba
 import numpy as np
-from numba import njit, prange
 
+from ..decorators import maybe_njit
 from .common import clean_qrels, fix_k
 
 
 # LOW LEVEL FUNCTIONS ==========================================================
-@njit(cache=True)
+@maybe_njit(cache=True)
 def _average_precision(qrels, run, k, rel_lvl):
     qrels = clean_qrels(qrels, rel_lvl)
     if len(qrels) == 0:
@@ -42,18 +41,44 @@ def _average_precision(qrels, run, k, rel_lvl):
     return np.sum(precision_scores) / qrels.shape[0]
 
 
-@njit(cache=True, parallel=True)
-def _average_precision_parallel(qrels, run, k, rel_lvl):
+# Handle parallel version with conditional compilation
+try:
+    from numba import njit, prange
+
+    @njit(cache=True, parallel=True)
+    def _average_precision_parallel_numba(qrels, run, k, rel_lvl):
+        scores = np.zeros((len(qrels)), dtype=np.float64)
+        for i in prange(len(qrels)):
+            scores[i] = _average_precision(qrels[i], run[i], k, rel_lvl)
+        return scores
+
+    NUMBA_AVAILABLE = True
+except ImportError:
+    NUMBA_AVAILABLE = False
+
+
+def _average_precision_numpy(qrels, run, k, rel_lvl):
+    """NumPy fallback implementation."""
     scores = np.zeros((len(qrels)), dtype=np.float64)
-    for i in prange(len(qrels)):
+    for i in range(len(qrels)):
         scores[i] = _average_precision(qrels[i], run[i], k, rel_lvl)
     return scores
 
 
+def _average_precision_parallel(qrels, run, k, rel_lvl):
+    """Dispatch to best available implementation."""
+    from ..config import use_numba
+
+    if NUMBA_AVAILABLE and use_numba():
+        return _average_precision_parallel_numba(qrels, run, k, rel_lvl)
+    else:
+        return _average_precision_numpy(qrels, run, k, rel_lvl)
+
+
 # HIGH LEVEL FUNCTIONS =========================================================
 def average_precision(
-    qrels: Union[np.ndarray, numba.typed.List],
-    run: Union[np.ndarray, numba.typed.List],
+    qrels: Union[np.ndarray, list],
+    run: Union[np.ndarray, list],
     k: int = 0,
     rel_lvl: int = 1,
 ) -> np.ndarray:
@@ -72,9 +97,9 @@ def average_precision(
     - $R$ is the total number of relevant documents.
 
     Args:
-        qrels (Union[np.ndarray, numba.typed.List]): IDs and relevance scores of _relevant_ documents.
+        qrels: IDs and relevance scores of _relevant_ documents.
 
-        run (Union[np.ndarray, numba.typed.List]): IDs and relevance scores of _retrieved_ documents.
+        run: IDs and relevance scores of _retrieved_ documents.
 
         k (int, optional): Number of retrieved documents to consider. k=0 means all retrieved documents will be considered. Defaults to 0.
 

ranx/metrics/common.py

@@ -1,12 +1,13 @@
 import numpy as np
-from numba import njit
 
+from ..decorators import maybe_njit
 
-@njit(cache=True)
+
+@maybe_njit(cache=True)
 def clean_qrels(qrels, rel_lvl):
     return qrels[np.argwhere(qrels[:, 1] >= rel_lvl).flatten()]
 
 
-@njit(cache=True)
+@maybe_njit(cache=True)
 def fix_k(k, run):
     return run.shape[0] if k == 0 or k > run.shape[0] else k