Skip to content

[PERF] Use prefix sums for conditional_join range aggregations (float dtypes) #1671

Description

@samukweku

Brief description

Follow-up to #1648, which restricts the prefix-sum rewrite of the forward
range-sum kernels used by join_agg to integer dtypes.

This issue tracks extending the same optimization to float32/float64,
which is deliberately out of scope for #1648.

Relevant Rust implementations:

  • src/aggs/sum/sum_starts.rs
  • src/aggs/sum/sum_ends.rs
  • src/aggs/sum/sum_starts_ends.rs

Relevant Python dispatch:

  • janitor/functions/_conditional_join/_agg_functions.py
  • janitor/functions/_conditional_join/_get_join_aggs.py

Why this is separate from #1648

Some current Rust kernels use compensated summation (e.g. Kahan/Neumaier) for
floats. A naive prefix-sum-and-subtract rewrite can change rounding or
cancellation behavior versus the current output, so this isn't a drop-in
replacement the way it is for integers.

Proposed scope

  1. Define and test an explicit numerical-accuracy contract for float prefix
    sums (e.g. acceptable ULP/relative-error bounds vs. the current
    compensated-summation kernels), or decide compensated summation must be
    preserved and design a prefix-sum variant that keeps it (e.g. a
    compensated prefix sum with a compensated subtraction step).
  2. Implement the NumPy (or Rust) prefix-sum helper for float32/float64
    covering suffix, prefix, and arbitrary [start:end) ranges, mirroring the
    integer implementation landed in [PERF] Use prefix sums for conditional_join range aggregations (integer dtypes) #1648.
  3. Add focused unit tests for rounding/cancellation parity, null handling,
    empty ranges, and extension-array reconstruction.
  4. Add benchmarks comparable to the ones gathered for [PERF] Use prefix sums for conditional_join range aggregations (integer dtypes) #1648.

Acceptance criteria

  • Float results match the current implementation within the agreed
    numerical-accuracy contract, across null placements, empty inputs, and
    range shapes.
  • The full test suite passes.
  • Benchmarks demonstrate an end-to-end improvement, not only a kernel-level
    improvement.
  • The changelog entry references this issue number, the implementing PR
    number, and @samukweku.

Related aggregation work:

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions