From 0414f81ae8ebc62d92b396b02d8842d0949a8916 Mon Sep 17 00:00:00 2001 From: samukweku Date: Sat, 22 Aug 2026 08:34:16 +1000 Subject: [PATCH 1/4] [PERF] Use NumPy prefix sums for integer conditional_join range-sum aggs Replace the O(sum of interval widths) Rust kernels backing join_agg(..., aggfunc=[(col, "sum")]) for the <-only, >-only, and arbitrary-interval dispatch paths with an O(n + m) NumPy prefix-sum, for integer dtypes. Float dtypes are untouched (tracked separately in #1671, since the existing float kernels use compensated summation). Also removes a pre-existing duplicate definition of _sum_starts_ends found while editing it. Issue #1648 Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01S1gZKDRiZoBLZnXXXjW3gt --- CHANGELOG.md | 4 + benchmarks/bench_join_agg_sum.py | 168 ++++++++++++++++ .../_conditional_join/_agg_functions.py | 97 ++++----- pixi.lock | 4 +- .../test_conditional_join_agg_int_sum.py | 185 ++++++++++++++++++ 5 files changed, 400 insertions(+), 58 deletions(-) create mode 100644 benchmarks/bench_join_agg_sum.py create mode 100644 tests/functions/test_conditional_join_agg_int_sum.py diff --git a/CHANGELOG.md b/CHANGELOG.md index e464228b9..825788617 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,10 @@ # Changelog ## [Unreleased] +- [PERF] Use an O(n + m) NumPy prefix-sum for integer `conditional_join` + range-sum aggregations (`join_agg(..., aggfunc=[(col, "sum")])`), + replacing the O(sum of interval widths) Rust kernels for the `<`-only, + `>`-only, and arbitrary-interval dispatch paths. - Issue #1648 @samukweku - [ENH] Avoid copying column data during `conditional_join` input validation. - Issue #1645, PR #1642 @samukweku - [ENH] Speed up `conditional_join` with an unsorted right join key and diff --git a/benchmarks/bench_join_agg_sum.py b/benchmarks/bench_join_agg_sum.py new file mode 100644 index 000000000..29c2e7c37 --- /dev/null +++ b/benchmarks/bench_join_agg_sum.py @@ -0,0 +1,168 @@ +"""Benchmark for the integer prefix-sum `join_agg` kernels (Issue #1648). + +Reproduces the shape of the issue's local numbers: + +- Kernel-level: the old Rust `compute_sum_start*`/`compute_sum_end*`/ + `compute_sum_start_end*` kernels (O(sum of interval widths)) versus the + new NumPy prefix-sum kernels in `_agg_functions` (O(n + m)), for + "suffixes" (`<`), "prefixes" (`>`), and "arbitrary intervals" (a range + join with no equality condition). +- End-to-end: a real `join_agg(..., aggfunc=[("value", "sum")])` call, + which now dispatches through the new prefix-sum kernels automatically. + +Run with: + + pixi run python benchmarks/bench_join_agg_sum.py + pixi run python benchmarks/bench_join_agg_sum.py --large # 10M/50M rows + +Large-N results (Apple Silicon, single run, int64), kernel-level only -- +the old Rust kernels are the O(sum of interval widths) algorithm this issue +replaces, so they are not re-run here; they are already intractable at the +20,000-row scale (see the small-N results above):: + + n = 10,000,000 + kernel suffixes (prefix-sum): 0.066 s + kernel prefixes (prefix-sum): 0.056 s + kernel intervals (prefix-sum): 0.098 s + end-to-end join_agg sum '<': 2.027 s (peak RSS ~1.6 GB) + + n = 50,000,000 + kernel suffixes (prefix-sum): 0.419 s + kernel prefixes (prefix-sum): 0.351 s + kernel intervals (prefix-sum): 0.652 s + end-to-end join_agg sum '<': 18.795 s (peak RSS ~6.1 GB) + +The prefix-sum kernel itself stays linear and sub-second through 50M rows. +The end-to-end number is dominated by `conditional_join`'s match-finding +(building `starts`/`right_index`), which is unrelated to this issue's +scope (the sum kernel) and unchanged by this PR. +""" + +import timeit + +import janitor_rs +import numpy as np +import pandas as pd + +import janitor # noqa: F401 registers the DataFrame accessor +from janitor.functions._conditional_join import _agg_functions as agg + + +def _bench(label: str, fn, number: int = 5) -> float: + """Time `fn`, print the per-call average, and return it in seconds.""" + seconds = timeit.timeit(fn, number=number) / number + print(f"{label:<28} {seconds * 1000:10.3f} ms") + return seconds + + +def kernel_level(n: int) -> None: + """Compare the old Rust kernels against the new prefix-sum kernels.""" + print(f"\n-- kernel-level, n={n:,} (int64) --") + rng = np.random.default_rng(0) + arr = rng.integers(-1000, 1000, size=n, dtype="int64") + booleans = np.zeros(n, dtype=bool) + + # suffixes: every row starts somewhere and runs to the end + starts = rng.integers(0, n, size=n, dtype="int64") + _bench( + "suffixes (rust)", + lambda: janitor_rs.compute_sum_start_int64( + arr=arr, starts=starts, booleans=booleans + ), + ) + _bench( + "suffixes (prefix-sum)", + lambda: agg._sum_starts(arr=arr, starts=starts, booleans=booleans), + ) + + # prefixes: every row runs from the start to somewhere + ends = rng.integers(0, n, size=n, dtype="int64") + _bench( + "prefixes (rust)", + lambda: janitor_rs.compute_sum_end_int64(arr=arr, ends=ends, booleans=booleans), + ) + _bench( + "prefixes (prefix-sum)", + lambda: agg._sum_ends(arr=arr, ends=ends, booleans=booleans), + ) + + # arbitrary intervals + lo = rng.integers(0, n, size=n, dtype="int64") + hi = rng.integers(0, n, size=n, dtype="int64") + starts2, ends2 = np.minimum(lo, hi), np.maximum(lo, hi) + _bench( + "intervals (rust)", + lambda: janitor_rs.compute_sum_start_end_int64( + arr=arr, starts=starts2, ends=ends2, booleans=booleans + ), + ) + _bench( + "intervals (prefix-sum)", + lambda: agg._sum_starts_ends( + arr=arr, starts=starts2, ends=ends2, booleans=booleans + ), + ) + + +def end_to_end(n: int) -> None: + """Time a real `join_agg(..., aggfunc=[("value", "sum")])` call.""" + print(f"\n-- end-to-end join_agg, n={n:,} (int64, '<') --") + rng = np.random.default_rng(0) + left = pd.DataFrame({"key": rng.integers(0, n, size=n)}) + right = pd.DataFrame( + { + "key": np.sort(rng.integers(0, n, size=n)), + "value": rng.integers(-1000, 1000, size=n), + } + ) + + def run(): + """Run one `join_agg` sum call.""" + return left.join_agg(right, ("key", "key", "<"), aggfunc=[("value", "sum")]) + + _bench("join_agg sum '<'", run, number=3) + + +def large_scale(n: int) -> None: + """Prefix-sum kernel only (the old Rust kernels are intractable here).""" + print(f"\n-- kernel-level, n={n:,} (int64), prefix-sum only --") + rng = np.random.default_rng(0) + arr = rng.integers(-1000, 1000, size=n, dtype="int64") + booleans = np.zeros(n, dtype=bool) + starts = rng.integers(0, n, size=n, dtype="int64") + ends = rng.integers(0, n, size=n, dtype="int64") + lo = rng.integers(0, n, size=n, dtype="int64") + hi = rng.integers(0, n, size=n, dtype="int64") + starts2, ends2 = np.minimum(lo, hi), np.maximum(lo, hi) + + _bench( + "suffixes (prefix-sum)", + lambda: agg._sum_starts(arr=arr, starts=starts, booleans=booleans), + number=3, + ) + _bench( + "prefixes (prefix-sum)", + lambda: agg._sum_ends(arr=arr, ends=ends, booleans=booleans), + number=3, + ) + _bench( + "intervals (prefix-sum)", + lambda: agg._sum_starts_ends( + arr=arr, starts=starts2, ends=ends2, booleans=booleans + ), + number=3, + ) + end_to_end(n) + + +if __name__ == "__main__": + import sys + + if "--large" in sys.argv: + for n in (10_000_000, 50_000_000): + large_scale(n) + else: + for n in (1_000, 20_000): + kernel_level(n) + for n in (1_000, 20_000): + end_to_end(n) diff --git a/janitor/functions/_conditional_join/_agg_functions.py b/janitor/functions/_conditional_join/_agg_functions.py index 96fa0b7a2..6cbce5305 100644 --- a/janitor/functions/_conditional_join/_agg_functions.py +++ b/janitor/functions/_conditional_join/_agg_functions.py @@ -1,6 +1,28 @@ import janitor_rs import numpy as np +_INTEGER_DTYPE_NAMES = frozenset( + {"int64", "int32", "int16", "int8", "uint64", "uint32", "uint16", "uint8"} +) + + +def _int64_prefix_sums(arr: np.ndarray, booleans: np.ndarray) -> np.ndarray: + """ + Running total of `arr` widened to int64, null positions zeroed, + with a leading zero so `prefix[i]` is the sum of `arr[:i]`. + + ELI5: write the running total once; any `[start:end)` range sum is + then just `prefix[end] - prefix[start]` -- two lookups and a + subtraction, instead of re-adding every element in the range again. + """ + widened = arr.astype(np.int64) + if booleans.any(): + widened[booleans] = 0 + prefix = np.empty(widened.size + 1, dtype=np.int64) + prefix[0] = 0 + np.cumsum(widened, out=prefix[1:]) + return prefix + def _sum_starts( arr: np.ndarray, @@ -10,19 +32,14 @@ def _sum_starts( """ Compute sum """ + dtype_name = arr.dtype.name + if dtype_name in _INTEGER_DTYPE_NAMES: + prefix = _int64_prefix_sums(arr=arr, booleans=booleans) + return prefix[-1] - prefix[starts] mapping = { - "int64": janitor_rs.compute_sum_start_int64, - "int32": janitor_rs.compute_sum_start_int32, - "int16": janitor_rs.compute_sum_start_int16, - "int8": janitor_rs.compute_sum_start_int8, - "uint64": janitor_rs.compute_sum_start_uint64, - "uint32": janitor_rs.compute_sum_start_uint32, - "uint16": janitor_rs.compute_sum_start_uint16, - "uint8": janitor_rs.compute_sum_start_uint8, "float64": janitor_rs.compute_sum_start_f64, "float32": janitor_rs.compute_sum_start_f32, } - dtype_name = arr.dtype.name try: func = mapping[dtype_name] except KeyError: @@ -38,19 +55,14 @@ def _sum_ends( """ Compute sum """ + dtype_name = arr.dtype.name + if dtype_name in _INTEGER_DTYPE_NAMES: + prefix = _int64_prefix_sums(arr=arr, booleans=booleans) + return prefix[ends] mapping = { - "int64": janitor_rs.compute_sum_end_int64, - "int32": janitor_rs.compute_sum_end_int32, - "int16": janitor_rs.compute_sum_end_int16, - "int8": janitor_rs.compute_sum_end_int8, - "uint64": janitor_rs.compute_sum_end_uint64, - "uint32": janitor_rs.compute_sum_end_uint32, - "uint16": janitor_rs.compute_sum_end_uint16, - "uint8": janitor_rs.compute_sum_end_uint8, "float64": janitor_rs.compute_sum_end_f64, "float32": janitor_rs.compute_sum_end_f32, } - dtype_name = arr.dtype.name try: func = mapping[dtype_name] except KeyError: @@ -799,19 +811,21 @@ def _sum_starts_ends( """ Compute sum """ + dtype_name = arr.dtype.name + if dtype_name in _INTEGER_DTYPE_NAMES: + prefix = _int64_prefix_sums(arr=arr, booleans=booleans) + result = prefix[ends] - prefix[starts] + # an empty (or inverted) range contributes nothing, matching the + # Rust `for nn in start_..end_` loop, which never iterates when + # start_ >= end_ + empty_range = starts >= ends + if empty_range.any(): + result[empty_range] = 0 + return result mapping = { - "int64": janitor_rs.compute_sum_start_end_int64, - "int32": janitor_rs.compute_sum_start_end_int32, - "int16": janitor_rs.compute_sum_start_end_int16, - "int8": janitor_rs.compute_sum_start_end_int8, - "uint64": janitor_rs.compute_sum_start_end_uint64, - "uint32": janitor_rs.compute_sum_start_end_uint32, - "uint16": janitor_rs.compute_sum_start_end_uint16, - "uint8": janitor_rs.compute_sum_start_end_uint8, "float64": janitor_rs.compute_sum_start_end_f64, "float32": janitor_rs.compute_sum_start_end_f32, } - dtype_name = arr.dtype.name try: func = mapping[dtype_name] except KeyError: @@ -886,35 +900,6 @@ def _prod_starts_ends_matches( ) -def _sum_starts_ends( - arr: np.ndarray, - starts: np.ndarray, - ends: np.ndarray, - booleans: np.ndarray, -) -> tuple: - """ - Compute sum - """ - mapping = { - "int64": janitor_rs.compute_sum_start_end_int64, - "int32": janitor_rs.compute_sum_start_end_int32, - "int16": janitor_rs.compute_sum_start_end_int16, - "int8": janitor_rs.compute_sum_start_end_int8, - "uint64": janitor_rs.compute_sum_start_end_uint64, - "uint32": janitor_rs.compute_sum_start_end_uint32, - "uint16": janitor_rs.compute_sum_start_end_uint16, - "uint8": janitor_rs.compute_sum_start_end_uint8, - "float64": janitor_rs.compute_sum_start_end_f64, - "float32": janitor_rs.compute_sum_start_end_f32, - } - dtype_name = arr.dtype.name - try: - func = mapping[dtype_name] - except KeyError: - raise KeyError(f"Unsupported data type -> {dtype_name}") - return func(arr=arr, starts=starts, ends=ends, booleans=booleans) - - def _sum_starts_ends_matches( arr: np.ndarray, starts: np.ndarray, diff --git a/pixi.lock b/pixi.lock index 8b8f1a44c..9c2eca750 100644 --- a/pixi.lock +++ b/pixi.lock @@ -22775,8 +22775,8 @@ packages: timestamp: 1774796815820 - pypi: ./ name: pyjanitor - version: 0.32.23 - sha256: 2c551db6e9b84114e74ba359a19fafa89130d1d5ce2b322f906d8b934e72926a + version: 0.32.24 + sha256: 048f055a9e41a71f61809b58114fe6dd79913a4232e76f8a5cf2e4b661ff05c5 requires_dist: - pandas>=3.0.0 - natsort>=8.4.0,<9 diff --git a/tests/functions/test_conditional_join_agg_int_sum.py b/tests/functions/test_conditional_join_agg_int_sum.py new file mode 100644 index 000000000..98396ecf7 --- /dev/null +++ b/tests/functions/test_conditional_join_agg_int_sum.py @@ -0,0 +1,185 @@ +"""Focused unit tests for the integer prefix-sum kernels backing `join_agg`. + +Covers `_int64_prefix_sums`, `_sum_starts`, `_sum_ends`, and +`_sum_starts_ends` in `janitor.functions._conditional_join._agg_functions` +(Issue #1648) -- the O(n + m) NumPy replacements for the Rust +`compute_sum_start*`, `compute_sum_end*`, and `compute_sum_start_end*` +kernels, for integer dtypes only. +""" + +import numpy as np +import pandas as pd +import pytest + +from janitor.functions._conditional_join import _agg_functions + +INTEGER_DTYPES = [ + "int64", + "int32", + "int16", + "int8", + "uint64", + "uint32", + "uint16", + "uint8", +] + +EXTENSION_DTYPE_NAMES = { + "int64": "Int64", + "int32": "Int32", + "int16": "Int16", + "int8": "Int8", + "uint64": "UInt64", + "uint32": "UInt32", + "uint16": "UInt16", + "uint8": "UInt8", +} + + +def naive_sum_starts(arr, starts, booleans): + """Reference implementation matching the Rust kernel loop exactly.""" + out = np.zeros(starts.size, dtype="int64") + n = arr.size + for pos, start in enumerate(starts): + total = 0 + for nn in range(start, n): + if booleans[nn]: + continue + total += int(arr[nn]) + out[pos] = total + return out + + +def naive_sum_ends(arr, ends, booleans): + out = np.zeros(ends.size, dtype="int64") + for pos, end in enumerate(ends): + total = 0 + for nn in range(0, end): + if booleans[nn]: + continue + total += int(arr[nn]) + out[pos] = total + return out + + +def naive_sum_starts_ends(arr, starts, ends, booleans): + out = np.zeros(starts.size, dtype="int64") + for pos, (start, end) in enumerate(zip(starts, ends)): + total = 0 + for nn in range(start, end): + if booleans[nn]: + continue + total += int(arr[nn]) + out[pos] = total + return out + + +@pytest.mark.parametrize("dtype", INTEGER_DTYPES) +def test_sum_starts_matches_naive_with_nulls(dtype): + """Nulls scattered at start/middle/end should be skipped, not zeroed-in.""" + arr = pd.array([1, 2, 3, None, 5, None, 7, 8], dtype=EXTENSION_DTYPE_NAMES[dtype]) + booleans = pd.isna(arr) + arr = arr.to_numpy(dtype=dtype, na_value=0, copy=False) + starts = np.array([0, 1, 3, 5, 7, 8], dtype="int64") + + expected = naive_sum_starts(arr, starts, booleans) + actual = _agg_functions._sum_starts(arr=arr, starts=starts, booleans=booleans) + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize("dtype", INTEGER_DTYPES) +def test_sum_ends_matches_naive_with_nulls(dtype): + arr = pd.array([1, 2, 3, None, 5, None, 7, 8], dtype=EXTENSION_DTYPE_NAMES[dtype]) + booleans = pd.isna(arr) + arr = arr.to_numpy(dtype=dtype, na_value=0, copy=False) + ends = np.array([0, 1, 3, 5, 7, 8], dtype="int64") + + expected = naive_sum_ends(arr, ends, booleans) + actual = _agg_functions._sum_ends(arr=arr, ends=ends, booleans=booleans) + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize("dtype", INTEGER_DTYPES) +def test_sum_starts_ends_matches_naive_with_nulls(dtype): + arr = pd.array([1, 2, 3, None, 5, None, 7, 8], dtype=EXTENSION_DTYPE_NAMES[dtype]) + booleans = pd.isna(arr) + arr = arr.to_numpy(dtype=dtype, na_value=0, copy=False) + starts = np.array([0, 0, 2, 4, 8, 5], dtype="int64") + ends = np.array([0, 8, 2, 3, 8, 2], dtype="int64") # includes empty/inverted ranges + + expected = naive_sum_starts_ends(arr, starts, ends, booleans) + actual = _agg_functions._sum_starts_ends( + arr=arr, starts=starts, ends=ends, booleans=booleans + ) + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize("dtype", INTEGER_DTYPES) +def test_all_null_range_is_zero(dtype): + arr = np.array([1, 2, 3, 4], dtype=dtype) + booleans = np.ones(4, dtype=bool) + starts = np.array([0], dtype="int64") + ends = np.array([4], dtype="int64") + + assert _agg_functions._sum_starts(arr=arr, starts=starts, booleans=booleans)[0] == 0 + assert _agg_functions._sum_ends(arr=arr, ends=ends, booleans=booleans)[0] == 0 + assert ( + _agg_functions._sum_starts_ends( + arr=arr, starts=starts, ends=ends, booleans=booleans + )[0] + == 0 + ) + + +@pytest.mark.parametrize("dtype", INTEGER_DTYPES) +def test_empty_array(dtype): + arr = np.array([], dtype=dtype) + booleans = np.array([], dtype=bool) + starts = np.array([0], dtype="int64") + ends = np.array([0], dtype="int64") + + assert _agg_functions._sum_starts(arr=arr, starts=starts, booleans=booleans)[0] == 0 + assert _agg_functions._sum_ends(arr=arr, ends=ends, booleans=booleans)[0] == 0 + assert ( + _agg_functions._sum_starts_ends( + arr=arr, starts=starts, ends=ends, booleans=booleans + )[0] + == 0 + ) + + +def test_uint64_values_above_int64_max_reinterpret_like_rust(): + """uint64 values past i64::MAX must bit-reinterpret to negative int64, + matching the Rust kernel's `current as i64` cast.""" + huge = np.iinfo("uint64").max # 2**64 - 1 -> -1 as int64 + arr = np.array([huge, huge - 1, 5], dtype="uint64") + booleans = np.zeros(3, dtype=bool) + starts = np.array([0], dtype="int64") + + result = _agg_functions._sum_starts(arr=arr, starts=starts, booleans=booleans) + # -1 + -2 + 5 == 2, computed in wrapping int64 arithmetic + assert result[0] == 2 + + +def test_int64_accumulation_wraps_like_rust_release_mode(): + """Overflowing the running total wraps (two's complement), it does not + raise -- matching a Rust release build (overflow-checks disabled).""" + n = 100 + arr = np.full(n, np.iinfo("int64").max // 2, dtype="int64") + booleans = np.zeros(n, dtype=bool) + starts = np.array([0], dtype="int64") + + result = _agg_functions._sum_starts(arr=arr, starts=starts, booleans=booleans) + assert result[0] == -100 + + +def test_starts_beyond_ends_is_empty_range(): + arr = np.array([1, 2, 3, 4, 5], dtype="int64") + booleans = np.zeros(5, dtype=bool) + starts = np.array([3], dtype="int64") + ends = np.array([1], dtype="int64") + + result = _agg_functions._sum_starts_ends( + arr=arr, starts=starts, ends=ends, booleans=booleans + ) + assert result[0] == 0 From 952bf4ff8e6f2bb3c34fa228877dbb9bae870a4c Mon Sep 17 00:00:00 2001 From: samukweku Date: Sat, 22 Aug 2026 08:34:47 +1000 Subject: [PATCH 2/4] docs: add PR number to changelog entry for #1648 Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01S1gZKDRiZoBLZnXXXjW3gt --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 825788617..b2bb4dd33 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ - [PERF] Use an O(n + m) NumPy prefix-sum for integer `conditional_join` range-sum aggregations (`join_agg(..., aggfunc=[(col, "sum")])`), replacing the O(sum of interval widths) Rust kernels for the `<`-only, - `>`-only, and arbitrary-interval dispatch paths. - Issue #1648 @samukweku + `>`-only, and arbitrary-interval dispatch paths. - Issue #1648, PR #1673 @samukweku - [ENH] Avoid copying column data during `conditional_join` input validation. - Issue #1645, PR #1642 @samukweku - [ENH] Speed up `conditional_join` with an unsorted right join key and From 1757af3f00618045294d38fce88fa3ebfbba4eb2 Mon Sep 17 00:00:00 2001 From: samukweku Date: Sat, 22 Aug 2026 12:46:58 +1000 Subject: [PATCH 3/4] chore: drop benchmark script, folded into PR description instead Repo has no existing benchmarks/ convention -- keeping the numbers reproducible in the PR description avoids introducing an uncommitted- to-CI script and a new top-level directory precedent. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01S1gZKDRiZoBLZnXXXjW3gt --- benchmarks/bench_join_agg_sum.py | 168 ------------------------------- 1 file changed, 168 deletions(-) delete mode 100644 benchmarks/bench_join_agg_sum.py diff --git a/benchmarks/bench_join_agg_sum.py b/benchmarks/bench_join_agg_sum.py deleted file mode 100644 index 29c2e7c37..000000000 --- a/benchmarks/bench_join_agg_sum.py +++ /dev/null @@ -1,168 +0,0 @@ -"""Benchmark for the integer prefix-sum `join_agg` kernels (Issue #1648). - -Reproduces the shape of the issue's local numbers: - -- Kernel-level: the old Rust `compute_sum_start*`/`compute_sum_end*`/ - `compute_sum_start_end*` kernels (O(sum of interval widths)) versus the - new NumPy prefix-sum kernels in `_agg_functions` (O(n + m)), for - "suffixes" (`<`), "prefixes" (`>`), and "arbitrary intervals" (a range - join with no equality condition). -- End-to-end: a real `join_agg(..., aggfunc=[("value", "sum")])` call, - which now dispatches through the new prefix-sum kernels automatically. - -Run with: - - pixi run python benchmarks/bench_join_agg_sum.py - pixi run python benchmarks/bench_join_agg_sum.py --large # 10M/50M rows - -Large-N results (Apple Silicon, single run, int64), kernel-level only -- -the old Rust kernels are the O(sum of interval widths) algorithm this issue -replaces, so they are not re-run here; they are already intractable at the -20,000-row scale (see the small-N results above):: - - n = 10,000,000 - kernel suffixes (prefix-sum): 0.066 s - kernel prefixes (prefix-sum): 0.056 s - kernel intervals (prefix-sum): 0.098 s - end-to-end join_agg sum '<': 2.027 s (peak RSS ~1.6 GB) - - n = 50,000,000 - kernel suffixes (prefix-sum): 0.419 s - kernel prefixes (prefix-sum): 0.351 s - kernel intervals (prefix-sum): 0.652 s - end-to-end join_agg sum '<': 18.795 s (peak RSS ~6.1 GB) - -The prefix-sum kernel itself stays linear and sub-second through 50M rows. -The end-to-end number is dominated by `conditional_join`'s match-finding -(building `starts`/`right_index`), which is unrelated to this issue's -scope (the sum kernel) and unchanged by this PR. -""" - -import timeit - -import janitor_rs -import numpy as np -import pandas as pd - -import janitor # noqa: F401 registers the DataFrame accessor -from janitor.functions._conditional_join import _agg_functions as agg - - -def _bench(label: str, fn, number: int = 5) -> float: - """Time `fn`, print the per-call average, and return it in seconds.""" - seconds = timeit.timeit(fn, number=number) / number - print(f"{label:<28} {seconds * 1000:10.3f} ms") - return seconds - - -def kernel_level(n: int) -> None: - """Compare the old Rust kernels against the new prefix-sum kernels.""" - print(f"\n-- kernel-level, n={n:,} (int64) --") - rng = np.random.default_rng(0) - arr = rng.integers(-1000, 1000, size=n, dtype="int64") - booleans = np.zeros(n, dtype=bool) - - # suffixes: every row starts somewhere and runs to the end - starts = rng.integers(0, n, size=n, dtype="int64") - _bench( - "suffixes (rust)", - lambda: janitor_rs.compute_sum_start_int64( - arr=arr, starts=starts, booleans=booleans - ), - ) - _bench( - "suffixes (prefix-sum)", - lambda: agg._sum_starts(arr=arr, starts=starts, booleans=booleans), - ) - - # prefixes: every row runs from the start to somewhere - ends = rng.integers(0, n, size=n, dtype="int64") - _bench( - "prefixes (rust)", - lambda: janitor_rs.compute_sum_end_int64(arr=arr, ends=ends, booleans=booleans), - ) - _bench( - "prefixes (prefix-sum)", - lambda: agg._sum_ends(arr=arr, ends=ends, booleans=booleans), - ) - - # arbitrary intervals - lo = rng.integers(0, n, size=n, dtype="int64") - hi = rng.integers(0, n, size=n, dtype="int64") - starts2, ends2 = np.minimum(lo, hi), np.maximum(lo, hi) - _bench( - "intervals (rust)", - lambda: janitor_rs.compute_sum_start_end_int64( - arr=arr, starts=starts2, ends=ends2, booleans=booleans - ), - ) - _bench( - "intervals (prefix-sum)", - lambda: agg._sum_starts_ends( - arr=arr, starts=starts2, ends=ends2, booleans=booleans - ), - ) - - -def end_to_end(n: int) -> None: - """Time a real `join_agg(..., aggfunc=[("value", "sum")])` call.""" - print(f"\n-- end-to-end join_agg, n={n:,} (int64, '<') --") - rng = np.random.default_rng(0) - left = pd.DataFrame({"key": rng.integers(0, n, size=n)}) - right = pd.DataFrame( - { - "key": np.sort(rng.integers(0, n, size=n)), - "value": rng.integers(-1000, 1000, size=n), - } - ) - - def run(): - """Run one `join_agg` sum call.""" - return left.join_agg(right, ("key", "key", "<"), aggfunc=[("value", "sum")]) - - _bench("join_agg sum '<'", run, number=3) - - -def large_scale(n: int) -> None: - """Prefix-sum kernel only (the old Rust kernels are intractable here).""" - print(f"\n-- kernel-level, n={n:,} (int64), prefix-sum only --") - rng = np.random.default_rng(0) - arr = rng.integers(-1000, 1000, size=n, dtype="int64") - booleans = np.zeros(n, dtype=bool) - starts = rng.integers(0, n, size=n, dtype="int64") - ends = rng.integers(0, n, size=n, dtype="int64") - lo = rng.integers(0, n, size=n, dtype="int64") - hi = rng.integers(0, n, size=n, dtype="int64") - starts2, ends2 = np.minimum(lo, hi), np.maximum(lo, hi) - - _bench( - "suffixes (prefix-sum)", - lambda: agg._sum_starts(arr=arr, starts=starts, booleans=booleans), - number=3, - ) - _bench( - "prefixes (prefix-sum)", - lambda: agg._sum_ends(arr=arr, ends=ends, booleans=booleans), - number=3, - ) - _bench( - "intervals (prefix-sum)", - lambda: agg._sum_starts_ends( - arr=arr, starts=starts2, ends=ends2, booleans=booleans - ), - number=3, - ) - end_to_end(n) - - -if __name__ == "__main__": - import sys - - if "--large" in sys.argv: - for n in (10_000_000, 50_000_000): - large_scale(n) - else: - for n in (1_000, 20_000): - kernel_level(n) - for n in (1_000, 20_000): - end_to_end(n) From 7e4583a03f106b0055572588f84c91296dff487d Mon Sep 17 00:00:00 2001 From: samukweku Date: Sat, 22 Aug 2026 13:06:33 +1000 Subject: [PATCH 4/4] perf: adapt integer range sum dispatch --- CHANGELOG.md | 8 +- .../_conditional_join/_agg_functions.py | 78 +++++++++++++++---- .../test_conditional_join_agg_int_sum.py | 71 ++++++++++++++++- 3 files changed, 134 insertions(+), 23 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b2bb4dd33..c6c7d375a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,10 +1,10 @@ # Changelog ## [Unreleased] -- [PERF] Use an O(n + m) NumPy prefix-sum for integer `conditional_join` - range-sum aggregations (`join_agg(..., aggfunc=[(col, "sum")])`), - replacing the O(sum of interval widths) Rust kernels for the `<`-only, - `>`-only, and arbitrary-interval dispatch paths. - Issue #1648, PR #1673 @samukweku +- [PERF] Adaptively use an O(n + m) NumPy prefix-sum for dense integer + `conditional_join` range-sum aggregations + (`join_agg(..., aggfunc=[(col, "sum")])`), while retaining the Rust + kernels for sparse ranges. - Issue #1648, PR #1673 @samukweku - [ENH] Avoid copying column data during `conditional_join` input validation. - Issue #1645, PR #1642 @samukweku - [ENH] Speed up `conditional_join` with an unsorted right join key and diff --git a/janitor/functions/_conditional_join/_agg_functions.py b/janitor/functions/_conditional_join/_agg_functions.py index 6cbce5305..7787df76b 100644 --- a/janitor/functions/_conditional_join/_agg_functions.py +++ b/janitor/functions/_conditional_join/_agg_functions.py @@ -4,6 +4,21 @@ _INTEGER_DTYPE_NAMES = frozenset( {"int64", "int32", "int16", "int8", "uint64", "uint32", "uint16", "uint8"} ) +_PREFIX_SUM_WORK_FACTOR = 3 +# This many or fewer valid ranges cannot exceed the work threshold, because +# each range is at most ``arr.size`` elements wide. + + +def _use_prefix_sums(arr_size: int, total_width: int) -> bool: + """ + Use a prefix sum only when repeated range scans cost more. + + ELI5: Rust rereads every requested section. NumPy writes down one running + total for the whole array. Make that extra list only when rereading the + requested sections would mean walking the whole array more than three + times. + """ + return total_width > (_PREFIX_SUM_WORK_FACTOR * arr_size) def _int64_prefix_sums(arr: np.ndarray, booleans: np.ndarray) -> np.ndarray: @@ -33,10 +48,20 @@ def _sum_starts( Compute sum """ dtype_name = arr.dtype.name - if dtype_name in _INTEGER_DTYPE_NAMES: - prefix = _int64_prefix_sums(arr=arr, booleans=booleans) - return prefix[-1] - prefix[starts] + if dtype_name in _INTEGER_DTYPE_NAMES and starts.size > _PREFIX_SUM_WORK_FACTOR: + total_width = (arr.size * starts.size) - starts.sum(dtype=np.int64) + if _use_prefix_sums(arr_size=arr.size, total_width=total_width): + prefix = _int64_prefix_sums(arr=arr, booleans=booleans) + return prefix[-1] - prefix[starts] mapping = { + "int64": janitor_rs.compute_sum_start_int64, + "int32": janitor_rs.compute_sum_start_int32, + "int16": janitor_rs.compute_sum_start_int16, + "int8": janitor_rs.compute_sum_start_int8, + "uint64": janitor_rs.compute_sum_start_uint64, + "uint32": janitor_rs.compute_sum_start_uint32, + "uint16": janitor_rs.compute_sum_start_uint16, + "uint8": janitor_rs.compute_sum_start_uint8, "float64": janitor_rs.compute_sum_start_f64, "float32": janitor_rs.compute_sum_start_f32, } @@ -56,10 +81,20 @@ def _sum_ends( Compute sum """ dtype_name = arr.dtype.name - if dtype_name in _INTEGER_DTYPE_NAMES: - prefix = _int64_prefix_sums(arr=arr, booleans=booleans) - return prefix[ends] + if dtype_name in _INTEGER_DTYPE_NAMES and ends.size > _PREFIX_SUM_WORK_FACTOR: + total_width = ends.sum(dtype=np.int64) + if _use_prefix_sums(arr_size=arr.size, total_width=total_width): + prefix = _int64_prefix_sums(arr=arr, booleans=booleans) + return prefix[ends] mapping = { + "int64": janitor_rs.compute_sum_end_int64, + "int32": janitor_rs.compute_sum_end_int32, + "int16": janitor_rs.compute_sum_end_int16, + "int8": janitor_rs.compute_sum_end_int8, + "uint64": janitor_rs.compute_sum_end_uint64, + "uint32": janitor_rs.compute_sum_end_uint32, + "uint16": janitor_rs.compute_sum_end_uint16, + "uint8": janitor_rs.compute_sum_end_uint8, "float64": janitor_rs.compute_sum_end_f64, "float32": janitor_rs.compute_sum_end_f32, } @@ -812,17 +847,28 @@ def _sum_starts_ends( Compute sum """ dtype_name = arr.dtype.name - if dtype_name in _INTEGER_DTYPE_NAMES: - prefix = _int64_prefix_sums(arr=arr, booleans=booleans) - result = prefix[ends] - prefix[starts] - # an empty (or inverted) range contributes nothing, matching the - # Rust `for nn in start_..end_` loop, which never iterates when - # start_ >= end_ - empty_range = starts >= ends - if empty_range.any(): - result[empty_range] = 0 - return result + if dtype_name in _INTEGER_DTYPE_NAMES and starts.size > _PREFIX_SUM_WORK_FACTOR: + widths = np.maximum(ends - starts, 0) + total_width = widths.sum(dtype=np.int64) + if _use_prefix_sums(arr_size=arr.size, total_width=total_width): + prefix = _int64_prefix_sums(arr=arr, booleans=booleans) + result = prefix[ends] - prefix[starts] + # an empty (or inverted) range contributes nothing, matching the + # Rust `for nn in start_..end_` loop, which never iterates when + # start_ >= end_ + empty_range = starts >= ends + if empty_range.any(): + result[empty_range] = 0 + return result mapping = { + "int64": janitor_rs.compute_sum_start_end_int64, + "int32": janitor_rs.compute_sum_start_end_int32, + "int16": janitor_rs.compute_sum_start_end_int16, + "int8": janitor_rs.compute_sum_start_end_int8, + "uint64": janitor_rs.compute_sum_start_end_uint64, + "uint32": janitor_rs.compute_sum_start_end_uint32, + "uint16": janitor_rs.compute_sum_start_end_uint16, + "uint8": janitor_rs.compute_sum_start_end_uint8, "float64": janitor_rs.compute_sum_start_end_f64, "float32": janitor_rs.compute_sum_start_end_f32, } diff --git a/tests/functions/test_conditional_join_agg_int_sum.py b/tests/functions/test_conditional_join_agg_int_sum.py index 98396ecf7..47d8025f8 100644 --- a/tests/functions/test_conditional_join_agg_int_sum.py +++ b/tests/functions/test_conditional_join_agg_int_sum.py @@ -2,9 +2,9 @@ Covers `_int64_prefix_sums`, `_sum_starts`, `_sum_ends`, and `_sum_starts_ends` in `janitor.functions._conditional_join._agg_functions` -(Issue #1648) -- the O(n + m) NumPy replacements for the Rust -`compute_sum_start*`, `compute_sum_end*`, and `compute_sum_start_end*` -kernels, for integer dtypes only. +(Issue #1648) -- the O(n + m) NumPy path for dense integer ranges; sparse +ranges continue to use the Rust `compute_sum_start*`, `compute_sum_end*`, and +`compute_sum_start_end*` kernels. """ import numpy as np @@ -183,3 +183,68 @@ def test_starts_beyond_ends_is_empty_range(): arr=arr, starts=starts, ends=ends, booleans=booleans ) assert result[0] == 0 + + +@pytest.mark.parametrize( + "func_name,rust_name,indexers", + [ + ("_sum_starts", "compute_sum_start_int64", {"starts": [99]}), + ("_sum_ends", "compute_sum_end_int64", {"ends": [1]}), + ( + "_sum_starts_ends", + "compute_sum_start_end_int64", + {"starts": [50], "ends": [51]}, + ), + ], +) +def test_sparse_ranges_use_rust(monkeypatch, func_name, rust_name, indexers): + """Selective ranges should not pay to scan and copy the full array.""" + expected = np.array([123], dtype=np.int64) + + def fake_rust(**kwargs): + return expected + + monkeypatch.setattr(_agg_functions.janitor_rs, rust_name, fake_rust) + monkeypatch.setattr( + _agg_functions, + "_int64_prefix_sums", + lambda **kwargs: pytest.fail("sparse ranges should stay in Rust"), + ) + actual = getattr(_agg_functions, func_name)( + arr=np.ones(100, dtype=np.int64), + booleans=np.zeros(100, dtype=bool), + **{name: np.array(values, dtype=np.int64) for name, values in indexers.items()}, + ) + + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize( + "func_name,rust_family,indexers", + [ + ("_sum_starts", "compute_sum_start", {"starts": [0, 0, 0, 0]}), + ("_sum_ends", "compute_sum_end", {"ends": [100, 100, 100, 100]}), + ( + "_sum_starts_ends", + "compute_sum_start_end", + {"starts": [0, 0, 0, 0], "ends": [100, 100, 100, 100]}, + ), + ], +) +@pytest.mark.parametrize("dtype", INTEGER_DTYPES) +def test_dense_ranges_use_prefix_sums( + monkeypatch, func_name, rust_family, indexers, dtype +): + """Heavily overlapping ranges should use one shared running total.""" + + def fail_rust(**kwargs): + pytest.fail("dense ranges should use prefix sums") + + monkeypatch.setattr(_agg_functions.janitor_rs, f"{rust_family}_{dtype}", fail_rust) + actual = getattr(_agg_functions, func_name)( + arr=np.ones(100, dtype=dtype), + booleans=np.zeros(100, dtype=bool), + **{name: np.array(values, dtype=np.int64) for name, values in indexers.items()}, + ) + + np.testing.assert_array_equal(actual, np.full(4, 100, dtype=np.int64))