Repository: pyjanitor-devs/pyjanitor

Issue:
[ENH] implement full numba version of a single conditional_join
- currently the implementation is dependent on the non-numba implementation(generic_cond_join) - instead implement the core part completely in numba


Discussion:

Resolving patch:
diff --git a/janitor/functions/_numba.py b/janitor/functions/_numba.py
--- a/janitor/functions/_numba.py
+++ b/janitor/functions/_numba.py
@@ -7,12 +7,26 @@
 
 from janitor.functions.utils import (
     _generic_func_cond_join,
-    _JoinOperator,
+    _null_checks_cond_join,
     greater_than_join_types,
     less_than_join_types,
 )
 
 
+def _convert_to_numpy(left: np.ndarray, right: np.ndarray) -> tuple:
+    """
+    Ensure array is a numpy array.
+    """
+    if is_extension_array_dtype(left):
+        array_dtype = left.dtype.numpy_dtype
+        left = left.astype(array_dtype)
+        right = right.astype(array_dtype)
+    if is_datetime64_dtype(left):
+        left = left.view(np.int64)
+        right = right.view(np.int64)
+    return left, right
+
+
 def _numba_equi_join(df, right, eqs, ge_gt, le_lt):
     """
     Compute indices when an equi join is present.
@@ -146,13 +160,7 @@ def _numba_equi_join(df, right, eqs, ge_gt, le_lt):
         left_column, right_column, op = ge_gt
         ge_arr1 = df.loc[left_index, left_column]._values
         ge_arr2 = right[right_column]._values
-        if is_extension_array_dtype(ge_arr1):
-            array_dtype = ge_arr1.dtype.numpy_dtype
-            ge_arr1 = ge_arr1.astype(array_dtype)
-            ge_arr2 = ge_arr2.astype(array_dtype)
-        if is_datetime64_dtype(ge_arr1):
-            ge_arr1 = ge_arr1.view(np.int64)
-            ge_arr2 = ge_arr2.view(np.int64)
+        ge_arr1, ge_arr2 = _convert_to_numpy(left=ge_arr1, right=ge_arr2)
         ge_strict = True if op == ">" else False
 
     le_arr1 = None
@@ -162,13 +170,7 @@ def _numba_equi_join(df, right, eqs, ge_gt, le_lt):
         left_column, right_column, op = le_lt
         le_arr1 = df.loc[left_index, left_column]._values
         le_arr2 = right[right_column]._values
-        if is_extension_array_dtype(le_arr1):
-            array_dtype = le_arr1.dtype.numpy_dtype
-            le_arr1 = le_arr1.astype(array_dtype)
-            le_arr2 = le_arr2.astype(array_dtype)
-        if is_datetime64_dtype(le_arr1):
-            le_arr1 = le_arr1.view(np.int64)
-            le_arr2 = le_arr2.view(np.int64)
+        le_arr1, le_arr2 = _convert_to_numpy(left=le_arr1, right=le_arr2)
         le_strict = True if op == "<" else False
 
     if le_lt and ge_gt:
@@ -244,8 +246,6 @@ def _numba_equi_le_join(
             if le_strict and (l1 == r1[start]):
                 start = np.searchsorted(r1, l1, side="right")
         if start == r1.size:
-            start = -1
-            starts[num] = start
             counts += 1
             booleans[num] = False
         else:
@@ -321,8 +321,6 @@ def _numba_equi_ge_join(
             if ge_strict and (l1 == r1[end - 1]):
                 end = np.searchsorted(r1, l1, side="left")
         if end == 0:
-            end = -1
-            ends[num] = end
             counts += 1
             booleans[num] = False
         else:
@@ -396,8 +394,6 @@ def _numba_equi_join_range_join(
             if ge_strict and (l1 == r1[end - 1]):
                 end = np.searchsorted(r1, l1, side="left")
         if end == 0:
-            end = -1
-            ends[num] = end
             counts += 1
             booleans[num] = False
         else:
@@ -449,8 +445,6 @@ def _numba_equi_join_range_join(
             if le_strict and (l1 == r1[start]):
                 start = np.searchsorted(r1, l1, side="right")
         if start == r1.size:
-            start = -1
-            starts[num] = start
             counts += 1
             booleans[num] = False
         else:
@@ -1099,122 +1093,295 @@ def _get_indices_multiple(
     return l_index, r_index
 
 
-def _numba_single_join(
-    left: pd.Series,
-    right: pd.Series,
-    op: str,
-    keep: str,
-) -> tuple:
-    """Return matching indices for single non-equi join."""
-
-    outcome = _generic_func_cond_join(
-        left=left,
-        right=right,
-        op=op,
-        multiple_conditions=True,
-        keep=keep,
-    )
-
-    if (outcome is None) or (op == _JoinOperator.NOT_EQUAL.value):
-        return outcome
+@njit()
+def _less_than_index_numba(arr: np.ndarray, value, low, high) -> int:
+    """
+    Get the index position of a value in an array.
+    """
+    while low < high:
+        mid = low + (high - low) // 2
+        if arr[mid] < value:
+            low = mid + 1
+        else:
+            high = mid
+    return low
 
-    left_index, right_index, search_indices = outcome
 
-    if op in greater_than_join_types:
-        starts = np.zeros(shape=search_indices.size, dtype=np.int8)
-        ends = search_indices
-        counts = search_indices
-    else:
-        ends = np.full(
-            shape=search_indices.size,
-            dtype=np.intp,
-            fill_value=right_index.size,
-        )
-        starts = search_indices
-        counts = ends - starts
-    if keep == "all":
-        return _get_indices_single(
-            l_index=left_index,
-            r_index=right_index,
-            counts=counts,
-            starts=starts,
-            ends=ends,
-        )
-
-    if (
-        (keep == "first")
-        and (op in less_than_join_types)
-        and pd.Series(right_index).is_monotonic_increasing
-    ):
-        return left_index, right_index[search_indices]
-    if (
-        (keep == "last")
-        and (op in greater_than_join_types)
-        and pd.Series(right_index).is_monotonic_increasing
-    ):
-        return left_index, right_index[search_indices - 1]
-    if keep == "first":
-        right_index = _numba_single_non_equi_keep_first(
-            right_index, starts, ends
-        )
-    else:
-        right_index = _numba_single_non_equi_keep_last(
-            right_index, starts, ends
-        )
-    return left_index, right_index
+@njit()
+def _greater_than_index_numba(arr: np.ndarray, value, low, high) -> int:
+    """
+    Get the index position of a value in an array.
+    """
+    while low < high:
+        mid = low + (high - low) // 2
+        if value < arr[mid]:
+            high = mid
+        else:
+            low = mid + 1
+    return low
 
 
 @njit(parallel=True)
-def _numba_single_non_equi_keep_first(
-    right_index: np.ndarray, starts: np.ndarray, ends: np.ndarray
+def _greater_than_indices_numba(
+    left_array: np.ndarray,
+    right_array: np.ndarray,
+    left_index: np.ndarray,
+    right_index: np.ndarray,
+    right_is_sorted: bool,
+    keep: str,
+    any_nulls: bool,
+    strict: bool = False,
 ) -> np.ndarray:
     """
