perf(conditional_join): compensated prefix sum for float join_agg ranges

[samukweku]

Summary

• Replaces the forward float32/float64 sum range kernels behind
  join_agg (suffix _sum_starts, prefix _sum_ends, and
  arbitrary-interval _sum_starts_ends) with a Neumaier-compensated
  prefix sum built once per column, so each range query becomes O(1)
  instead of Rust re-scanning every overlapping range from scratch.
• Restricted to float dtypes; the integer counterpart is tracked
  separately in [PERF] Use prefix sums for conditional_join range aggregations (integer dtypes) #1648 (a plain np.cumsum is enough there — no
  compensation needed).
• Gated by a work-estimate heuristic (_prefix_sum_is_worthwhile): the
  one-time O(n) Python build only pays off once the total width summed
  across all queried ranges clears roughly 150x the array size —
  otherwise Rust's native per-range scan is cheaper overall, even though
  it re-scans overlaps. Below that threshold this defers straight to
  Rust.
• Two additional correctness guards fall back to Rust whenever
  single-level compensation can't be trusted for the whole array: a real
  +/-inf or a genuinely overflowing partial sum, and extreme
  within-array dynamic range (see ELI5).
• Fixes a pre-existing duplicate _sum_starts_ends definition found
  while touching this function (the second copy silently shadowed the
  first; both were identical, so no behavior change).

ELI5

Summing a bunch of overlapping ranges the old way means re-adding the
same numbers over and over — once per range. A prefix sum writes down
the running total once, so any range's sum is just two lookups and a
subtraction. That's the win in #1648 for integers.

Floats are trickier because the existing Rust kernels already use
compensated (Kahan) summation to stay accurate — they track a
"leftover" alongside the running total to avoid losing precision when a
tiny number gets added to a much bigger one. So this PR builds a
compensated prefix instead of a plain one, using the same idea
(Neumaier's variant) so range subtraction stays close to summing that
slice directly.

Two things can break that, both found via randomized testing rather than
by inspection:

1. Overflow. Two huge-but-finite numbers can sum to +/-inf (that's
   normal float behavior). But once the prefix's running total hits
   +/-inf, subtracting two +/-inf prefix entries for a later range
   that never touched the huge numbers gives inf - inf = NaN — wrong.
   Rust's kernels don't have this problem because they recompute each
   range from scratch.
2. Compensation of the compensation. The "leftover" tracker is itself
   just an ordinary float. If it's already absorbed a moderate correction
   (say, size ~1e12), it can no longer represent a much later, much
   tinier correction (say ~1e-6) — that increment is too small to
   register against the leftover's own precision, the same way 1e12 + 1
   doesn't change a value that's already around 1e20. In testing, a big
   excursion followed by another huge, opposite-signed excursion followed
   by a small value made a single-element window silently return 0.0
   instead of the true (tiny) value — 100% relative error, with nothing
   ever overflowing to catch.

Both are handled by falling back to Rust for the whole array when either
condition is detected, rather than risking either failure mode.

Kahan and Neumaier can still round differently on cancellation-heavy data even
when both are accurate. The numerical contract therefore uses math.fsum as
the high-precision oracle and bounds forward error by
8 * float64_epsilon * sum(abs(inputs)). A result-relative-only tolerance is
not meaningful here because cancellation can make the result arbitrarily close
to zero. The tests apply this same scale-aware contract to both the new prefix
path and the current Rust baseline.

Benchmark

Suffix sum (_sum_starts, single < join), median of 3 runs,
float64, no nulls:

Case	Rows (n)	Queries (m)	Rust	This PR	Speed-up
Many overlapping ranges	1,000	1,000	~1 ms	~0.4 ms	~2.5x
Many overlapping ranges	20,000	20,000	~409 ms	~5 ms	~82x
Single query, large array	200,000	1	~0.4 ms	~0.4 ms	1x (deferred to Rust, no regression)
Single query, large array	2,000,000	1	~4 ms	~4 ms	1x (deferred to Rust, no regression)

The "many overlapping ranges" case is what #1648/#1671 target — Rust's
cost there grows with the sum of all range widths, not the array size,
so it degrades badly as overlap increases (note it's already ~1000x
slower than the O(n+m) approach going from n=1,000 to n=20,000 with full
overlap). The "single query" rows are the case an earlier review pass on
this PR caught as a regression before the work-estimate heuristic was
added: a small number of queries against a large array used to pay the
full O(n) prefix-build cost for almost no reuse (~100x slower than
Rust in that case); the heuristic now correctly defers to Rust instead.
These are development-machine measurements, not a committed benchmark
suite (per this repo's convention, the one-off script used to gather
them isn't included in the diff).

Test plan

• New tests/functions/test_conditional_join_prefix_sum_float_agg.py
  (28 tests): scale-aware accuracy-contract tests against math.fsum for
  both the prefix path and the existing Rust baseline across
  float32/float64 and all three kernels; bitwise parity for
  well-scaled data; null/all-null/empty-range handling; regressions for
  overflow, extreme dynamic range, and cancellation-heavy duplicate
  queries that cross the dispatch threshold; heuristic coverage; and
  end-to-end join_agg coverage of both dispatch paths.
• Full existing test_conditional_join.py suite (241 tests,
  including all hypothesis-based *_agg property tests) passes
  unchanged.
• 20,000-trial randomized adversarial validation against math.fsum
  (base scales from 1e-80 to 1e80, within-array spreads up to 1e15,
  random null placement and ranges): every accepted prefix stayed
  within the 8 * eps * sum(abs(inputs)) contract; worst observed
  error was ~0.995 epsilon-scaled units.
• All of the above pass under -W error::RuntimeWarning (no leaked
  overflow/invalid-value warnings from the expected-and-handled
  overflow/dynamic-range paths).
• pixi run lint clean on touched files.
• Adversarially reviewed by a fresh-context pass before this PR was
  opened, which caught the single-query-against-large-array
  regression addressed by the work-estimate heuristic.

Related: #1648 (integer counterpart), #1653/#1654 (min/max/prod range
algorithms), pyjanitor-devs/janitor-rs#23 (reverse aggregation kernels,
out of scope here).

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://pyjanitor-devs.github.io/pyjanitor/pr-preview/pr-1675/
Built to branch gh-pages at 2026-08-22 07:52 UTC. Preview will be ready when the GitHub Pages deployment is complete.