[PERF] Fuse join discovery and aggregation for join_agg

[samukweku]

Summary

Investigate and implement a direct fused execution path for join_agg that
performs join discovery and aggregation in the same kernel, without first
constructing the intermediate positions, matches, or materialized pair
arrays where the join shape permits it.

This is the stronger form of the Groupjoin idea: the join operator does not
emit every matching pair and a later operator does not rescan those pairs. It
updates aggregate state as matching candidates are discovered.

Related janitor-rs issue: #86 (perf: fuse multiple reverse aggregations over one join tape).

Motivation

The current join_agg path already avoids materializing a complete joined
DataFrame in many cases. However, for irregular multi-condition range joins it
can still construct intermediate structures such as:

• candidate range boundaries (starts and ends);
• a flattened positions tape;
• a flattened matches survivor mask;
• counts_array and total metadata;
• repeated output-label reconstruction for each aggregation.

These structures are useful for ordinary conditional_join, because ordinary
joins must eventually emit matching pairs. They are unnecessary when the
caller requests only join_agg and only needs one aggregate result per driving
row or per reverse output group.

The cost is especially important when:

• candidate ranges are broad;
• the final predicate is selective;
• the join is many-to-many;
• several aggregations are requested;
• the aggregate input column is narrow but the match tape uses int64 slots;
• the result is much smaller than the logical join cardinality.

Current execution model

For a reverse aggregation, a representative compact input can look like:

{
    "left_index": np.array([0, 1]),
    "right_index": np.array([100, 101, 102, 103]),
    "starts": np.array([0, 2]),
    "ends": np.array([2, 3]),
    "positions": np.array([0, 1, 1]),
    "counts_array": np.array([2, 1]),
    "total": 3,
}

This represents:

left row 0 -> right slots 0, 1
left row 1 -> right slot  1

The positional representation is substantially better than materialized
(left_index, right_index) pairs, but positions still contains one integer
per logical match. Its size is:

len(positions) == total == counts_array.sum()

The proposed path would update aggregate state while producing these logical
matches, or avoid producing them altogether when the join algorithm can test
the predicates directly.

Proposed execution strategies

1. Equality/groupjoin fast path

For equality joins, build a hash table keyed by the right join key. Each left
row probes the table and updates aggregate state for its matching right rows.
Do not construct pair indices unless the caller also requests ordinary join
output.

This is the closest direct application of the database Groupjoin operator.
It is a particularly attractive target because join_agg already permits
equality-only joins even though ordinary conditional_join does not.

Candidate flow:

build right-key lookup
        |
probe one left row
        |
visit matching right slots
        |
update sum/prod/min/max/size state
        |
emit one result row for the left group

2. Contiguous sorted-range fast path

When a sorted range predicate produces a contiguous right-side interval, use
the interval boundaries directly. For suitable aggregates, use prefix/suffix
state or a direct range kernel without constructing positions.

This path should be considered separately for:

• sum, including compensated floating-point sum where required;
• size;
• min and max using the existing range-extreme machinery;
• prod, subject to its identity, null, and overflow contracts.

3. Fused predicate-and-aggregate path

For multi-condition joins, scan each candidate only once and evaluate the
remaining predicates in that same loop. If a candidate survives, update the
aggregate state immediately.

Conceptually:

for each driving row:
    initialize aggregate state
    for each candidate right row:
        if all predicates pass:
            update aggregate state
    emit aggregate result

This avoids a matches survivor mask and a subsequent aggregation pass. It is
not always preferable: if the same candidate tape is reused by multiple
operations, retaining a compact representation may still be cheaper.

4. Adaptive fallback

Do not remove the existing compact paths. Select among:

• direct fused execution;
• contiguous range kernels;
• compact positions/matches execution;
• ordinary materialized join execution.

The decision should be based on join shape and requested operation, not on an
unvalidated assumption that fusion is always faster.

Candidate internal boundary

The public Python API should remain unchanged:

df.join_agg(
    right,
    ("left_key", "right_key", ">="),
    ("left_key_2", "right_key_2", "<="),
    aggfunc=[("value", "sum"), ("value", "min"), ("id", "size")],
)

Possible internal Rust boundaries include:

compute_join_agg_<shape>_<dtype>(join columns, aggregate columns, null masks)

or a lower-level builder/consumer split:

join candidate iterator -> aggregate accumulator

The latter may make it easier to share join traversal among multiple
aggregates while retaining independent kernels as a fallback.

The design should avoid exposing a Python object for every candidate and avoid
allocating Python lists in the hot loop.

Multi-aggregation design

janitor-rs #86 focuses on fusing several aggregations over an already-built
join tape. This issue should evaluate the additional benefit of fusing join
discovery with that aggregate traversal.

Possible levels of fusion:

