diff --git a/CHANGELOG.md b/CHANGELOG.md index e464228b9..c6c7d375a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,10 @@ # Changelog ## [Unreleased] +- [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 96fa0b7a2..7787df76b 100644 --- a/janitor/functions/_conditional_join/_agg_functions.py +++ b/janitor/functions/_conditional_join/_agg_functions.py @@ -1,6 +1,43 @@ import janitor_rs import numpy as np +_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: + """ + 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,6 +47,12 @@ def _sum_starts( """ Compute sum """ + dtype_name = arr.dtype.name + 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, @@ -22,7 +65,6 @@ def _sum_starts( "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,6 +80,12 @@ def _sum_ends( """ Compute sum """ + dtype_name = arr.dtype.name + 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, @@ -50,7 +98,6 @@ def _sum_ends( "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,6 +846,20 @@ def _sum_starts_ends( """ Compute sum """ + dtype_name = arr.dtype.name + 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, @@ -811,7 +872,6 @@ def _sum_starts_ends( "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 +946,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..47d8025f8 --- /dev/null +++ b/tests/functions/test_conditional_join_agg_int_sum.py @@ -0,0 +1,250 @@ +"""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 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 +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 + + +@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))