diff --git a/CHANGELOG.md b/CHANGELOG.md index 9877dc507..85d93f5d2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,8 @@ # Changelog ## [Unreleased] +- [ENH] Avoid materializing all unequal pairs in `conditional_join` with + `keep="first"` or `keep="last"`. - Issue #1651, PR #1681 @tunglambk - [PERF] Reduce peak memory in `_build_indexer_reorder_contents` for wide frames (reps >= 8) using single NumPy allocation; tall frames with few repetitions retain the original reshape path. - Issue #1655 @Anupam2400 - [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/_not_equal_indices.py b/janitor/functions/_conditional_join/_not_equal_indices.py index f0489b80c..b59822554 100644 --- a/janitor/functions/_conditional_join/_not_equal_indices.py +++ b/janitor/functions/_conditional_join/_not_equal_indices.py @@ -24,6 +24,9 @@ def _not_equal_indices( and strictly greater than indices. """ + if keep in {"first", "last"}: + return _not_equal_keep_one(left=left, right=right, keep=keep) + dummy = np.array([], dtype=np.intp) # deal with nulls @@ -106,3 +109,89 @@ def _not_equal_indices( outcome = _keep_output(keep, left, right) left_index, right_index = outcome return {"left_index": left_index, "right_index": right_index} + + +def _not_equal_keep_one( + left: pd.Series, + right: pd.Series, + keep: str, +) -> dict: + """Return the first or last unequal right position for each left row.""" + dummy = np.array([], dtype=np.intp) + if left.empty or right.empty: + return {"left_index": dummy, "right_index": dummy} + + # The first unequal item is either the first item or, when those values + # match, the first item with a different value. The same observation works + # backwards for ``keep="last"``, so only two right-side candidates are + # needed. + reverse = keep == "last" + first_offset = -1 if reverse else 0 + first_value = right.iloc[first_offset] + first_is_null = pd.isna(first_value) + + if first_is_null: + first_equals_right = right.isna() + first_equals_left = pd.Series(False, index=left.index) + else: + first_equals_right = right.eq(first_value).fillna(False) + first_equals_left = left.eq(first_value).fillna(False) + + first_equals_right = first_equals_right.to_numpy(dtype=bool, na_value=False) + different_offsets = np.flatnonzero(~first_equals_right) + second_position = None + if different_offsets.size: + second_offset = different_offsets[-1] if reverse else different_offsets[0] + second_position = right.index[second_offset] + + left_index = left.index.to_numpy(copy=False) + first_equals_left = first_equals_left.to_numpy(dtype=bool, na_value=False) + if second_position is None: + keep_rows = ~first_equals_left + left_index = left_index[keep_rows] + right_index = np.full(left_index.size, right.index[first_offset], dtype=np.intp) + else: + right_index = np.where( + first_equals_left, + second_position, + right.index[first_offset], + ).astype(np.intp, copy=False) + + left_order = _not_equal_left_order(left=left, right=right) + selected = pd.Index(left_index).get_indexer(left_order) + selected = selected[selected >= 0] + left_index = left_index[selected] + right_index = right_index[selected] + + return {"left_index": left_index, "right_index": right_index} + + +def _not_equal_left_order(left: pd.Series, right: pd.Series) -> np.ndarray: + """Return left positions in the order produced by the materialized join.""" + left_nulls = left.isna().to_numpy() + right_nulls = right.isna().to_numpy() + seen = np.zeros(left.size, dtype=bool) + order = [] + nonnull_right = right[~right_nulls] + + if not nonnull_right.empty: + less_than_max = left.lt(nonnull_right.max()).fillna(False) + less_than_max = less_than_max.to_numpy(dtype=bool, na_value=False) + less_than_max = less_than_max & ~left_nulls + order.append(left.index[less_than_max].to_numpy(copy=False)) + seen |= less_than_max + + greater_than_min = left.gt(nonnull_right.min()).fillna(False) + greater_than_min = greater_than_min.to_numpy(dtype=bool, na_value=False) + greater_than_min = greater_than_min & ~left_nulls & ~seen + order.append(left.index[greater_than_min].to_numpy(copy=False)) + seen |= greater_than_min + + unseen_nulls = left_nulls & ~seen + order.append(left.index[unseen_nulls].to_numpy(copy=False)) + seen |= unseen_nulls + + if right_nulls.any(): + order.append(left.index[~seen].to_numpy(copy=False)) + + return np.concatenate(order) diff --git a/janitor/functions/conditional_join.py b/janitor/functions/conditional_join.py index a3f2f5c14..067e349fb 100644 --- a/janitor/functions/conditional_join.py +++ b/janitor/functions/conditional_join.py @@ -99,7 +99,9 @@ def conditional_join( The operator can be any of `==`, `!=`, `<=`, `<`, `>=`, `>`. - There is no optimisation for the `!=` operator. + For a single `!=` condition with `keep="first"` or `keep="last"`, + matching positions are selected without materializing all unequal pairs. + Other `!=` joins are not optimized. The join is done only on the columns. diff --git a/tests/functions/test_conditional_join.py b/tests/functions/test_conditional_join.py index 3f4241434..26a0069f5 100644 --- a/tests/functions/test_conditional_join.py +++ b/tests/functions/test_conditional_join.py @@ -1427,6 +1427,113 @@ def test_single_condition_not_equal_datetime(df, right): assert_frame_equal(expected, actual) +@pytest.mark.parametrize("keep", ["first", "last"]) +@pytest.mark.parametrize( + ("left", "right"), + [ + ( + pd.Series([1, 2, 3, 4], dtype="int64"), + pd.Series([2, 1, 2, 3], dtype="int64"), + ), + ( + pd.Series([2, 3], dtype="Int64"), + pd.Series([2, 2, 2], dtype="Int64"), + ), + ( + pd.Series([pd.NA, 2, 3, 4], dtype="Int64"), + pd.Series([pd.NA, 2, pd.NA, 3], dtype="Int64"), + ), + ( + pd.Series(pd.to_datetime([None, "2025-01-02", "2025-01-03"])), + pd.Series(pd.to_datetime(["2025-01-03", None, "2025-01-02"])), + ), + ( + pd.Series([np.nan, 2.0, 3.0]), + pd.Series([2.0, np.nan, 3.0]), + ), + ( + pd.Series([np.nan, np.nan]), + pd.Series([np.nan, np.nan]), + ), + ( + pd.Series([], dtype="float64"), + pd.Series([1.0, 2.0]), + ), + ( + pd.Series([1.0, 2.0]), + pd.Series([], dtype="float64"), + ), + ], + ids=[ + "duplicates", + "one-distinct-value", + "nullable-integers", + "datetimes", + "float-nan", + "all-null", + "empty-left", + "empty-right", + ], +) +def test_single_condition_not_equal_keep_one(left, right, keep): + """First and last selection match a direct scan in original row order.""" + expected_left = [] + expected_right = [] + right_positions = range(len(right)) + if keep == "last": + right_positions = reversed(right_positions) + + for left_position, left_value in enumerate(left): + for right_position in right_positions: + right_value = right.iloc[right_position] + left_is_null = pd.isna(left_value) + right_is_null = pd.isna(right_value) + if left_is_null or right_is_null: + unequal = True + else: + unequal = left_value != right_value + if unequal: + expected_left.append(left_position) + expected_right.append(right_position) + break + right_positions = range(len(right)) + if keep == "last": + right_positions = reversed(right_positions) + + left_frame = left.to_frame("left").assign(left_position=range(len(left))) + right_frame = right.to_frame("right").assign(right_position=range(len(right))) + actual = left_frame.conditional_join( + right_frame, + ("left", "right", "!="), + keep=keep, + ) + + actual = actual.sort_values("left_position", ignore_index=True) + assert actual["left_position"].tolist() == expected_left + assert actual["right_position"].tolist() == expected_right + + +@pytest.mark.parametrize( + ("keep", "right_positions"), + [("first", [1, 0, 0, 0]), ("last", [3, 3, 3, 2])], +) +def test_single_condition_not_equal_keep_one_preserves_output_order( + keep, right_positions +): + """Optimized selection preserves the legacy materialized-pair order.""" + left = pd.DataFrame({"left": [2, -1, 3, 0], "left_position": range(4)}) + right = pd.DataFrame({"right": [2, -3, -2, 3], "right_position": range(4)}) + + actual = left.conditional_join( + right, + ("left", "right", "!="), + keep=keep, + ) + + assert actual["left_position"].tolist() == [0, 1, 3, 2] + assert actual["right_position"].tolist() == right_positions + + @pytest.mark.turtle @settings(deadline=None, max_examples=10) @given(df=conditional_df(), right=conditional_right())