[PERF] Reduce or eliminate intermediate matches arrays in conditional_join

[samukweku]

Summary

conditional_join uses a flat matches array as a survivor mask while applying multiple predicates over candidate ranges. This avoids materializing all (left, right) pairs, but the mask can still be as large as the sum of all candidate range widths and is rescanned/reallocated across conditions.

Relevant implementation:

• janitor/functions/_conditional_join/_helpers.py: _get_positive_matches and build_indices_matches
• Rust comparison kernels receiving matches and counts_array

Problem

For a left row with candidate range [start, end), the current representation reserves one mask entry for every candidate, including candidates eliminated by earlier predicates. With broad ranges and several conditions, this can create substantial memory traffic even when the final result is selective.

The same intermediate state is also unnecessary for keep="first" and keep="last", where only one surviving right position per left row is required.

Possible directions

Investigate and benchmark the following alternatives:

1. Fused predicate evaluation in Rust: evaluate all conditions in one candidate scan without an intermediate matches tape.
2. Direct first/last selection: scan candidates in the requested direction and stop at the first complete match.
3. Compact survivor representation: use CSR-style row offsets plus surviving right positions after selective predicates.
4. Bit-packed masks: replace the int8 mask with a uint64 bitset if the tape remains necessary.
5. Interval/segment survivors for predicates whose matches remain contiguous under sorted right-side data.

Suggested scope

Start with a benchmark and a fused/direct-selection implementation for keep="first" and keep="last"; retain the current path as a fallback for keep="all", irregular predicates, and aggregation paths until equivalent behavior is established.

Acceptance criteria

• Preserve output ordering and null semantics, including nullable dtypes and !=.
• Preserve join_agg, include_join_positions, and return_building_blocks behavior.
• Compare peak memory and runtime against the current implementation for broad ranges, selective predicates, multiple conditions, and high-cardinality inputs.
• Add regression tests covering empty ranges, duplicate values, nulls, sorted/unsorted right inputs, and all keep modes.
• Document when the existing matches representation remains the preferred path.

[samukweku]

Implementation sequence

The optimization sequence is now recorded here:

1. Predicate projection during matching — [PERF] Project conditional_join predicate columns during matching #1647 (merged).
2. Selectivity-aware anchor predicate selection — [PERF] Selectivity-aware anchor predicate selection for conditional_join #1658 (merged).
3. Reduce/eliminate intermediate matches arrays — this issue:
   • baseline benchmark (started in the implementation branch);
   • implement the Rust primitive for direct first/last selection;
   • integrate it into pyjanitor for those keep modes;
   • add correctness, runtime, and memory validation;
   • retain the current tape for keep="all" and join_agg until a separate survivor representation is justified.
4. Follow-up: compact survivor representation for keep="all"/aggregation paths.
5. Consider bit-packed masks or interval/segment survivors only if measurements justify them.

The required janitor-rs work is tracked in janitor-rs #101. This is a separate dependency/API task; pyjanitor integration should follow its implementation and release.

[samukweku]

Why direct first/last selection needs a Rust kernel

For each left row, the first predicate produces a candidate range, but later predicates can eliminate the early candidates. For example, candidates [10, 11, 12, 13] may become [10, 11, 12], then [11, 12], then [12]. The current implementation carries a flat matches tape through each predicate:

[1, 1, 1, 1] -> [0, 1, 1, 0] -> [0, 0, 1, 0]

That is unnecessary for keep="first"/keep="last", where only the first/last complete match is needed. The direct path should scan each candidate range in the requested direction, evaluate the full predicate chain for each candidate, stop at the first complete match, and emit one right position (or no match).

This belongs in janitor-rs because the candidate scan is the hot loop, predicates may use different dtypes, and selecting a winner after each individual predicate would be incorrect. The existing matches path remains for keep="all" and join_agg, where every surviving pair is required.

Coordination

