diff --git a/janitor/functions/_conditional_join/_range_indices.py b/janitor/functions/_conditional_join/_range_indices.py index f58840c9b..2ba7f46a2 100644 --- a/janitor/functions/_conditional_join/_range_indices.py +++ b/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} diff --git a/tests/functions/test_conditional_join_range_rmq.py b/tests/functions/test_conditional_join_range_rmq.py new file mode 100644 index 000000000..467c663b8 --- /dev/null +++ b/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