-    Generate all indices when keep = `first`
-    Applies only to >, >= , <, <= operators.
+    Get indices for greater than join
     """
-    r_index = np.empty(starts.size, np.intp)
-    for num in prange(starts.size):
-        indexer = slice(starts[num], ends[num])
-        r_index[num] = right_index[indexer].min()
-    return r_index
+    low = 0
+    high = len(right_array)
+    length = len(left_array)
+    counts = 0
+    exclude = 0
+    positions = np.empty(length, dtype=np.intp)
+    booleans = np.ones(length, dtype=np.bool_)
+    for num in prange(length):
+        value = left_array[num]
+        pos = _greater_than_index_numba(right_array, value, low, high)
+        if pos == low:
+            exclude += 1
+            booleans[num] = False
+            continue
+        if strict and (right_array[pos - 1] == value):
+            pos = _less_than_index_numba(right_array, value, low, pos - 1)
+            if pos == low:
+                booleans[num] = False
+                exclude += 1
+                continue
+        positions[num] = pos
+        counts += pos
+    if counts == 0:
+        return None, None
+
+    if exclude > 0:
+        positions = positions[booleans]
+        left_index = left_index[booleans]
+        booleans = None
+    length = len(positions)
+    if right_is_sorted and (keep == "last"):
+        right_arr = np.empty(length, dtype=np.intp)
+        if any_nulls:
+            for num in prange(length):
+                right_arr[num] = right_index[positions[num] - 1]
+            return left_index, right_arr
+        return left_index, positions - 1
+
+    if keep == "first":
+        right_arr = np.empty(length, dtype=np.intp)
+        for num in prange(length):
+            pos = positions[num]
+            minimum = right_index[0]
+            for n in range(pos):
+                value = right_index[n]
+                if value < minimum:
+                    minimum = value
+            right_arr[num] = minimum
+        return left_index, right_arr
+
+    if keep == "last":
+        right_arr = np.empty(length, dtype=np.intp)
+        for num in prange(length):
+            pos = positions[num]
+            maximum = right_index[0]
+            for n in range(pos):
+                value = right_index[n]
+                if value > maximum:
+                    maximum = value
+            right_arr[num] = maximum
+        return left_index, right_arr
+
+    starts = np.empty(length, dtype=np.intp)
+    starts[0] = 0
+    starts[1:] = np.cumsum(positions[:-1])
+    left_arr = np.empty(counts, dtype=np.intp)
+    right_arr = np.empty(counts, dtype=np.intp)
+    for num in prange(length):
+        start = starts[num]
+        pos = positions[num]
+        l_index = left_index[num]
+        for n in range(pos):
+            left_arr[start + n] = l_index
+            right_arr[start + n] = right_index[n]
+
+    return left_arr, right_arr
 
 
 @njit(parallel=True)
-def _numba_single_non_equi_keep_last(
-    right_index: np.ndarray, starts: np.ndarray, ends: np.ndarray
+def _less_than_indices_numba(
+    left_array: np.ndarray,
+    right_array: np.ndarray,
+    left_index: np.ndarray,
+    right_index: np.ndarray,
+    right_is_sorted: bool,
+    keep: str,
+    any_nulls: bool,
+    strict: bool = False,
 ) -> np.ndarray:
     """
