perf: use range query kernel for wide dual ranges

janitor/functions/_conditional_join/_range_indices.py

@@ -6,6 +6,55 @@
 
 from janitor.functions._conditional_join import _binary_search, _helpers
 
+# ELI5: building the Rust tree costs roughly one pass over the right table,
+# while the existing Python path revisits every candidate interval. Require
+# more than eight right-table passes' worth of candidate work before paying
+# the tree-build and Python/Rust call overhead. This is a conservative,
+# benchmark-derived crossover heuristic, not a correctness or complexity
+# invariant; revisit it if end-to-end workloads change.
+_RANGE_RMQ_WORK_FACTOR = 8.0
+
+
+def _range_rmq(
+    right_index: np.ndarray,
+    starts: np.ndarray,
+    ends: np.ndarray,
+    keep: str,
+) -> np.ndarray | None:
+    """
+    Use the Rust range-query tree only when it is likely to amortize its build.
+
+    ELI5: the current Python path reopens every interval. The Rust tree walks
+    the right index once to prepare reusable block answers, so it is useful
+    only when the total interval width is several times the right-table size.
+    Returning ``None`` keeps all existing fast paths and dtype combinations
+    unchanged when the new extension is unavailable or unsuitable.
+    """
+    if (
+        right_index.dtype != np.dtype(np.int64)
+        or right_index.size < 32
+        or starts.size < 32
+    ):
+        return None
+    total_width = float(np.asarray(ends - starts, dtype=np.int64).sum(dtype=np.float64))
+    if total_width <= _RANGE_RMQ_WORK_FACTOR * right_index.size:
+        return None
+    function_name = (
+        "index_starts_and_ends_keep_first_direct"
+        if keep == "first"
+        else "index_starts_and_ends_keep_last_direct"
+    )
+    function = getattr(janitor_rs, function_name, None)
+    if function is None:
+        return None
+    return np.asarray(
+        function(
+            index=right_index,
+            starts=np.asarray(starts, dtype=np.int64),
+            ends=np.asarray(ends, dtype=np.int64),
+        )
+    )
+
 
 def _range_indices(
     df: pd.DataFrame, right: pd.DataFrame, ge_gt: tuple, le_lt: tuple, is_sorted: bool
@@ -96,10 +145,16 @@ def _build_indices(
     if (keep == "last") and right_is_sorted:
         return {"left_index": left_index, "right_index": right_index[ends - 1]}
     if keep == "first":
+        right_rmq = _range_rmq(right_index, starts, ends, keep)
+        if right_rmq is not None:
+            return {"left_index": left_index, "right_index": right_rmq}
         right = [right_index[start:end] for start, end in zip(starts, ends)]
         right = [arr.min() for arr in right]
         return {"left_index": left_index, "right_index": right}
     if keep == "last":
+        right_rmq = _range_rmq(right_index, starts, ends, keep)
+        if right_rmq is not None:
+            return {"left_index": left_index, "right_index": right_rmq}
         right = [right_index[start:end] for start, end in zip(starts, ends)]
         right = [arr.max() for arr in right]
         return {"left_index": left_index, "right_index": right}

tests/functions/test_conditional_join_range_rmq.py

@@ -0,0 +1,84 @@
+import numpy as np
+
+from janitor.functions._conditional_join import _range_indices
+
+
+def _fixture():
+    right_index = np.array([24, 58, 2, 13, 91, 7, 40, 5] * 8, dtype=np.int64)
+    starts = np.zeros(64, dtype=np.int64)
+    ends = np.full(64, 16, dtype=np.int64)
+    return right_index, starts, ends
+
+
+def test_range_rmq_dispatches_first_to_rust(monkeypatch):
+    right_index, starts, ends = _fixture()
+    called = {}
+
+    def fake_kernel(**kwargs):
+        called.update(kwargs)
+        return np.full(starts.size, 2, dtype=np.int64)
+
+    monkeypatch.setattr(
+        _range_indices.janitor_rs,
+        "index_starts_and_ends_keep_first_direct",
+        fake_kernel,
+        raising=False,
+    )
+
+    result = _range_indices._range_rmq(right_index, starts, ends, "first")
+
+    assert np.array_equal(result, np.full(starts.size, 2, dtype=np.int64))
+    assert called["index"] is right_index
+    assert np.array_equal(called["starts"], starts)
+    assert np.array_equal(called["ends"], ends)
+
+
+def test_range_rmq_dispatches_last_to_rust(monkeypatch):
+    right_index, starts, ends = _fixture()
+
+    monkeypatch.setattr(
+        _range_indices.janitor_rs,
+        "index_starts_and_ends_keep_last_direct",
+        lambda **kwargs: np.full(starts.size, 91, dtype=np.int64),
+        raising=False,
+    )
+
+    result = _range_indices._range_rmq(right_index, starts, ends, "last")
+
+    assert np.array_equal(result, np.full(starts.size, 91, dtype=np.int64))
+
+
+def test_range_rmq_falls_back_for_narrow_workloads(monkeypatch):
+    right_index = np.arange(64, dtype=np.int64)
+    starts = np.arange(64, dtype=np.int64)
+    ends = starts + 1
+
+    def unexpected_kernel(**kwargs):
+        raise AssertionError("narrow ranges must retain the existing path")
+
+    monkeypatch.setattr(
+        _range_indices.janitor_rs,
+        "index_starts_and_ends_keep_first_direct",
+        unexpected_kernel,
+        raising=False,
+    )
+
+    assert _range_indices._range_rmq(right_index, starts, ends, "first") is None
+
+
+def test_range_rmq_falls_back_for_non_int64_index(monkeypatch):
+    right_index = np.arange(64, dtype=np.int32)
+    starts = np.zeros(64, dtype=np.int64)
+    ends = np.full(64, 32, dtype=np.int64)
+
+    def unexpected_kernel(**kwargs):
+        raise AssertionError("unsupported index dtypes must retain the existing path")
+
+    monkeypatch.setattr(
+        _range_indices.janitor_rs,
+        "index_starts_and_ends_keep_first_direct",
+        unexpected_kernel,
+        raising=False,
+    )
+
+    assert _range_indices._range_rmq(right_index, starts, ends, "first") is None