[PERF] Selective predicate selection for equi+non-equi joins

[samukweku]

Summary

• Fixes [PERF] Selective predicate selection doesn't reach equi+non-equi combined joins #1664: _get_indices_equi.py (joins combining an == predicate with non-equi predicates) consumes the same mapping["le_or_ge"]/["le_lt"]/["ge_gt"] as the pure non-equi dispatch tree, but [PERF] Selectivity-aware anchor predicate selection for conditional_join #1658's _select_anchor and [PERF] Selective bound selection for range joins #1663's _maybe_select_better_range_bounds never reached it - same class of pathology (Issue [PERF] Optimize conditional_join anchor predicate selection #1641), different dispatch path.
• Adds _maybe_select_better_equi_predicates in _get_indices_equi.py, mirroring both prior fixes: the single-anchor case for _equi_not_range_join.py (reuses _le_ge_1_or_more._sample_candidate_cost directly, same helper [PERF] Selective bound selection for range joins #1663 reused), and the two-bound case for _equi_range_join.py.
• Scoped to keep in ('first', 'last'), matching [PERF] Selectivity-aware anchor predicate selection for conditional_join #1658/[PERF] Selective bound selection for range joins #1663's own scoping.
• _equi_uniq_join.py needs no fix, and is correctly left untouched - it flattens every non-equi predicate into one unordered post-filter set (dict.fromkeys(rest)) regardless of which mapping key it came from, since the == predicate already reduces to at most one candidate per left row via a hash indexer - there's no anchor/window narrowing to optimize there in the first place. The original issue text assumed all three consumers needed the fix; verified this by reading the file directly rather than assuming.
• Applying the selection upstream of the whole dispatch (not just the two paths that need it) is deliberate and safe: reordering mapping before _equi_uniq_join's try-block is harmless there too (order-insensitive), which is simpler than special-casing which downstream function runs first.

ELI5

conditional_join has a fast path for joins that combine an equality condition (like matching on an id or category) with inequality conditions (like a date range). That fast path picks one of the inequality conditions to do the real narrowing and treats the rest as simple checks - and it had the same "always use whichever one you typed first" bug #1641/#1658/#1659 already fixed elsewhere, just in this different code path that nobody had gotten to yet. This PR applies the same fix here too.

Benchmark

Equi join on a ~20-group column, plus 2 non-equi predicates (one much less selective than the other), before this fix:

n	broad-first	selective-first	ratio
1,000	2.14ms	2.00ms	1.07x
3,000	2.56ms	2.13ms	1.20x
10,000	6.80ms	2.81ms	2.42x
30,000	40.38ms	4.51ms	8.95x

After this fix, flat at ~1.00-1.06x across the same range.

Test plan

• pixi run pytest tests/functions/test_conditional_join.py -v -n auto - 294 passed, 1 skipped (full existing suite, unchanged)
• New tests: order-invariance for both the equi+non-range (_equi_not_range_join.py) and equi+range (_equi_range_join.py) dispatch paths, across keep='first'/'last'; a mock-based guard proving the selection is never invoked for keep='all'; a no-op check for the single-candidate case (the common case).
• pixi run pytest --doctest-modules janitor/functions/conditional_join.py
• pixi run lint

Draft PR: depends on the unmerged #1658/#1663 chain (based on issue-1659-range-join-anchor, which is based on issue-1641-anchor-selection) - not mergeable until those land first.

🤖 Generated with Claude Code