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

CHANGELOG.md

@@ -1,6 +1,7 @@
 # Changelog
 
 ## [Unreleased]
+-   [ENH] `conditional_join` calls that combine an `==` predicate with 2+ non-equi predicates (same-direction or a range join) now get the same selective-predicate selection as the pure non-equi path, when `keep` is `'first'` or `'last'` - fixing the same class of pathological slowdown Issue #1641/#1659 fixed elsewhere, for this dispatch path too. - Issue #1664 @samukweku
 -   [ENH] Range joins (a `<`/`<=` predicate combined with a `>`/`>=` predicate) in `conditional_join` now pick the most selective candidate independently for each bound, when there are 2+ eligible predicates of either type and `keep` is `'first'` or `'last'` - fixing the same class of pathological slowdown from unfavorable predicate ordering as Issue #1641, for both `join_algorithm='default'` and `'regions'`. - Issue #1659 @samukweku
 -   [ENH] `conditional_join` now picks the most selective `<`/`<=`/`>`/`>=` predicate as its binary-search anchor (instead of the first one supplied) when `keep` is `'first'` or `'last'`, fixing pathological slowdowns from unfavorable predicate ordering; the choice is estimated from a fixed-size sample so the selection cost no longer scales with input size; `keep='all'` output is unaffected. - Issue #1641 @samukweku
 -   [BUG] Fix `conditional_join` crash for `keep='last'` joins with a single `<`/`<=` window and multiple non-equi conditions - a missing `counts` argument to the Rust `index_starts_only_keep_last` call. - Issue #1641 @samukweku

janitor/functions/_conditional_join/_get_indices_equi.py

@@ -10,9 +10,86 @@
     _equi_range_join,
     _equi_uniq_join,
     _helpers,
+    _le_ge_1_or_more,
 )
 
 
