diff --git a/janitor/functions/_conditional_join/_compare.py b/janitor/functions/_conditional_join/_compare.py index 7b05f6af9..52701e567 100644 --- a/janitor/functions/_conditional_join/_compare.py +++ b/janitor/functions/_conditional_join/_compare.py @@ -588,3 +588,40 @@ def _compare_positions_ne( is_extension_array=is_extension_array, op=op, ) + + +def _select_start_end_direct( + left: list[np.ndarray], + right: list[np.ndarray], + left_index: np.ndarray, + right_index: np.ndarray, + starts: np.ndarray, + ends: np.ndarray, + ops: list[int], + first: bool, +) -> tuple[np.ndarray, np.ndarray] | None: + """Select first/last complete matches without allocating a tape. + + This is deliberately a capability probe: older janitor-rs wheels do not + expose the direct kernel yet, in which case callers retain the matches- + tape implementation. + """ + mapping = { + "int64": "select_start_end_direct_int64", + "int32": "select_start_end_direct_int32", + "int16": "select_start_end_direct_int16", + "int8": "select_start_end_direct_int8", + "uint64": "select_start_end_direct_uint64", + "uint32": "select_start_end_direct_uint32", + "uint16": "select_start_end_direct_uint16", + "uint8": "select_start_end_direct_uint8", + "float64": "select_start_end_direct_f64", + "float32": "select_start_end_direct_f32", + } + dtype_name = left[0].dtype.name + try: + function_name = mapping[dtype_name] + func = getattr(janitor_rs, function_name) + except (KeyError, AttributeError): + return None + return func(left, right, left_index, right_index, starts, ends, ops, first) diff --git a/janitor/functions/_conditional_join/_equi_range_join.py b/janitor/functions/_conditional_join/_equi_range_join.py index 730ff5699..2269b1046 100644 --- a/janitor/functions/_conditional_join/_equi_range_join.py +++ b/janitor/functions/_conditional_join/_equi_range_join.py @@ -115,6 +115,17 @@ def _get_indices( keep=keep, right_is_sorted=check and is_sorted, ) + direct = _helpers._get_direct_indices_conditions( + df=df, + right=right, + conditions=rest, + left_index=left_index, + starts=starts, + ends=ends, + keep=keep, + ) + if direct is not None: + return direct outcome = _helpers._get_positive_matches_conditions( df=df, right=right, diff --git a/janitor/functions/_conditional_join/_helpers.py b/janitor/functions/_conditional_join/_helpers.py index 28dcdabec..3fd47d646 100644 --- a/janitor/functions/_conditional_join/_helpers.py +++ b/janitor/functions/_conditional_join/_helpers.py @@ -528,6 +528,62 @@ def _get_positive_matches_conditions( return {"matches": matches, "counts_array": counts_array, "total": total} +def _get_direct_indices_conditions( + df: pd.DataFrame, + right: pd.DataFrame, + conditions: list, + left_index: np.ndarray, + starts: np.ndarray, + ends: np.ndarray, + keep: str, +): + """Try direct first/last selection for safe homogeneous inputs. + + The direct janitor-rs kernel has one concrete dtype parameter for all + predicate columns. Mixed dtypes, null-bearing columns, and unavailable + extension exports deliberately return ``None`` so the caller can use the + established matches-tape path with identical pandas semantics. + """ + if keep not in {"first", "last"}: + return None + + left_arrays = [] + right_arrays = [] + ops = [] + dtype_name = None + for left_on, right_on, op in conditions: + left_series = df.loc[left_index, left_on] + right_series = right[right_on] + if left_series.isna().any() or right_series.isna().any(): + return None + left_array = _convert_array_to_numpy(array=left_series._values) + right_array = _convert_array_to_numpy(array=right_series._values) + if left_array.dtype != right_array.dtype: + return None + if dtype_name is None: + dtype_name = left_array.dtype.name + elif dtype_name != left_array.dtype.name: + return None + left_arrays.append(left_array) + right_arrays.append(right_array) + ops.append(operator_mapping[op]) + + result = _compare._select_start_end_direct( + left=left_arrays, + right=right_arrays, + left_index=left_index, + right_index=right.index._values, + starts=starts, + ends=ends, + ops=ops, + first=keep == "first", + ) + if result is None: + return None + selected_left, selected_right = result + return {"left_index": selected_left, "right_index": selected_right} + + def _get_boolean_args_for_ne( op: str, left: np.ndarray | None, right: np.ndarray | None ) -> tuple: diff --git a/janitor/functions/_conditional_join/_range_join_default.py b/janitor/functions/_conditional_join/_range_join_default.py index 651d3e20d..f5b63809c 100644 --- a/janitor/functions/_conditional_join/_range_join_default.py +++ b/janitor/functions/_conditional_join/_range_join_default.py @@ -53,6 +53,17 @@ def _get_indices( keep=keep, right_is_sorted=right_is_sorted, ) + direct = _helpers._get_direct_indices_conditions( + df=df, + right=right, + conditions=rest, + left_index=outcome["left_index"], + starts=outcome["starts"], + ends=outcome["ends"], + keep=keep, + ) + if direct is not None: + return direct out = _helpers._get_positive_matches_conditions( df=df, right=right, diff --git a/tests/functions/test_conditional_join.py b/tests/functions/test_conditional_join.py index 3f4241434..5327db50d 100644 --- a/tests/functions/test_conditional_join.py +++ b/tests/functions/test_conditional_join.py @@ -10,7 +10,7 @@ from pandas.testing import assert_frame_equal import janitor as jn -from janitor.functions._conditional_join import _le_ge_1_or_more +from janitor.functions._conditional_join import _helpers, _le_ge_1_or_more from janitor.testing_utils.strategies import ( conditional_df, conditional_right, @@ -118,6 +118,36 @@ def test_multiple_conditions_preserve_non_condition_columns(): assert_frame_equal(actual, expected) +def test_reverse_join_agg_handles_empty_and_non_empty_candidate_ranges(): + """A zero-width candidate row must not corrupt the matches tape.""" + left = pd.DataFrame({"value": [0, 10], "payload": [4, 6]}) + right = pd.DataFrame( + { + "lower": [0, 2], + "upper": [5, 20], + } + ) + + actual = left.join_agg( + right, + ("value", "lower", ">"), + ("value", "upper", "<"), + reverse=True, + aggfunc=[("payload", "sum"), ("payload", "size")], + ) + + expected = pd.DataFrame( + { + ("payload", "sum"): [6], + ("payload", "size"): [1], + }, + index=pd.Index([1]), + ) + expected.columns = pd.MultiIndex.from_tuples(expected.columns) + + assert_frame_equal(actual, expected) + + def test_df_columns_right_columns_both_None(dummy, series): """Raise if both df_columns and right_columns is None""" with pytest.raises( @@ -518,6 +548,76 @@ def test_dtype_different_non_equi(): left.conditional_join(right, ("A", "B", "<")) +def test_direct_selection_kernel_is_used_for_safe_range_rest(monkeypatch): + """Use the tape-free kernel only for homogeneous, non-null inputs.""" + calls = [] + + def fake_direct(left, right, left_index, right_index, starts, ends, ops, first): + calls.append((left, right, left_index, right_index, starts, ends, ops, first)) + return np.array([0], dtype=np.int64), np.array([1], dtype=np.int64) + + monkeypatch.setattr( + _helpers.janitor_rs, + "select_start_end_direct_int64", + fake_direct, + raising=False, + ) + left = pd.DataFrame({"start": [2], "payload": ["left"]}) + right = pd.DataFrame( + { + "lower": [1, 2, 3], + "upper": [3, 3, 3], + "tag": [7, 7, 8], + "payload": ["zero", "one", "two"], + } + ) + + actual = left.conditional_join( + right, + ("start", "lower", ">="), + ("start", "upper", "<="), + ("start", "tag", "<="), + keep="first", + ) + + assert len(calls) == 1 + assert calls[0][-1] is True + assert list(actual[("right", "payload")]) == ["one"] + + +def test_keep_all_retains_every_match_without_direct_selection(monkeypatch): + """The direct kernel must remain limited to first/last selection.""" + + def fail_if_called(*args, **kwargs): + raise AssertionError("keep='all' must use the matches-tape path") + + monkeypatch.setattr( + _helpers.janitor_rs, + "select_start_end_direct_int64", + fail_if_called, + raising=False, + ) + left = pd.DataFrame({"start": [2]}) + right = pd.DataFrame( + { + "lower": [1, 2, 3], + "upper": [3, 3, 3], + "tag": [7, 7, 8], + "payload": ["zero", "one", "two"], + } + ) + + actual = left.conditional_join( + right, + ("start", "lower", ">="), + ("start", "upper", "<="), + ("start", "tag", "<="), + keep="all", + ) + + assert list(actual["payload"]) == ["zero", "one"] + + @pytest.mark.parametrize("op", ["<", "<=", ">", ">="]) @pytest.mark.parametrize("keep", ["first", "last"]) def test_single_inequality_unsorted_right_keep(op, keep):