-    Generate all indices when keep = `last`
-    Applies only to >, >= , <, <= operators.
+    Get indices for less than join
     """
-    r_index = np.empty(starts.size, np.intp)
-    for num in prange(starts.size):
-        indexer = slice(starts[num], ends[num])
-        r_index[num] = right_index[indexer].max()
-    return r_index
-
-
-@njit(cache=True, parallel=True)
-def _get_indices_single(
-    l_index: np.ndarray,
-    r_index: np.ndarray,
-    counts: int,
-    starts: np.ndarray,
-    ends: np.ndarray,
-):
-    """Compute indices when starts and ends are already known"""
-    lengths = np.cumsum(counts)
-    left_index = np.empty(lengths[-1], np.intp)
-    right_index = np.empty(lengths[-1], np.intp)
-    start_indices = np.empty(lengths.size, np.intp)
-    start_indices[0] = 0
-    start_indices[1:] = lengths[:-1]
-    for num in prange(lengths.size):
-        start = start_indices[num]
-        width = counts[num]
-        l_indexer = slice(start, start + width)
-        left_index[l_indexer] = l_index[num]
-        r_indexer = slice(starts[num], ends[num])
-        right_index[l_indexer] = r_index[r_indexer]
+    low = 0
+    high = len(right_array)
+    length = len(left_array)
+    positions = np.empty(length, dtype=np.intp)
+    booleans = np.ones(length, dtype=np.bool_)
+    counts = 0
+    exclude = 0
+    for num in prange(length):
+        value = left_array[num]
+        pos = _less_than_index_numba(right_array, value, low, high)
+        if pos == high:
+            booleans[num] = False
+            exclude += 1
+            continue
+        if strict and (right_array[pos] == value):
+            pos = _greater_than_index_numba(right_array, value, pos + 1, high)
+            if pos == high:
+                booleans[num] = False
+                exclude += 1
+                continue
+        positions[num] = pos
+        counts += high - pos
+    if counts == 0:
+        return None, None
 
-    return left_index, right_index
+    if exclude > 0:
+        positions = positions[booleans]
+        left_index = left_index[booleans]
+        booleans = None
+    length = len(positions)
+    if right_is_sorted and (keep == "first"):
+        right_arr = np.empty(length, dtype=np.intp)
+        if any_nulls:
+            for num in prange(length):
+                right_arr[num] = right_index[positions[num]]
+            return left_index, right_arr
+        return left_index, positions
+
+    if keep == "first":
+        right_arr = np.empty(length, dtype=np.intp)
+        for num in prange(length):
+            pos = positions[num]
+            minimum = right_index[pos]
+            for n in range(pos, high):
+                value = right_index[n]
+                if value < minimum:
+                    minimum = value
+            right_arr[num] = minimum
+        return left_index, right_arr
+
+    if keep == "last":
+        right_arr = np.empty(length, dtype=np.intp)
+        for num in prange(length):
+            pos = positions[num]
+            maximum = right_index[pos]
+            for n in range(pos, high):
+                value = right_index[n]
+                if value > maximum:
+                    maximum = value
+            right_arr[num] = maximum
+        return left_index, right_arr
+
+    starts = np.empty(length, dtype=np.intp)
+    starts[0] = 0
+    starts[1:] = np.cumsum((high - positions)[:-1])
+    left_arr = np.empty(counts, dtype=np.intp)
+    right_arr = np.empty(counts, dtype=np.intp)
+    for num in prange(length):
+        start = starts[num]
+        pos = positions[num]
+        l_index = left_index[num]
+        for n in range(high - pos):
+            left_arr[start + n] = l_index
+            right_arr[start + n] = right_index[pos + n]
+
+    return left_arr, right_arr
+
+
+def _numba_single_join(
+    left: pd.Series,
+    right: pd.Series,
+    op: str,
+    keep: str,
+) -> tuple:
+    """Return matching indices for single non-equi join."""
+
+    if op == "!=":
+        return _generic_func_cond_join(
+            left=left,
+            right=right,
+            op=op,
+            multiple_conditions=True,
+            keep=keep,
+        )
+
+    strict = True if op in {"<", ">"} else False
+
+    if op in less_than_join_types:
+        # quick break, avoiding the hassle
+        if left.min() > right.max():
+            return None
+
+        outcome = _null_checks_cond_join(left=left, right=right)
+        if not outcome:
+            return None
+        (
+            left,
+            right,
+            left_index,
+            right_index,
+            right_is_sorted,
+            any_nulls,
+        ) = outcome
+
+        left, right = _convert_to_numpy(left=left, right=right)
+        outcome = _less_than_indices_numba(
+            left_array=left,
+            right_array=right,
+            strict=strict,
+            left_index=left_index,
+            right_index=right_index,
+            keep=keep,
+            right_is_sorted=right_is_sorted,
+            any_nulls=any_nulls,
+        )
+
+    elif op in greater_than_join_types:
+        # quick break, avoiding the hassle
+        if left.max() < right.min():
+            return None
+
+        outcome = _null_checks_cond_join(left=left, right=right)
+        if not outcome:
+            return None
+        (
+            left,
+            right,
+            left_index,
+            right_index,
+            right_is_sorted,
+            any_nulls,
+        ) = outcome
+
+        left, right = _convert_to_numpy(left=left, right=right)
+        outcome = _greater_than_indices_numba(
+            left_array=left,
+            right_array=right,
+            strict=strict,
+            left_index=left_index,
+            right_index=right_index,
+            keep=keep,
+            right_is_sorted=right_is_sorted,
+            any_nulls=any_nulls,
+        )
+
+    if outcome[0] is None:
+        return None
+    return outcome
diff --git a/janitor/functions/conditional_join.py b/janitor/functions/conditional_join.py
--- a/janitor/functions/conditional_join.py
+++ b/janitor/functions/conditional_join.py
@@ -7,12 +7,11 @@
 import numpy as np
 import pandas as pd
 import pandas_flavor as pf
-from pandas.core.dtypes.common import (
+from pandas.api.types import (
     is_datetime64_dtype,
     is_dtype_equal,
     is_extension_array_dtype,
     is_numeric_dtype,
-    is_string_dtype,
     is_timedelta64_dtype,
 )
 from pandas.core.reshape.merge import _MergeOperation
@@ -281,6 +280,7 @@ def conditional_join(
             merge key is found in both DataFrames.
         force: If `True`, force the non-equi join conditions to execute before the equi join.
 
+
     Returns:
         A pandas DataFrame of the two merged Pandas objects.
     """  # noqa: E501
