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

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

janitor/functions/_conditional_join/_agg_functions.py

@@ -1,15 +1,184 @@
 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,
     starts: np.ndarray,
     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,