Please do not close this issue automatically from any PR. I will close #1687 manually after all related PRs, including the janitor-rs dependency work, are merged.

[samukweku]

Ordering correction

Clarification to the direct-selection description: keep="first"/"last" means the smallest/largest original right-row position, not the first/last element in an internally sorted candidate range. A direct kernel therefore generally scans the full candidate range and compares the original right indices; it may stop early only when that range is proven monotonic in original-row order.

[samukweku]

Implementation decision and current status

The first implementation of this optimization follows a deliberately narrow, fallback-safe path:

• janitor-rs provides a fused direct-selection kernel that evaluates the full predicate chain in one candidate scan, without a matches tape.
• It is used initially only for keep="first" and keep="last".
• The pyjanitor integration is gated to homogeneous, supported, non-null dtypes. Mixed dtypes, nullable inputs, unavailable janitor-rs exports, and unsupported cases continue to use the existing matches-tape implementation.
• keep="all" and join_agg remain on the existing tape path because they require all surviving pairs; a compact survivor representation is a separate follow-up.
• The Rust kernel selects the minimum or maximum original right-row index, rather than relying on internal candidate ordering, preserving the existing first/last semantics for unsorted right inputs.

Related implementation work:

• janitor-rs direct kernel: #108
• pyjanitor integration: #1691

The kernel has correctness coverage for invalid ranges, duplicates, ordering, empty/no-survivor cases, and benchmark coverage across small, large, and very-large inputs. The remaining work is review, merge/release coordination, and then separate investigation of a compact survivor representation for keep="all"/aggregation.

This issue remains open until all related implementation and integration PRs are merged; it should not be auto-closed by an individual PR.

[samukweku]

Findings from independently re-reading the current implementation