@@ -439,74 +439,35 @@ def _conditional_join_type_check(
     left_column: pd.Series, right_column: pd.Series, op: str, use_numba: bool
 ) -> None:
     """
-    Raise error if column type is not any of
-    numeric or timedelta or datetime or string.
+    Dtype check for columns in the join.
+    Checks are not conducted for the equi-join columns,
+    except when use_numba is set to True.
     """
 
     if (
-        (op == _JoinOperator.STRICTLY_EQUAL.value)
-        and use_numba
+        ((op != _JoinOperator.STRICTLY_EQUAL.value) or use_numba)
         and not is_numeric_dtype(left_column)
         and not is_datetime64_dtype(left_column)
         and not is_timedelta64_dtype(left_column)
     ):
         raise TypeError(
             "Only numeric, timedelta and datetime types "
-            "are supported in an equi-join "
-            "when use_numba is set to True"
+            "are supported in a non equi-join, "
+            "or if use_numba is set to True. "
+            f"{left_column.name} in condition "
+            f"({left_column.name}, {right_column.name}, {op}) "
+            f"has a dtype {left_column.dtype}."
         )
 
-    is_categorical_dtype = isinstance(left_column.dtype, pd.CategoricalDtype)
-
-    if not is_categorical_dtype:
-        permitted_types = {
-            is_datetime64_dtype,
-            is_numeric_dtype,
-            is_string_dtype,
-            is_timedelta64_dtype,
-        }
-        for func in permitted_types:
-            if func(left_column.dtype):
-                break
-        else:
-            raise TypeError(
-                "conditional_join only supports "
-                "string, category, numeric, "
-                "date (without timezone) - ",
-                "or timedelta dtypes."
-                f"'{left_column.name} is of type "
-                f"{left_column.dtype}.",
-            )
-
-    if is_categorical_dtype:
-        if not left_column.array._categories_match_up_to_permutation(
-            right_column.array
-        ):
-            raise TypeError(
-                f"'{left_column.name}' and '{right_column.name}' "
-                "should have the same categories, and the same order."
-            )
-    elif not is_dtype_equal(left_column, right_column):
+    if (
+        (op != _JoinOperator.STRICTLY_EQUAL.value) or use_numba
+    ) and not is_dtype_equal(left_column, right_column):
         raise TypeError(
             f"Both columns should have the same type - "
             f"'{left_column.name}' has {left_column.dtype} type;"
             f"'{right_column.name}' has {right_column.dtype} type."
         )
 