+def _select_best_candidate(candidates: list, df: pd.DataFrame, right: pd.DataFrame):
+    """
+    Pick the cheapest-looking candidate from a pool via the same sampling
+    approach as `_le_ge_1_or_more._select_anchor` /
+    `_get_indices_non_equi._select_range_bound`. Ties are broken toward
+    the earliest candidate, so an already-optimal ordering is left
+    untouched.
+
+    Only called when `candidates` has 2+ entries - see the call sites.
+    """
+    best_cost = None
+    best = None
+    for candidate in candidates:
+        cost = _le_ge_1_or_more._sample_candidate_cost(candidate, df, right)
+        if best_cost is None or cost < best_cost:
+            best_cost = cost
+            best = candidate
+    return best
+
+
+def _maybe_select_better_equi_predicates(
+    mapping: dict, df: pd.DataFrame, right: pd.DataFrame
+):
+    """
+    Selectivity-aware refinement of `mapping["le_or_ge"]`/`["le_lt"]`/
+    `["ge_gt"]` for the equi + non-equi dispatch tree (`_get_indices_equi
+    .py`), mirroring #1658's `_select_anchor` (for `_equi_not_range_join
+    .py`'s single-anchor case) and #1663's `_maybe_select_better_range_
+    bounds` (for `_equi_range_join.py`'s two-bound case) - see issue
+    #1664. `_equi_uniq_join.py` needs no such refinement: it flattens
+    every non-equi predicate into one unordered post-filter set
+    regardless of which mapping key it came from, so which candidate
+    holds which role never affects it - reordering here is harmless for
+    that path, not just unneeded.
+
+    Mutates and returns `mapping`. Matched row content is unaffected
+    either way (same invariance as #1641/#1658/#1663); only performance
+    - and, for `keep='all'`, row order - can change, so this is only
+    worth calling when there's an actual choice and `keep in ('first',
+    'last')`, mirroring #1658/#1663's own scoping.
+    """
+    if mapping["is_range_join"]:
+        # _equi_range_join.py consumes le_lt/ge_gt directly and treats
+        # any extra le_or_ge candidates as unordered post-filters
+        # already (like _equi_uniq_join.py) - only le_lt/ge_gt matter.
+        for bound_key, candidates_key in (
+            ("le_lt", "le_lt_candidates"),
+            ("ge_gt", "ge_gt_candidates"),
+        ):
+            candidates = mapping[candidates_key]
+            if len(candidates) < 2:
+                continue
+            current = mapping[bound_key]
+            best = _select_best_candidate(candidates, df, right)
+            if best == current:
+                continue
+            # `best` is a le_or_ge member being promoted to `bound_key`;
+            # `current` demotes into its slot. Order within le_or_ge
+            # doesn't matter - it's always applied as an unordered set of
+            # independent post-filters - so an in-place swap is enough.
+            mapping["le_or_ge"][mapping["le_or_ge"].index(best)] = current
+            mapping[bound_key] = best
+        return mapping
+    # _equi_not_range_join.py picks the first le_or_ge entry as anchor,
+    # so unlike the range-join case above, position 0 matters here - swap
+    # the winner into it (mutates mapping["le_or_ge"] in place).
+    candidates = mapping["le_or_ge"]
+    if len(candidates) < 2:
+        return mapping
+    best = _select_best_candidate(candidates, df, right)
+    if best != candidates[0]:
+        best_pos = candidates.index(best)
+        candidates[0], candidates[best_pos] = candidates[best_pos], candidates[0]
+    return mapping
+
+
 def _get_indices(
     df: pd.DataFrame,
     right: pd.DataFrame,
@@ -49,6 +126,10 @@ def _get_indices(
             "left_index": empty_array,
             "right_index": empty_array,
         }
+    if keep in ("first", "last"):
+        mapping = _maybe_select_better_equi_predicates(
+            mapping=mapping, df=df, right=right
+        )
     try:
         # this section assumes one-to-one or many-to-one
         # no need to capture indices for aggfunc here

tests/functions/test_conditional_join.py

@@ -4956,6 +4956,127 @@ def test_range_join_bound_selection_noop_with_single_candidate():
     assert result["le_or_ge"] == []
 
 
+# --- issue #1664: selective predicate selection for equi + non-equi joins ---
+#
+# `_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 #1658/#1663's fixes never
+# reached it. `_maybe_select_better_equi_predicates` mirrors both: the
+# single-anchor case (for `_equi_not_range_join.py`) and the two-bound
+# case (for `_equi_range_join.py`). `_equi_uniq_join.py` needs no fix -
+# it flattens every non-equi predicate into one unordered post-filter
+# set regardless of which mapping key it came from.
+
+
+def _skewed_equi_broad_selective_frames(seed, n=3000, n_groups=20):
+    """Like `_skewed_broad_selective_frames`, plus an equi column - for
+    the `_equi_not_range_join.py` dispatch path."""
+    rng = np.random.default_rng(seed)
+    df = pd.DataFrame(
+        {
+            "grp": rng.integers(0, n_groups, size=n),
+            "l_broad": rng.integers(0, 10, size=n),
+            "l_selective": rng.integers(n - 10, n, size=n),
+        }
+    )
+    right = pd.DataFrame(
+        {
+            "grp": rng.integers(0, n_groups, size=n),
+            "r_broad": rng.integers(0, n, size=n),
+            "r_selective": rng.integers(0, n, size=n),
+        }
+    )
+    return df, right
+
+
+def _skewed_equi_range_frames(seed, n=3000, n_groups=20):
+    """Like `_skewed_range_frames`, plus an equi column - for the
+    `_equi_range_join.py` dispatch path."""
+    rng = np.random.default_rng(seed)
+    df = pd.DataFrame(
+        {"grp": rng.integers(0, n_groups, size=n), "a": rng.integers(0, n, size=n)}
+    )
+    lo = rng.integers(0, n, size=n)
+    right = pd.DataFrame(
+        {
+            "grp": rng.integers(0, n_groups, size=n),
+            "lo": lo,
+            "hi_narrow": lo + rng.integers(1, 5, size=n),
+            "hi_broad": lo + rng.integers(1, n, size=n),
+        }
+    )
+    return df, right
+
+
+@pytest.mark.parametrize("keep", ["first", "last"])
+def test_equi_non_range_predicate_order_invariant(keep):
+    """Output for an equi join combined with 2+ non-equi predicates must
+    not depend on which order the non-equi predicates are supplied in -
+    exercises `_equi_not_range_join.py` via `_maybe_select_better_equi_
+    predicates`."""
+    df, right = _skewed_equi_broad_selective_frames(seed=0)
+    grp = ("grp", "grp", "==")
+    broad = ("l_broad", "r_broad", "<")
+    selective = ("l_selective", "r_selective", "<=")
+
+    bad_order = df.conditional_join(
+        right, grp, broad, selective, keep=keep, how="inner"
+    )
+    good_order = df.conditional_join(
+        right, grp, selective, broad, keep=keep, how="inner"
+    )
+    assert_frame_equal(bad_order, good_order)
+
+
+@pytest.mark.parametrize("keep", ["first", "last"])
+def test_equi_range_join_predicate_order_invariant(keep):
+    """Output for an equi join combined with a range join (2 eligible
+    `<`-type candidates for the upper bound) must not depend on argument
+    order - exercises `_equi_range_join.py`."""
+    df, right = _skewed_equi_range_frames(seed=1)
+    grp = ("grp", "grp", "==")
+    ge_gt = ("a", "lo", ">")
+    narrow = ("a", "hi_narrow", "<")
+    broad = ("a", "hi_broad", "<")
+
+    bad_order = df.conditional_join(
+        right, grp, ge_gt, broad, narrow, keep=keep, how="inner"
+    )
+    good_order = df.conditional_join(
+        right, grp, ge_gt, narrow, broad, keep=keep, how="inner"
+    )
+    assert_frame_equal(bad_order, good_order)
+
+
+def test_equi_predicate_selection_skipped_for_keep_all():
+    """`_maybe_select_better_equi_predicates` must never run for
+    `keep='all'` - mirrors #1658/#1663's scoping."""
+    df, right = _skewed_equi_broad_selective_frames(seed=2, n=50)
+    grp = ("grp", "grp", "==")
+    broad = ("l_broad", "r_broad", "<")
+    selective = ("l_selective", "r_selective", "<=")
+
+    with mock.patch(
+        "janitor.functions._conditional_join._get_indices_equi."
+        "_maybe_select_better_equi_predicates"
+    ) as patched:
+        df.conditional_join(right, grp, broad, selective, keep="all", how="inner")
+        patched.assert_not_called()
+
+
+def test_equi_predicate_selection_noop_with_single_candidate():
+    """With a single non-equi candidate (the common case), the selection
+    step must be a no-op."""
+    from janitor.functions._conditional_join import _get_indices_equi as m
+
+    df, right = _skewed_equi_broad_selective_frames(seed=3, n=50)
+    narrow = ("l_selective", "r_selective", "<=")
+
+    mapping = {"is_range_join": False, "le_or_ge": [narrow]}
+    result = m._maybe_select_better_equi_predicates(mapping, df, right)
+    assert result["le_or_ge"] == [narrow]
+
+
 @pytest.mark.turtle
 @settings(deadline=None, max_examples=10)
 @given(df=conditional_df(), right=conditional_right())