Traced this against main as of the "Selectivity-aware anchor predicate selection" (#1658) and "direct selection" (3c7d099) commits, rather than starting from the directions listed above.

Two tape representations already exist, not one

• matches — an i8 boolean AND-mask, one byte per candidate, walked in lockstep with an implicit contiguous right[start:end] slice per row (_helpers.py::_get_positive_matches, Rust core in compare/comp.rs::compare_start_end_core). Used by _range_join_default.py, _le_ge_1_or_more.py, _equi_range_join.py — the paths where the right side is sorted and a row's candidates really are one contiguous slice.
• positions — an i64 tape holding the actual right-row index or a -1 "eliminated" sentinel (comp_posns.rs), used wherever candidates aren't contiguous (_dual_non_equi.py-derived paths: _range_join_regions.py, _not_range_join_default.py, _not_range_join_regions.py).

Both are consumed the same way by _get_positive_matches_conditions / _get_positive_matches_conditions_posns in _helpers.py: each extra predicate in rest gets a freshly allocated, full-width tape, walked start-to-end again. counts_array[i] shrinks as rows lose candidates, but the physical width handed to the next predicate never does — a dead slot still costs a branch-check (if matches[n]==0: skip) in every subsequent condition. This is the actual mechanism behind the "rescanned/reallocated across conditions" cost, and it's identical in both representations, not just the i8 one.

What's already landed

• #1658 — selectivity-aware anchor predicate ordering, so the cheapest predicate runs first and narrows the tape before the expensive ones run.
• 3c7d099 — a genuine tape-free path (select_start_end_direct in Rust), covering directions great idea! #1/Examples needed #2 above. It only fires for keep in {"first", "last"}, non-null predicate columns, and a single shared dtype across all predicate columns (_helpers.py::_get_direct_indices_conditions). keep="all" and mixed-dtype/nullable cases still fall through to the tape.

Remaining gap and the worst case

The gap is direction #3 for keep="all", and specifically: _le_ge_1_or_more.py (a single one-sided anchor, e.g. only >), where ends is implicitly right_len for every row. Tape width there can be O(n × right_len) regardless of how selective the extra predicates are, and no existing optimization (anchor ordering, direct selection) touches it, since anchor ordering only reorders which predicate sets starts/ends, and direct selection is keep="first"/"last"-only. This is the scenario I'd anchor a benchmark on.

Proposed direction: compact into the existing positions tape, adaptively

Rather than a third representation, reuse the positions tape that already exists for the regions algorithm, but apply it generally as an intermediate compaction step between predicates (not just as the regions algorithm's native output):

• After a predicate whose survivor fraction drops below roughly 1/8 of the current width, materialize a compact positions array (surviving right-indices only) plus fresh CSR starts/ends via cumsum(counts_array), and route the next predicate through compare_posns/compare_posns_ne instead of compare_start_end/compare_start.
• The 1/8 threshold isn't arbitrary: positions entries are i64 (8 bytes) vs matches' i8 (1 byte), so compacting only pays for itself memory-wise once more than 7/8 of candidates have been eliminated. Below that, compacting increases memory — worth calling out since a naive "always compact" strategy would regress the common low-selectivity case.
• Complementary to direction [DOC] Add "missing data" viz to example notebooks #4: bit-pack the un-compacted matches tape (u64 words) for the case where compaction doesn't pay off. That's a pure storage-format change, independent of the compaction logic above.

The real cost: aggregation kernels don't have a positions-based counterpart

join_agg's Rust kernels (aggs/sum, min, max, prod, forward and reverse — on the order of 50 files) only know the (starts, ends) + matches contract; there's no positions-based aggregation kernel today. Compacting mid-chain would silently break aggregation unless every agg family gets a _posns sibling. I'd stage this as two PRs:

1. Compaction for the filter/index-building path only (immediate win for keep="all"; no agg-kernel changes needed, since conditional_join.py's dispatch already routes aggfunc cases through a separate path).
2. Port agg kernels to consume positions in a follow-up, once (1) proves out on the benchmark above.

This keeps the "preserve join_agg... until equivalent behavior is established" acceptance criterion satisfiable without blocking the keep="all" filtering win on the much larger agg-kernel port.

[samukweku]

I went through the actual kernel source on both sides (comparison kernels in janitor-rs/src/compare/, aggregation kernels in janitor-rs/src/aggs/*_rev/) independently of this issue, prompted by a separate multi-PR bug review of the _matches aggregation family. Sharing what I found since it's directly relevant to the direction here — in particular, it changes the framing of this from a pure perf issue to a perf-plus-correctness one.

The matches tape's real cost isn't just memory — it's a fragile cross-array contract

The producer (comp_first_starts.rs) walks each row's [start, end) range and writes one i8 flag per candidate into a flat array sized to sum(end_i - start_i), computed once at build time. The consumer (every compute_*_matches kernel) re-walks the same row ranges independently and steps through matches via its own separately-tracked running counter. Nothing ties these two walks together structurally — both sides just have to independently re-derive the identical width formula and agree on it exactly.

That's not a hypothetical risk: it's the root mechanism behind a whole class of bugs I found across ~15 PRs in the _starts_matches/_ends_matches/_starts_ends_matches kernel families — off-by-one row boundaries, disputes over whether an all-zero-width batch should error, per-row zero-width rows wrongly rejected. All of them trace back to this same "two independently-computed quantities must agree exactly" shape. So eliminating the tape isn't only a memory/perf win — it removes a recurring, easy-to-reintroduce defect surface.

It also doesn't scale the way you'd want under multiple chained conditions: comp_starts.rs (used for the 2nd+ non-equi predicate) rebuilds a full-size tape and re-walks the entire original range again per condition, ANDing against the previous pass. N chained conditions costs O(N × total candidate width), with no narrowing between passes — worth calling out explicitly in the "compare peak memory and runtime" acceptance criterion, since it's not just the base construction that's expensive.

A working alternative already lives in the same codebase: the positions family

There's a parallel kernel family (comp_posns.rs, and aggregation-side consumers like max_positions_core) that doesn't use a boolean survivor tape at all. Each row's (start, end) bounds a sub-range directly inside the positions array itself, validated per-row via checked_range(start, end, positions.len()) — which returns None (a silent per-row skip) rather than requiring any externally-derived width to match. The loop then indexes positions[nn] straight through, storing -1 for excluded candidates. There's no separate running counter and no cross-array length contract to violate.

I flagged this because in three independent adversarial review passes over the aggregation PRs, the positions family never once produced this bug class — its only real defect was unrelated length-kwarg plumbing. That's decent evidence this shape is structurally sound in practice, not just in theory.

Caveat, since it matters for direction 3 in the issue: positions as it exists today does not compact between chained conditions — it stays full-size and masks with -1 rather than shrinking. So on its own it doesn't give the O(N × width) → O(width + survivors) improvement that a true CSR redesign implies. Getting that win requires actually compacting survivors after each predicate (e.g. via np.compress/boolean-mask filtering, or a Rust-side compaction pass), which isn't in either family as it stands today.

On the specific "possible directions" listed

• (3) CSR-style compact survivors — this is where I'd put the emphasis, but built by extending the existing positions pattern (self-describing per-row bounds, safe indexing) with real compaction, rather than inventing a new representation. It gets you both the perf win and inherits the robustness that pattern has already demonstrated.
• (4) Bit-packed uint64 masks — this shrinks the memory quantum but keeps the exact same fragile "two independently-computed widths must agree" contract that's driving the correctness bugs. I'd treat it as a fallback if (3) doesn't land, not a first choice.
• Memory honesty: for != (the dominant real-world predicate for this kernel family in current usage), selectivity is typically low — it excludes only exact-value matches and keeps almost everything — so a dense i8 tape is already close to memory-optimal for that specific case, and a naive dense positions swap (i64, 8 bytes/candidate) would actually cost more memory with no compaction. The memory win from moving to positions only materializes once genuine compaction is added, and mainly benefits the more selective predicates, not !=. Worth reflecting in the benchmark matrix — !=-heavy workloads vs. selective-range workloads will show opposite memory deltas.

Happy to go deeper on any of the kernel files above if useful for scoping the CSR work.

[samukweku]

Focused recommendation from source review

The current matches object is a flat survivor mask over the sum of all candidate range widths. It is not merely an output array: every additional predicate receives the full-width tape, skips dead slots, and allocates a fresh full-width result tape. Thus chained predicates cost O(number_of_conditions × candidate_width), with peak memory temporarily including both the input and output tapes.

The preferred design for keep="all" is genuine survivor compaction:

• retain starts/ends as row boundaries;
• after a selective predicate, compact surviving right positions into a new tape;
• rebuild per-row offsets/counts;
• evaluate subsequent predicates only over survivors.

The existing positions kernel family is a good structural basis because it makes row bounds explicit and avoids the separate implicit tape-width contract. However, simply switching from int8 matches to an int64 positions tape without compaction is not an improvement: it is approximately 8x larger and still rescans dead slots.

A practical adaptive policy is:

• keep the dense int8 mask while survivors are dense;
• compact to int64 positions once more than 7/8 of candidates have been eliminated, since int64 offsets only become memory-cheaper below roughly 1/8 density;
• add bit-packed uint64 masks as an independent optimization for cases where compaction is not worthwhile.

For join_agg, compaction needs corresponding positions-based aggregation kernels or a carefully isolated conversion back to the range/mask contract. Therefore I would stage the work:

1. benchmark and implement adaptive compaction for ordinary index-building/keep="all" paths;
2. separately extend aggregation kernels;
3. use fused direct-selection kernels for first/last, where all surviving pairs are not required;
4. treat bit-packing as a fallback/orthogonal optimization, not the primary architectural change.

This preserves the existing contiguous-range fast paths while eliminating unnecessary work when predicates are selective.