diff --git a/CHANGELOG.md b/CHANGELOG.md index 2495389f7..a77a93d39 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/janitor/functions/_conditional_join/_get_indices_equi.py b/janitor/functions/_conditional_join/_get_indices_equi.py index d3dfdc1cd..393a1fd04 100644 --- a/janitor/functions/_conditional_join/_get_indices_equi.py +++ b/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 diff --git a/tests/functions/test_conditional_join.py b/tests/functions/test_conditional_join.py index b49170729..19c1a97d6 100644 --- a/tests/functions/test_conditional_join.py +++ b/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())