diff --git a/CHANGELOG.md b/CHANGELOG.md index e464228b9..71f7428e0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,14 @@ # Changelog ## [Unreleased] +- [PERF] Replace the forward `float32`/`float64` range-`sum` kernels behind + `join_agg` (suffix, prefix, and arbitrary-interval ranges) with a + Neumaier-compensated prefix sum when the total queried width is large + enough relative to the array for the one-time O(n) build to pay off; + otherwise, or when the whole array can't be trusted for prefix + subtraction (a real +/-inf, an overflowing partial sum, or extreme + within-array dynamic range), falls back to the existing Rust kernel. + - Issue #1671, PR #1675 @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..19931ed10 100644 --- a/janitor/functions/_conditional_join/_agg_functions.py +++ b/janitor/functions/_conditional_join/_agg_functions.py @@ -1,6 +1,159 @@ import janitor_rs import numpy as np +# float32/float64 sum ranges are computed with a Neumaier-compensated +# prefix sum instead of a Rust round-trip - see _compensated_prefix_sum. +_FLOAT_PREFIX_DTYPES = frozenset(("float64", "float32")) + + +def _compensated_prefix_sum( + arr: np.ndarray, + booleans: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + """ + Build a Neumaier-compensated running-sum prefix, upcast to float64. + + ELI5: a plain running total loses precision when a tiny number gets + added to a much bigger one - the tiny part just gets rounded away. + This keeps a second "leftover" running total alongside the main one, + tracking what got rounded away at each step, so that a later + `prefix[end] - prefix[start]` subtraction stays close to summing that + slice directly instead of drifting the way a naive running total would. + `prefix[0]` is `(0.0, 0.0)`; `prefix[i]` holds the compensated sum of + the first `i` elements (with null positions treated as `0.0`). + """ + values = arr.astype(np.float64, copy=False) + n = values.size + hi = np.empty(n + 1, dtype=np.float64) + lo = np.empty(n + 1, dtype=np.float64) + hi[0] = 0.0 + lo[0] = 0.0 + total = 0.0 + compensation = 0.0 + # A running total near +/-inf can legitimately overflow, and the + # caller (_build_prefix_if_safe) already detects and handles that - + # silence the transient overflow/NaN warnings that would otherwise + # surface for a case this code correctly falls back on. + with np.errstate(over="ignore", invalid="ignore"): + for pos in range(n): + value = 0.0 if booleans[pos] else values[pos] + new_total = total + value + if abs(total) >= abs(value): + compensation += (total - new_total) + value + else: + compensation += (value - new_total) + total + total = new_total + hi[pos + 1] = total + lo[pos + 1] = compensation + return hi, lo + + +def _range_sum_from_prefix( + hi: np.ndarray, + lo: np.ndarray, + starts, + ends, +) -> np.ndarray: + """ + Answer [start, end) range-sum queries from a compensated prefix. + + ELI5: each prefix position has a main total (`hi`) and a leftovers + total (`lo`). Subtract both ledgers at `start` from both at `end`, then + combine what remains to recover the requested slice. + + This need not be bit-for-bit equal to Rust summing the slice afresh: + Kahan summation and compensated-prefix subtraction can round + cancellation differently. Correctness is instead bounded against a + high-precision `math.fsum` reference by a small multiple of float64 + epsilon times `sum(abs(inputs))`. Scaling by the input magnitudes is + intentional; dividing only by a cancellation-heavy, near-zero answer + would make two accurate results appear arbitrarily far apart. + """ + return (hi[ends] - hi[starts]) + (lo[ends] - lo[starts]) + + +# Neumaier compensation is itself just a float64, so once it has absorbed +# a moderate-magnitude correction (say ~1e12) it can no longer represent a +# *later* tiny correction (say ~1e-6) - that increment underflows below +# the compensation term's own ULP. Empirically (see PR discussion), a big +# excursion followed by another huge, opposite-signed excursion followed +# by tiny values can silently lose 100% of a small window's true value, +# well before the running total itself ever overflows. Random dynamic +# range alone stays safe under the scale-aware accuracy contract exercised +# against `math.fsum` in the tests. A result-relative error is deliberately +# not used: cancellation can make the true result arbitrarily close to zero +# even when both summation methods are accurate. The 1e15 cutoff remains +# several orders of magnitude below the observed catastrophic regime. +_MAX_SAFE_DYNAMIC_RANGE = 1e15 + + +def _has_safe_dynamic_range(arr: np.ndarray, booleans: np.ndarray) -> bool: + """ + Whether the non-null magnitudes in `arr` are too spread out to trust. + + ELI5: the leftover ledger is also a float. If the biggest value is far + too large compared with the smallest, that ledger can lose the small + correction it was created to protect, so the caller uses Rust instead. + """ + non_null = arr if not booleans.any() else arr[~booleans] + magnitudes = np.abs(non_null.astype(np.float64, copy=False)) + magnitudes = magnitudes[magnitudes > 0] + if magnitudes.size <= 1: + return True + with np.errstate(over="ignore"): + threshold = magnitudes.min() * _MAX_SAFE_DYNAMIC_RANGE + return bool(magnitudes.max() <= threshold) + + +def _build_prefix_if_safe( + arr: np.ndarray, + booleans: np.ndarray, +) -> tuple[np.ndarray, np.ndarray] | None: + """ + Build the compensated prefix, or signal it isn't safe to use. + + Two distinct ways the fast path can silently go wrong, both handled + here by falling back to Rust (which recomputes each range from + scratch and so isn't exposed to either failure mode): + + 1. A real +/-inf already in `arr` (which then poisons every later + `hi` entry, since finite + inf stays infinite forever), or two + ordinary finite values whose partial sum genuinely overflows + float64 range. Either way, once `hi` holds a +/-inf, subtracting + two such entries for a range that never touched the offending + value(s) produces `inf - inf = NaN`. Checked by requiring every + `hi` entry beyond the leading zero to be finite. + 2. Extreme dynamic range within `arr` - see `_MAX_SAFE_DYNAMIC_RANGE`. + + Passing these guards means the prefix satisfies the scale-aware + accuracy contract described in `_range_sum_from_prefix`; it does not + promise identical rounding to Rust's per-range Kahan implementation. + """ + if not _has_safe_dynamic_range(arr=arr, booleans=booleans): + return None + hi, lo = _compensated_prefix_sum(arr=arr, booleans=booleans) + if not np.isfinite(hi).all(): + return None + return hi, lo + + +# Building the prefix costs ~250ns/element in a pure Python loop, vs. +# ~2ns/element for Rust's native per-range scan (measured directly) - +# roughly 120x slower per element. The one-time O(n) build only pays off +# once the *total* width summed across every query range exceeds that +# same ratio times the array size; below that, Rust doing direct O(width) +# work per range is cheaper overall, even though it re-scans overlapping +# regions. A margin above the measured ~120x ratio keeps this conservative +# (i.e. biased toward Rust when it's a close call). +_MIN_TOTAL_WIDTH_RATIO = 150 + + +def _prefix_sum_is_worthwhile(array_size: int, total_width: int) -> bool: + """ + Whether the one-time O(n) prefix build is expected to pay off overall. + """ + return total_width > _MIN_TOTAL_WIDTH_RATIO * array_size + def _sum_starts( arr: np.ndarray, @@ -8,8 +161,24 @@ def _sum_starts( booleans: np.ndarray, ) -> tuple: """ - Compute sum + Sum each suffix selected by `starts`. + + ELI5: a suffix is "everything from here onward." For every + `starts[i]`, this computes `arr[starts[i]:]`, skipping nulls. When many + suffixes overlap, the compensated prefix lets us subtract the saved + total before `start` from the saved total at the end instead of adding + the same tail repeatedly. """ + dtype_name = arr.dtype.name + if dtype_name in _FLOAT_PREFIX_DTYPES: + total_width = starts.size * arr.size - int(starts.sum()) + if _prefix_sum_is_worthwhile(arr.size, total_width): + prefix = _build_prefix_if_safe(arr=arr, booleans=booleans) + if prefix is not None: + hi, lo = prefix + return _range_sum_from_prefix( + hi=hi, lo=lo, starts=starts, ends=arr.size + ) mapping = { "int64": janitor_rs.compute_sum_start_int64, "int32": janitor_rs.compute_sum_start_int32, @@ -22,7 +191,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: @@ -36,8 +204,21 @@ def _sum_ends( booleans: np.ndarray, ) -> tuple: """ - Compute sum + Sum each prefix selected by `ends`. + + ELI5: a prefix is "everything from the beginning up to here." For + every `ends[i]`, this computes `arr[:ends[i]]`, skipping nulls. A + compensated prefix has already saved exactly that running total, so a + dense workload can answer each query with a lookup. """ + dtype_name = arr.dtype.name + if dtype_name in _FLOAT_PREFIX_DTYPES: + total_width = int(ends.sum()) + if _prefix_sum_is_worthwhile(arr.size, total_width): + prefix = _build_prefix_if_safe(arr=arr, booleans=booleans) + if prefix is not None: + hi, lo = prefix + return _range_sum_from_prefix(hi=hi, lo=lo, starts=0, ends=ends) mapping = { "int64": janitor_rs.compute_sum_end_int64, "int32": janitor_rs.compute_sum_end_int32, @@ -50,7 +231,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: @@ -797,8 +977,20 @@ def _sum_starts_ends( booleans: np.ndarray, ) -> tuple: """ - Compute sum + Sum each arbitrary interval selected by `starts` and `ends`. + + ELI5: `[start:end)` means "begin at `start`, stop just before `end`." + Subtracting the saved prefix at `start` from the saved prefix at `end` + removes everything before the interval and leaves only its sum. """ + dtype_name = arr.dtype.name + if dtype_name in _FLOAT_PREFIX_DTYPES: + total_width = int((ends - starts).sum()) + if _prefix_sum_is_worthwhile(arr.size, total_width): + prefix = _build_prefix_if_safe(arr=arr, booleans=booleans) + if prefix is not None: + hi, lo = prefix + return _range_sum_from_prefix(hi=hi, lo=lo, starts=starts, ends=ends) mapping = { "int64": janitor_rs.compute_sum_start_end_int64, "int32": janitor_rs.compute_sum_start_end_int32, @@ -811,7 +1003,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 +1077,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_prefix_sum_float_agg.py b/tests/functions/test_conditional_join_prefix_sum_float_agg.py new file mode 100644 index 000000000..e072e2365 --- /dev/null +++ b/tests/functions/test_conditional_join_prefix_sum_float_agg.py @@ -0,0 +1,467 @@ +"""Tests for the compensated-prefix-sum float path in join_agg. + +See issue #1671: the forward `sum` range kernels for float32/float64 use +a Neumaier-compensated prefix sum instead of a Rust round-trip per range, +gated by a work-estimate heuristic (`_prefix_sum_is_worthwhile`) since the +one-time O(n) build only pays off when the total width of the queried +ranges is large relative to the array size - otherwise Rust's per-range +native scan is cheaper overall. + +Numerical-correctness tests below exercise the compensated-prefix math +directly via `_prefix_range_sum` (bypassing the performance heuristic, +which is orthogonal to correctness) against `math.fsum` as a high-precision +reference. The error budget scales with the sum of the input magnitudes, +not the possibly near-zero result: cancellation makes result-relative error +ill-conditioned and can make two accurate summation methods look arbitrarily +far apart. Separate tests cover the heuristic itself and full `join_agg` +integration at a scale where the fast path is actually selected. +""" + +import math + +import janitor_rs +import numpy as np +import pandas as pd +import pytest +from pandas.testing import assert_frame_equal + +import janitor # noqa: F401 - registers the join_agg DataFrame accessor +from janitor.functions._conditional_join import _agg_functions as af + +ROUNDING_ERROR_FACTOR = 8 +# Compensated sums should stay within a small multiple of machine epsilon +# times sum(abs(inputs)); eight leaves headroom for prefix construction plus +# range subtraction while still catching a lost contribution. + +RUST_SUM_STARTS = { + "float64": janitor_rs.compute_sum_start_f64, + "float32": janitor_rs.compute_sum_start_f32, +} +RUST_SUM_ENDS = { + "float64": janitor_rs.compute_sum_end_f64, + "float32": janitor_rs.compute_sum_end_f32, +} +RUST_SUM_STARTS_ENDS = { + "float64": janitor_rs.compute_sum_start_end_f64, + "float32": janitor_rs.compute_sum_start_end_f32, +} + + +def _assert_accurate(actual, arr, booleans, start, end): + """Compare with a high-precision sum using a cancellation-safe budget.""" + values = arr[start:end] + mask = booleans[start:end] + values = values[~mask] + exact = math.fsum(float(value) for value in values) + scale = math.fsum(abs(float(value)) for value in values) + error_bound = ROUNDING_ERROR_FACTOR * np.finfo(np.float64).eps * scale + assert abs(float(actual) - exact) <= error_bound, ( + actual, + exact, + error_bound, + start, + end, + ) + + +def _prefix_range_sum(arr, booleans, starts, ends): + """Exercise the compensated-prefix path directly, bypassing the + performance heuristic in _sum_starts/_sum_ends/_sum_starts_ends - + for tests about the math, not about when it's chosen. Returns None + if the safety guard legitimately rejects this array (e.g. dynamic + range right at the threshold, which randomized test data can land on + by chance) - callers should skip that trial rather than treat it as + a failure.""" + prefix = af._build_prefix_if_safe(arr=arr, booleans=booleans) + if prefix is None: + return None + hi, lo = prefix + return af._range_sum_from_prefix(hi=hi, lo=lo, starts=starts, ends=ends) + + +@pytest.mark.parametrize("dtype", ["float64", "float32"]) +def test_sum_starts_ends_matches_high_precision_reference(dtype): + """Randomized, mixed-magnitude ranges stay within the accuracy contract.""" + rng = np.random.default_rng(20260822) + rust_fn = RUST_SUM_STARTS_ENDS[dtype] + checked = 0 + for n in (1, 2, 5, 37, 500, 999, 4096, 10_007): + for _ in range(25): + arr = rng.standard_normal(n).astype(dtype) + scale = rng.choice([1e-6, 1, 1e6, 1e9], size=n).astype(dtype) + arr = (arr * scale).astype(dtype) + booleans = rng.random(n) < 0.05 + start = int(rng.integers(0, n + 1)) + end = int(rng.integers(start, n + 1)) + starts = np.array([start], dtype=np.int64) + ends = np.array([end], dtype=np.int64) + + actual = _prefix_range_sum(arr, booleans, starts, ends) + if actual is None: + continue # safety guard legitimately rejected this draw + baseline = rust_fn(arr=arr, starts=starts, ends=ends, booleans=booleans)[0] + _assert_accurate(actual[0], arr, booleans, start, end) + _assert_accurate(baseline, arr, booleans, start, end) + checked += 1 + assert checked > 0 + + +@pytest.mark.parametrize("dtype", ["float64", "float32"]) +def test_sum_starts_matches_high_precision_reference(dtype): + """Suffix sums (single '<' join) stay within the accuracy contract.""" + rng = np.random.default_rng(20260823) + rust_fn = RUST_SUM_STARTS[dtype] + checked = 0 + for n in (1, 5, 500, 4096): + for _ in range(10): + arr = rng.standard_normal(n).astype(dtype) + scale = rng.choice([1e-6, 1, 1e6], size=n).astype(dtype) + arr = (arr * scale).astype(dtype) + booleans = rng.random(n) < 0.05 + starts = rng.integers(0, n + 1, size=3).astype(np.int64) + + actual = _prefix_range_sum(arr, booleans, starts, arr.size) + if actual is None: + continue + baseline = rust_fn(arr=arr, starts=starts, booleans=booleans) + for start, value, rust_value in zip(starts, actual, baseline): + _assert_accurate(value, arr, booleans, int(start), arr.size) + _assert_accurate(rust_value, arr, booleans, int(start), arr.size) + checked += 1 + assert checked > 0 + + +@pytest.mark.parametrize("dtype", ["float64", "float32"]) +def test_sum_ends_matches_high_precision_reference(dtype): + """Prefix sums (single '>' join) stay within the accuracy contract.""" + rng = np.random.default_rng(20260824) + rust_fn = RUST_SUM_ENDS[dtype] + checked = 0 + for n in (1, 5, 500, 4096): + for _ in range(10): + arr = rng.standard_normal(n).astype(dtype) + scale = rng.choice([1e-6, 1, 1e6], size=n).astype(dtype) + arr = (arr * scale).astype(dtype) + booleans = rng.random(n) < 0.05 + ends = rng.integers(0, n + 1, size=3).astype(np.int64) + + actual = _prefix_range_sum(arr, booleans, 0, ends) + if actual is None: + continue + baseline = rust_fn(arr=arr, ends=ends, booleans=booleans) + for end, value, rust_value in zip(ends, actual, baseline): + _assert_accurate(value, arr, booleans, 0, int(end)) + _assert_accurate(rust_value, arr, booleans, 0, int(end)) + checked += 1 + assert checked > 0 + + +def test_moderate_dynamic_range_still_uses_fast_path(): + """The dynamic-range guard shouldn't be so conservative that it defeats + the point: everyday mixed-magnitude data must still take the fast + path, not fall back to Rust for every call. Magnitudes are bounded + away from zero (unlike a raw Gaussian*scale draw, whose tail can land + arbitrarily close to zero by chance and blow out the *empirical* + dynamic range regardless of the intended scale spread).""" + rng = np.random.default_rng(2) + magnitudes = rng.uniform(1.0, 1e9, size=500) + signs = rng.choice([-1.0, 1.0], size=500) + arr = signs * magnitudes + booleans = np.zeros(arr.size, dtype=np.bool_) + assert af._has_safe_dynamic_range(arr, booleans) + + +@pytest.mark.parametrize("dtype", [np.float64, np.float32]) +def test_well_scaled_arrays_match_rust_exactly(dtype): + """Same order of magnitude throughout -> no meaningful cancellation, + so the compensated prefix sum should agree with Rust bit-for-bit.""" + rng = np.random.default_rng(1) + arr = rng.standard_normal(2000).astype(dtype) + booleans = np.zeros(arr.size, dtype=np.bool_) + dtype_name = np.dtype(dtype).name + starts = np.array([0, 5, 1000], dtype=np.int64) + ends = np.array([2000, 1800, 1999], dtype=np.int64) + + expected = RUST_SUM_STARTS_ENDS[dtype_name]( + arr=arr, starts=starts, ends=ends, booleans=booleans + ) + actual = _prefix_range_sum(arr, booleans, starts, ends) + np.testing.assert_array_equal(actual, expected) + + +def test_query_multiplicity_preserves_accuracy_across_dispatch_paths(): + """Adding duplicate queries may select the prefix path, but both paths + must remain accurate even when cancellation makes their rounded results + differ. This guards the case found during review: Rust's per-range Kahan + sum returns -0.0999755859375, while the prefix path returns the more + accurate -0.1 once enough duplicate ranges cross the work threshold.""" + right = pd.DataFrame({"key": [0.0, 1.0, 2.0], "value": [-0.1, -1e12, 1e12]}) + one = pd.DataFrame({"key": [-1.0]}) + many = pd.DataFrame({"key": np.full(151, -1.0)}) + + direct = one.join_agg(right, ("key", "key", "<"), aggfunc=[("value", "sum")]).iloc[ + 0, 0 + ] + prefix = many.join_agg(right, ("key", "key", "<"), aggfunc=[("value", "sum")]).iloc[ + 0, 0 + ] + + arr = right["value"].to_numpy() + booleans = np.zeros(arr.size, dtype=np.bool_) + assert not af._prefix_sum_is_worthwhile(arr.size, arr.size) + assert af._prefix_sum_is_worthwhile(arr.size, 151 * arr.size) + _assert_accurate(direct, arr, booleans, 0, arr.size) + _assert_accurate(prefix, arr, booleans, 0, arr.size) + + +@pytest.mark.parametrize("dtype", [np.float64, np.float32]) +def test_null_mask_is_authoritative(dtype): + """A `booleans=True` position is ignored regardless of its raw value - + matches the Rust kernels' verified null-handling semantics.""" + arr = np.array([5.0, 999.0, 5.0], dtype=dtype) + booleans = np.array([False, True, False]) + out = _prefix_range_sum(arr, booleans, 0, arr.size) + assert out == pytest.approx(10.0) + + +@pytest.mark.parametrize("dtype", [np.float64, np.float32]) +def test_all_null_window_sums_to_zero(dtype): + """An all-null window sums to 0.0, not NaN (matches Rust).""" + arr = np.array([1.0, 2.0, 3.0], dtype=dtype) + booleans = np.array([True, True, True]) + out = _prefix_range_sum(arr, booleans, 0, arr.size) + assert out == 0.0 + + +@pytest.mark.parametrize("dtype", [np.float64, np.float32]) +def test_empty_range_sums_to_zero(dtype): + """start == end -> an empty range sums to 0.0.""" + arr = np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=dtype) + booleans = np.zeros(arr.size, dtype=np.bool_) + out = _prefix_range_sum(arr, booleans, 2, 2) + assert out == 0.0 + + +def test_compensated_sum_more_accurate_than_naive_cumsum(): + """Sanity check that compensation is actually doing something: a + repeated-0.1 sum should land on the exact value a naive cumsum misses.""" + n = 100_000 + arr = np.full(n, 0.1, dtype=np.float64) + booleans = np.zeros(n, dtype=np.bool_) + + naive = np.cumsum(arr)[-1] + compensated = _prefix_range_sum(arr, booleans, 0, n) + + assert naive != 10_000.0 # confirms the test actually stresses precision + assert compensated == 10_000.0 + + +def test_int_dtypes_still_use_rust_path(): + """Integer dtypes are untouched by this change (tracked in #1648).""" + arr = np.array([1, 2, 3], dtype=np.int64) + booleans = np.zeros(3, dtype=np.bool_) + starts = np.array([0], dtype=np.int64) + out = af._sum_starts(arr=arr, starts=starts, booleans=booleans) + assert out[0] == 6 + + +def test_overflow_during_summation_is_rejected_by_safety_guard(): + """Regression test: two large-but-finite values whose sum overflows + float64 range must saturate to +/-inf (matching NumPy's own overflow + behavior and Rust), not produce a spurious NaN from the compensation + bookkeeping's inf - inf. Found via hypothesis property tests. + `_build_prefix_if_safe` must refuse to hand back a prefix built over + such an array, so callers fall back to Rust.""" + arr = np.array([-8.988466e307, -8.988466e307], dtype=np.float64) + booleans = np.zeros(2, dtype=np.bool_) + starts = np.array([0], dtype=np.int64) + ends = np.array([2], dtype=np.int64) + + assert af._build_prefix_if_safe(arr=arr, booleans=booleans) is None + expected = RUST_SUM_STARTS_ENDS["float64"]( + arr=arr, starts=starts, ends=ends, booleans=booleans + ) + actual = af._sum_starts_ends(arr=arr, starts=starts, ends=ends, booleans=booleans) + assert actual[0] == expected[0] == -np.inf + + +def test_extreme_dynamic_range_is_rejected_by_safety_guard(): + """Regression test: Neumaier compensation is itself just a float64, so + once it has absorbed a moderate correction it can no longer represent + a much later, much tinier one - that increment underflows below the + compensation term's own ULP. A big excursion, another huge opposite- + signed excursion, then a tiny value can silently lose 100% of that + tiny value's true contribution, well before anything overflows to + +/-inf (so the finiteness guard alone doesn't catch it). Found via + randomized adversarial testing at extreme magnitude spreads + (~1e99 down to ~1e-6) - a single-element window recovered 0.0 instead + of the true ~1.06e-6. `_build_prefix_if_safe` must refuse this array.""" + arr = np.array( + [ + 3.39748102e99, + 1.98320141e12, + -6.32453040e99, + 5.40393231e-07, + 1.06052966e-06, + ], + dtype=np.float64, + ) + booleans = np.zeros(arr.size, dtype=np.bool_) + starts = np.array([4], dtype=np.int64) + ends = np.array([5], dtype=np.int64) + + assert not af._has_safe_dynamic_range(arr, booleans) + assert af._build_prefix_if_safe(arr=arr, booleans=booleans) is None + expected = RUST_SUM_STARTS_ENDS["float64"]( + arr=arr, starts=starts, ends=ends, booleans=booleans + ) + actual = af._sum_starts_ends(arr=arr, starts=starts, ends=ends, booleans=booleans) + assert actual[0] == expected[0] == arr[4] + + +@pytest.mark.parametrize("dtype", [np.float64, np.float32]) +def test_range_excluding_earlier_infinity_matches_rust(dtype): + """Regression test: once the running total hits +/-inf, naively + subtracting two +/-inf prefix entries for a later range that never + touched the infinity yields NaN. A range that excludes an earlier + infinity must still fall back to Rust and get the correct finite sum.""" + arr = np.array([1.0, -np.inf, 2.0, 3.0], dtype=dtype) + booleans = np.zeros(arr.size, dtype=np.bool_) + starts = np.array([2], dtype=np.int64) + ends = np.array([4], dtype=np.int64) + + assert af._build_prefix_if_safe(arr=arr, booleans=booleans) is None + expected = RUST_SUM_STARTS_ENDS[np.dtype(dtype).name]( + arr=arr, starts=starts, ends=ends, booleans=booleans + ) + actual = af._sum_starts_ends(arr=arr, starts=starts, ends=ends, booleans=booleans) + assert np.isfinite(actual[0]) + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize("dtype", [np.float64, np.float32]) +def test_masked_infinity_still_passes_safety_guard(dtype): + """An infinity at a null (`booleans=True`) position doesn't trip the + safety guard, since it never contributes to the compensated prefix.""" + arr = np.array([1.0, np.inf, 2.0], dtype=dtype) + booleans = np.array([False, True, False]) + out = _prefix_range_sum(arr, booleans, 0, arr.size) + assert out == pytest.approx(3.0) + + +def test_prefix_sum_is_worthwhile_boundary(): + """Unit-tests the work-estimate heuristic directly at its threshold: + the O(n) build should only be judged worthwhile once the total + queried width clears `_MIN_TOTAL_WIDTH_RATIO * array_size`.""" + array_size = 1000 + threshold = af._MIN_TOTAL_WIDTH_RATIO * array_size + assert not af._prefix_sum_is_worthwhile(array_size, threshold) + assert af._prefix_sum_is_worthwhile(array_size, threshold + 1) + + +def test_small_query_count_against_large_array_defers_to_rust(): + """The scenario a prior review flagged as a performance regression: a + single (or few) query against a large array. The one-time O(n) Python + build (~250ns/element) would be far slower than Rust's native + per-range scan (~2ns/element) here, so this must defer to Rust rather + than pay the build cost for essentially no reuse.""" + n = 200_000 + arr = np.random.default_rng(5).standard_normal(n) + booleans = np.zeros(n, dtype=np.bool_) + starts = np.array([0], dtype=np.int64) + + total_width = n # a single full-width suffix query + assert not af._prefix_sum_is_worthwhile(n, total_width) + + expected = RUST_SUM_STARTS["float64"](arr=arr, starts=starts, booleans=booleans) + actual = af._sum_starts(arr=arr, starts=starts, booleans=booleans) + np.testing.assert_array_equal(actual, expected) + + +def test_join_agg_float_sum_end_to_end(): + """Small, deterministic end-to-end join_agg check for float `sum`, + covering suffix ('<' -> _sum_starts), prefix ('>' -> _sum_ends), and + range-join (-> _sum_starts_ends) aggregation. Deliberately uses + well-scaled values rather than hypothesis-generated extremes: Rust's + Kahan kernels are already known (independent of this change) to + disagree with plain pandas/NumPy summation right at float64's overflow + boundary, so pandas is only a valid oracle away from that edge. This + dataset is far too small to trigger the fast path (see + test_join_agg_large_overlapping_ranges_uses_fast_path below for that), + but it exercises the full join_agg -> _agg_join_right -> _sum_* wiring.""" + left = pd.DataFrame({"a": [3, 7, 1], "b": [9, 2, 6]}) + right = pd.DataFrame( + { + "x": [1, 3, 5, 7, 9, 11], + "y": [0, 2, 4, 6, 8, 10], + "value": [0.1, 0.2, 0.3, 0.4, 0.5, 0.6], + } + ) + + expected_lt = ( + left.reset_index(names="l") + .merge(right, how="cross") + .query("a < x") + .groupby("l") + .agg({"value": ["sum"]}) + ) + expected_lt.index.names = [None] + actual_lt = left.join_agg(right, ("a", "x", "<"), aggfunc=[("value", "sum")]) + assert_frame_equal(expected_lt, actual_lt) + + expected_gt = ( + left.reset_index(names="l") + .merge(right, how="cross") + .query("a > x") + .groupby("l") + .agg({"value": ["sum"]}) + ) + expected_gt.index.names = [None] + actual_gt = left.join_agg(right, ("a", "x", ">"), aggfunc=[("value", "sum")]) + assert_frame_equal(expected_gt, actual_gt) + + expected_range = ( + left.reset_index(names="l") + .merge(right, how="cross") + .query("b > y and a < x") + .groupby("l") + .agg({"value": ["sum"]}) + ) + expected_range.index.names = [None] + actual_range = left.join_agg( + right, ("b", "y", ">"), ("a", "x", "<"), aggfunc=[("value", "sum")] + ) + actual_range = actual_range.loc[expected_range.index] + assert_frame_equal(expected_range, actual_range) + + +def test_join_agg_large_overlapping_ranges_uses_fast_path(): + """Genuine end-to-end integration coverage of the fast path itself: + a large right table with many overlapping suffix ranges (the case + #1671 targets) should both trigger `_prefix_sum_is_worthwhile` and + produce results matching a direct, per-row Rust computation.""" + n_right = 5000 + n_left = 2000 + rng = np.random.default_rng(11) + left = pd.DataFrame({"key": rng.uniform(0, n_right, size=n_left)}) + right = pd.DataFrame( + { + "key": np.arange(n_right, dtype=float), + "value": rng.standard_normal(n_right), + } + ) + + total_width = int(n_left * n_right - left["key"].to_numpy().sum()) + assert af._prefix_sum_is_worthwhile(n_right, total_width) + + actual = left.join_agg(right, ("key", "key", "<"), aggfunc=[("value", "sum")]) + + arr = right["value"].to_numpy() + booleans = np.zeros(n_right, dtype=np.bool_) + for l_idx, l_key in left["key"].items(): + start = int(np.searchsorted(right["key"].to_numpy(), l_key, side="right")) + if l_idx in actual.index: + got = actual.loc[l_idx, ("value", "sum")] + _assert_accurate(got, arr, booleans, start, n_right) + else: + assert start == n_right # no matches for this row