1. independent aggregate kernels over the existing tape;
2. one shared tape traversal with multiple accumulator states;
3. one shared join traversal that evaluates predicates and updates multiple
   accumulator states directly.

The implementation should support a clear compatibility matrix. For example,
it may fuse multiple aggregates only when they share:

• the same driving side;
• the same join shape;
• the same aggregate input column and dtype;
• compatible null semantics;
• compatible output grouping.

Different columns or dtypes may remain on the existing path until a benchmark
justifies a wider fused state structure.

Correctness requirements

The implementation must preserve the existing join_agg contract for:

• duplicate left rows;
• duplicate right labels;
• many-to-many joins;
• equality, inequality, and mixed equality/range predicates;
• sorted and unsorted right inputs;
• reverse=False and reverse=True;
• nullable integer, boolean, floating, datetime, and extension-array inputs;
• null skipping for value aggregates;
• size counting semantics;
• empty and zero-width candidate ranges;
• all-null aggregate inputs;
• min/max empty-result sentinels;
• integer wrapping/overflow behavior;
• floating-point compensated summation behavior;
• product identity and overflow behavior;
• deterministic output ordering;
• output labels containing sparse original dataframe positions.

The fused path must not change ordinary conditional_join behavior. If the
caller asks for join indices, include_join_positions, keep="all", or any
other output that requires pairs, the existing materialization-compatible path
must remain available.

Memory accounting

Every benchmark must report at least:

• candidate range metadata;
• positions or matches tape bytes, when present;
• hash-table or label-to-slot state;
• aggregate accumulator state;
• output arrays;
• peak resident/live allocation where measurable.

Fusion is not automatically a memory win. A fused kernel may keep several
aggregate states alive simultaneously. It should be compared against both:

one aggregate over a compact tape

and:

several independent aggregates over the same compact tape

An adaptive policy may be appropriate when aggregate state is large.

Benchmark matrix

Compare the current implementation with each candidate fused implementation
for:

• one, two, three, and five aggregations;
• sum, prod, min, max, and size;
• one input column and multiple input columns;
• integer and floating-point dtypes;
• null-free, partially-null, and all-null inputs;
• equality joins;
• single range joins;
• dual range joins;
• mixed equality/range joins;
• selective and highly overlapping predicates;
• unique and duplicate join labels;
• sorted and unsorted right inputs;
• tiny, large, very-large, and super-large inputs.

For each case record:

• wall-clock runtime;
• peak memory;
• logical match count;
• output size;
• whether a positions or matches allocation occurred;
• output equivalence against the current implementation.

Include a break-even analysis. A fused kernel that helps only when five
aggregations are requested should not replace the one-aggregation path.

Suggested implementation phases

Phase 1: baseline and instrumentation

• Add a reusable benchmark harness for join_agg.
• Measure current pair, compact-tape, and aggregate costs separately.
• Record allocation counts and peak memory where available.
• Establish correctness fixtures for every supported aggregation and dtype.

Phase 2: equality direct-fusion prototype

• Implement a direct equality join_agg path.
• Avoid pair-index construction.
• Compare it with the current equality-plus-aggregation path.
• Verify duplicate-key and duplicate-label behavior.

Phase 3: single contiguous range prototype

• Add a direct path for a sorted range whose matches remain contiguous.
• Reuse existing prefix/range-extreme kernels where possible.
• Establish when this path is safe for each aggregate.

Phase 4: multi-condition fused prototype

• Evaluate predicates and aggregate in one candidate traversal.
• Start with size and integer sum.
• Add null-aware and floating-point variants only after contracts are tested.

Phase 5: multi-aggregation fusion

• Coordinate with janitor-rs filtering with nice strings in between #86.
• Share traversal and candidate lookup among compatible aggregate states.
• Add an adaptive fallback based on aggregate-state size and requested count.

Non-goals

• Removing positions from ordinary joins that must emit all pairs.
• Replacing all range representations with bit masks.
• Assuming equality/groupjoin techniques solve arbitrary non-equi joins.
• Changing the public join_agg API without benchmark evidence.
• Treating a reduction in intermediate arrays as sufficient evidence of a
  runtime improvement.

Acceptance criteria

• A documented baseline exists for all major join_agg shapes.
• At least one direct fused path demonstrates a measured benefit over the
  current implementation.
• Correctness is proven against the existing implementation across duplicate,
  null, empty, overflow, dtype, and ordering cases.
• Peak memory is reported, including aggregate state and any retained tape.
• The dispatch conditions are explicit and conservative.
• Unsupported or unprofitable cases continue to use the existing kernels.
• The pyjanitor and janitor-rs issue/PR relationship is documented, including
  the distinction between direct fusion and janitor-rs filtering with nice strings in between #86’s shared traversal
  over an existing tape.