[PERF] Avoid materializing not-equal pairs for keep first or last

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

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)

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.
 

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())