-    if (
-        (op != _JoinOperator.STRICTLY_EQUAL.value)
-        and not is_numeric_dtype(left_column)
-        and not is_datetime64_dtype(left_column)
-        and not is_timedelta64_dtype(left_column)
-    ):
-        raise TypeError(
-            "non-equi joins are supported "
-            "only for datetime, timedelta and numeric dtypes. "
-            f"{left_column.name} in condition "
-            f"({left_column.name}, {right_column.name}, {op}) "
-            f"has a dtype {left_column.dtype}."
-        )
-
     return None
 
 
@@ -572,7 +533,12 @@ def _conditional_join_compute(
     if len(conditions) > 1:
         if eq_check:
             result = _multiple_conditional_join_eq(
-                df, right, conditions, keep, use_numba, force
+                df,
+                right,
+                conditions,
+                keep,
+                use_numba,
+                force,
             )
         elif le_lt_check:
             result = _multiple_conditional_join_le_lt(
@@ -727,6 +693,7 @@ def _multiple_conditional_join_eq(
         for left_on, right_on, op in conditions:
             if op == _JoinOperator.STRICTLY_EQUAL.value:
                 eqs = (left_on, right_on, op)
+                break
 
         le_lt = None
         ge_gt = None
@@ -780,7 +747,23 @@ def _multiple_conditional_join_eq(
             return None
         if any_nulls.any():
             right_df = right.loc[~any_nulls]
-        right_df = right_df.sort_values(right_columns)
+        equi_col = right_columns[0]
+        # check if the first column is sorted
+        # if sorted, check if the second column is sorted
+        # per group in the first column
+        right_is_sorted = right_df[equi_col].is_monotonic_increasing
+        if right_is_sorted:
+            grp = right_df.groupby(equi_col, sort=False)
+            non_equi_col = right_columns[1]
+            # groupby.is_monotonic_increasing uses apply under the hood
+            # this circumvents the Series creation (which isn't required here)
+            # and just gets a sequence of booleans, before calling `all`
+            # to get a single True or False.
+            right_is_sorted = all(
+                arr.is_monotonic_increasing for _, arr in grp[non_equi_col]
+            )
+        if not right_is_sorted:
+            right_df = right_df.sort_values(right_columns)
         indices = _numba_equi_join(left_df, right_df, eqs, ge_gt, le_lt)
 
         if not rest or (indices is None):
diff --git a/janitor/functions/utils.py b/janitor/functions/utils.py
--- a/janitor/functions/utils.py
+++ b/janitor/functions/utils.py
@@ -697,6 +697,37 @@ class _JoinOperator(Enum):
 }
 
 
+def _null_checks_cond_join(
+    left: pd.Series, right: pd.Series
+) -> Union[tuple, None]:
+    """
+    Checks for nulls in the arrays before conducting binary search.
+
+    Relevant to _less_than_indices and _greater_than_indices
+    """
+    any_nulls = left.isna()
+    if any_nulls.all():
+        return None
+    if any_nulls.any():
+        left = left[~any_nulls]
+    any_nulls = right.isna()
+    if any_nulls.all():
+        return None
+    if any_nulls.any():
+        right = right[~any_nulls]
+    any_nulls = any_nulls.any()
+    right_is_sorted = right.is_monotonic_increasing
+    if not right_is_sorted:
+        right = right.sort_values(kind="stable")
+
+    left_index = left.index._values
+    left = left._values
+    right_index = right.index._values
+    right = right._values
+
+    return left, right, left_index, right_index, right_is_sorted, any_nulls
+
+
 def _less_than_indices(
     left: pd.Series,
     right: pd.Series,
@@ -720,25 +751,10 @@ def _less_than_indices(
     if left.min() > right.max():
         return None
 
-    any_nulls = left.isna()
-    if any_nulls.all():
+    outcome = _null_checks_cond_join(left=left, right=right)
+    if not outcome:
         return None
-    if any_nulls.any():
-        left = left[~any_nulls]
-    any_nulls = right.isna()
-    if any_nulls.all():
-        return None
-    if any_nulls.any():
-        right = right[~any_nulls]
-    any_nulls = any_nulls.any()
-    right_is_sorted = right.is_monotonic_increasing
-    if not right_is_sorted:
-        right = right.sort_values(kind="stable")
-
-    left_index = left.index._values
-    left = left._values
-    right_index = right.index._values
-    right = right._values
+    left, right, left_index, right_index, right_is_sorted, any_nulls = outcome
 
     search_indices = right.searchsorted(left, side="left")
 
@@ -830,25 +846,10 @@ def _greater_than_indices(
     if left.max() < right.min():
         return None
 
-    any_nulls = left.isna()
-    if any_nulls.all():
-        return None
-    if any_nulls.any():
-        left = left[~any_nulls]
-    any_nulls = right.isna()
-    if any_nulls.all():
+    outcome = _null_checks_cond_join(left=left, right=right)
+    if not outcome:
         return None
-    if any_nulls.any():
-        right = right[~any_nulls]
-    any_nulls = any_nulls.any()
-    right_is_sorted = right.is_monotonic_increasing
-    if not right_is_sorted:
-        right = right.sort_values(kind="stable")
-
-    left_index = left.index._values
-    left = left._values
-    right_index = right.index._values
-    right = right._values
+    left, right, left_index, right_index, right_is_sorted, any_nulls = outcome
 
     search_indices = right.searchsorted(left, side="right")
     # if any of the positions in `search_indices`
@@ -1022,7 +1023,7 @@ def _generic_func_cond_join(
             keep=keep,
         )
     if op == _JoinOperator.NOT_EQUAL.value:
-        return _not_equal_indices(left, right, keep)
+        return _not_equal_indices(left=left, right=right, keep=keep)
 
 
 def _keep_output(keep: str, left: np.ndarray, right: np.ndarray):


Test patch:
diff --git a/tests/functions/test_conditional_join.py b/tests/functions/test_conditional_join.py
--- a/tests/functions/test_conditional_join.py
+++ b/tests/functions/test_conditional_join.py
@@ -308,23 +308,6 @@ def test_check_keep_value(dummy, series):
         dummy.conditional_join(series, ("id", "B", "<"), keep="ALL")
 
 
-def test_unequal_categories(dummy):
-    """
-    Raise TypeError if the dtypes are both categories
-    and do not match.
-    """
-    match = "'S' and 'Strings' should have the same categories,"
-    match = match + " and the same order."
-    with pytest.raises(TypeError, match=match):
-        dummy.astype({"S": "category"}).conditional_join(
-            dummy.rename(columns={"S": "Strings"}).encode_categorical(
-                Strings="appearance"
-            ),
-            ("S", "Strings", "=="),
-            ("id", "value_1", "<="),
-        )
-
-
 def test_dtype_not_permitted(dummy, series):
     """
     Raise TypeError if dtype of column in `df`
@@ -333,8 +316,9 @@ def test_dtype_not_permitted(dummy, series):
     dummy["F"] = pd.IntervalIndex.from_tuples(
         [(0, 10), (10, 20), (20, 30), (30, 40), (40, 50), (50, 60)]
     )
-    match = "conditional_join only supports string, "
-    match = match + "category, numeric, date.+"
+    match = "Only numeric, timedelta and datetime types "
+    match = "are supported in a non equi-join, "
+    match = "or if use_numba is set to True.+"
     with pytest.raises(TypeError, match=match):
         dummy.conditional_join(series, ("F", "B", "<"))
 
@@ -344,9 +328,10 @@ def test_dtype_str(dummy, series):
     Raise TypeError if dtype of column in `df`
     does not match the dtype of column from `right`.
     """
-    with pytest.raises(
-        TypeError, match="Both columns should have the same type.+"
-    ):
+    match = "Only numeric, timedelta and datetime types "
+    match = "are supported in a non equi-join, "
+    match = "or if use_numba is set to True.+"
+    with pytest.raises(TypeError, match=match):
         dummy.conditional_join(series, ("S", "B", "<"))
 
 
@@ -355,8 +340,9 @@ def test_dtype_strings_non_equi(dummy):
     Raise TypeError if the dtypes are both strings
     on a non-equi operator.
     """
-    match = "non-equi joins are supported only "
-    match = match + "for datetime, timedelta and numeric dtypes.+"
+    match = "Only numeric, timedelta and datetime types "
+    match = "are supported in a non equi-join, "
+    match = "or if use_numba is set to True.+"
     with pytest.raises(
         TypeError,
         match=match,
@@ -371,16 +357,27 @@ def test_dtype_category_non_equi():
     Raise TypeError if dtype is category,
     and op is non-equi.
     """
-    match = (
-        "non-equi joins are supported only for datetime, "
-        "timedelta and numeric dtypes.+"
-    )
+    match = "Only numeric, timedelta and datetime types "
+    match = "are supported in a non equi-join, "
+    match = "or if use_numba is set to True.+"
     with pytest.raises(TypeError, match=match):
         left = pd.DataFrame({"A": [1, 2, 3]}, dtype="category")
         right = pd.DataFrame({"B": [1, 2, 3]}, dtype="category")
         left.conditional_join(right, ("A", "B", "<"))
 
 
+def test_dtype_different_non_equi():
+    """
+    Raise TypeError if dtype is different,
+    and op is non-equi.
+    """
+    match = "Both columns should have the same type.+"
+    with pytest.raises(TypeError, match=match):
+        left = pd.DataFrame({"A": [1, 2, 3]}, dtype="int64")
+        right = pd.DataFrame({"B": [1, 2, 3]}, dtype="int8")
+        left.conditional_join(right, ("A", "B", "<"))
+
+
 @pytest.mark.turtle
 @settings(deadline=None, max_examples=10)
 @given(df=conditional_df(), right=conditional_right())
@@ -415,9 +412,114 @@ def test_single_condition_less_than_floats_keep_first(df, right):
 @pytest.mark.turtle
 @settings(deadline=None, max_examples=10)
 @given(df=conditional_df(), right=conditional_right())
-def test_single_condition_less_than_floats_keep_last(df, right):
+def test_single_condition_greater_than_floats_keep_last(df, right):
     """Test output for a single condition. "<"."""
 
+    df = df.sort_values("B").dropna(subset=["B"])
+    expected = pd.merge_asof(
+        df[["B"]],
+        right[["Numeric"]].sort_values("Numeric").dropna(subset=["Numeric"]),
+        left_on="B",
+        right_on="Numeric",
+        direction="backward",
+        allow_exact_matches=False,
+    )
+    expected.index = range(len(expected))
+    actual = (
+        df[["B"]]
+        .conditional_join(
+            right[["Numeric"]].sort_values("Numeric"),
+            ("B", "Numeric", ">"),
+            how="left",
+            sort_by_appearance=False,
+            keep="last",
+        )
+        .sort_values(["B", "Numeric"], ignore_index=True)
+    )
+
+    assert_frame_equal(expected, actual)
+
+
+@pytest.mark.turtle
+@settings(deadline=None, max_examples=10)
+@given(df=conditional_df(), right=conditional_right())
+def test_single_condition_greater_than_floats_keep_last_numba(df, right):
+    """
+    Test the functionality of conditional_join with a single
+    'greater than' condition on floating-point data,
+    while keeping the last match using Numba.
+
+    This test sorts and filters dataframes 'df' and 'right'
+    by columns 'B' and 'Numeric' respectively,
+    removing NaN values.
+    It then performs a backward merge_asof operation
+    on these sorted dataframes.
+    The expected outcome is a dataframe
+    where each row from 'df' is merged
+    with the last row from 'right'
+      where 'Numeric' is greater than 'B'.
+
+    The actual outcome is produced
+    by the conditional_join method
+    with a 'greater than' condition,
+    left join type, sorted by appearance,
+    keeping the last match,
+    and utilizing Numba for performance optimization.
+    The test asserts that the actual dataframe matches
+    the expected dataframe,
+    ensuring correct functionality of the conditional_join
+    under these specific parameters.
+    """
+
+    df = df.sort_values("B").dropna(subset=["B"])
+    expected = pd.merge_asof(
+        df[["B"]],
+        right[["Numeric"]].sort_values("Numeric").dropna(subset=["Numeric"]),
+        left_on="B",
+        right_on="Numeric",
+        direction="backward",
+        allow_exact_matches=False,
+    )
+    expected.index = range(len(expected))
+    actual = (
+        df[["B"]]
+        .conditional_join(
+            right[["Numeric"]].sort_values("Numeric"),
+            ("B", "Numeric", ">"),
+            how="left",
+            sort_by_appearance=False,
+            keep="last",
+            use_numba=True,
+        )
+        .sort_values(["B", "Numeric"], ignore_index=True)
+    )
+
+    assert_frame_equal(expected, actual)
+
+
+@pytest.mark.turtle
+@settings(deadline=None, max_examples=10)
+@given(df=conditional_df(), right=conditional_right())
+def test_single_condition_less_than_floats_keep_last(df, right):
+    """
+    Test the functionality of conditional_join
+    with a single 'greater than' condition on floating-point data,
+    while keeping the last match using Numba.
+
+    This test sorts and filters dataframes 'df' and 'right'
+    by columns 'B' and 'Numeric' respectively, removing NaN values.
+    It then performs a backward merge_asof operation on these sorted dataframes.
+    The expected outcome is a dataframe where each row from 'df'
+    is merged with the last row from 'right' where 'Numeric' is greater than 'B'.
+
+    The actual outcome is produced by the conditional_join method
+    with a 'greater than' condition, left join type, sorted by appearance,
+    keeping the last match, without utilizing Numba for performance optimization.
+    The test asserts that the actual dataframe matches the expected dataframe,
+    ensuring correct functionality of the conditional_join
+    under these specific parameters.
+    """
+
     df = df.sort_values("B").dropna(subset=["B"])
     right = right.sort_values("Numeric").dropna(subset=["Numeric"])
     expected = pd.merge_asof(