[PERF] NumPy prefix sums for integer conditional_join range-sum aggregations

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

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,