Repository: pyjanitor-devs/pyjanitor

Issue:
`conditional_join` fails on mac for `equi-join` and numba
# Brief Description

`conditional_join` fails when an equi join is present, a range join is involved, and numba is used as the computation backend.

# System Information

- Operating system: macOS
- Python version (required):3.11

# Minimally Reproducible Code

```
A = pd.DataFrame({'product': {17510: 84, 99969: 44, 51210: 92},
 'start_date': {17510: Timestamp('2014-06-08 00:00:00'),
  99969: Timestamp('2014-04-21 00:00:00'),
  51210: Timestamp('2015-05-25 00:00:00')},
 'end_date': {17510: Timestamp('2016-04-06 00:00:00'),
  99969: Timestamp('2014-07-17 00:00:00'),
  51210: Timestamp('2016-05-31 00:00:00')}})

A 
       product start_date   end_date
17510       84 2014-06-08 2016-04-06
99969       44 2014-04-21 2014-07-17
51210       92 2015-05-25 2016-05-31

(A
.conditional_join(
    A, 
    ('end_date','start_date','>='),
    ('start_date','end_date','<='), 
    ('product','product','==')
    )
)
```
# Error Messages



```
RecursionError: Failed in nopython mode pipeline (step: fuse parfors)
maximum recursion depth exceeded
```



Discussion:
works fine in dev - cant figure out why it failed in 0.26 🤷 

update:
- error stems from the python version (python 3.11). python 3.8 and 3.9 works fine. python version 3.10 upwards fails though for equi join + range join + numba

Resolving patch:
diff --git a/janitor/functions/_numba.py b/janitor/functions/_numba.py
--- a/janitor/functions/_numba.py
+++ b/janitor/functions/_numba.py
@@ -137,12 +137,17 @@ def _numba_equi_join(df, right, eqs, ge_gt, le_lt):
     # 	2	  3	    2	   2	       4
     #
     left_column, right_column, _ = eqs
-    left_arr = df[left_column]._values
+    # steal some perf here within the binary search
+    # search for uniques
+    # and later index them with left_positions
+    left_positions, left_arr = df[left_column].factorize(sort=False)
     right_arr = right[right_column]._values
     left_index = df.index._values
     right_index = right.index._values
     slice_starts = right_arr.searchsorted(left_arr, side="left")
+    slice_starts = slice_starts[left_positions]
     slice_ends = right_arr.searchsorted(left_arr, side="right")
+    slice_ends = slice_ends[left_positions]
     # check if there is a search space
     # this also lets us know if there are equi matches
     keep_rows = slice_starts < slice_ends
@@ -174,6 +179,22 @@ def _numba_equi_join(df, right, eqs, ge_gt, le_lt):
         le_strict = True if op == "<" else False
 
     if le_lt and ge_gt:
+        group = right.groupby(eqs[1])[le_lt[1]]
+        # is the last column (le_lt) monotonic increasing?
+        # fast path if it is
+        all_monotonic_increasing = all(
+            arr.is_monotonic_increasing for _, arr in group
+        )
+        if all_monotonic_increasing:
+            cum_max_arr = le_arr2[:]
+        else:
+            cum_max_arr = group.cummax()._values
+            if is_extension_array_dtype(cum_max_arr):
+                array_dtype = cum_max_arr.dtype.numpy_dtype
+                cum_max_arr = cum_max_arr.astype(array_dtype)
+            if is_datetime64_dtype(cum_max_arr):
+                cum_max_arr = cum_max_arr.view(np.int64)
+
         left_index, right_index = _numba_equi_join_range_join(
             left_index,
             right_index,
@@ -185,6 +206,8 @@ def _numba_equi_join(df, right, eqs, ge_gt, le_lt):
             le_arr1,
             le_arr2,
             le_strict,
+            all_monotonic_increasing,
+            cum_max_arr,
         )
 
     elif le_lt:
@@ -375,6 +398,8 @@ def _numba_equi_join_range_join(
     le_arr1,
     le_arr2,
     le_strict,
+    all_monotonic_increasing,
+    cum_max_arr,
 ):
     """
     Get indices for an equi join
@@ -400,26 +425,6 @@ def _numba_equi_join_range_join(
             ends[num] = slice_start + end
     if counts == length:
         return None, None
-    # cumulative array
-    # used to find the first possible match
-    # for the less than section below
-    max_arr = np.empty_like(le_arr2)
-    counter = 0  # are all groups monotonic increasing?
-    for num in prange(length):
-        slice_start = slice_starts[num]
-        slice_end = slice_ends[num]
-        r1 = le_arr2[slice_start:slice_end]
-        start = r1[0]
-        max_arr[slice_start] = start
-        if r1.size > 1:
-            for n in range(1, r1.size):
-                new_value = r1[n]
-                check = start < new_value
-                if check:
-                    start = new_value
-                else:
-                    counter += 1
-                max_arr[slice_start + n] = start
 
     if counts > 0:
         left_index = left_index[booleans]
@@ -432,14 +437,14 @@ def _numba_equi_join_range_join(
     length = left_index.size
     starts = np.empty(length, dtype=np.intp)
     booleans = np.ones(length, dtype=np.bool_)
-    if counter == 0:
+    if all_monotonic_increasing:
         sizes = np.empty(length, dtype=np.intp)
     counts = 0
     for num in prange(length):
         l1 = le_arr1[num]
         slice_start = slice_starts[num]
         slice_end = slice_ends[num]
-        r1 = max_arr[slice_start:slice_end]
+        r1 = cum_max_arr[slice_start:slice_end]
         start = np.searchsorted(r1, l1, side="left")
         if start < r1.size:
             if le_strict and (l1 == r1[start]):
@@ -449,7 +454,7 @@ def _numba_equi_join_range_join(
             booleans[num] = False
         else:
             starts[num] = slice_start + start
-            if counter == 0:
+            if all_monotonic_increasing:
                 sizes[num] = r1.size - start
     if counts == length:
         return None, None
@@ -458,7 +463,7 @@ def _numba_equi_join_range_join(
         le_arr1 = le_arr1[booleans]
         starts = starts[booleans]
         slice_ends = slice_ends[booleans]
-        if counter == 0:
+        if all_monotonic_increasing:
             sizes = sizes[booleans]
 
     slice_starts = starts
@@ -467,7 +472,7 @@ def _numba_equi_join_range_join(
     # no need to run a comparison
     # since all groups are monotonic increasing
     # simply create left and right indices
-    if counter == 0:
+    if all_monotonic_increasing:
         cum_sizes = np.cumsum(sizes)
         starts = np.empty(slice_ends.size, dtype=np.intp)
         starts[0] = 0
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
@@ -772,7 +772,6 @@ def _multiple_conditional_join_eq(
         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):
             return indices
 
@@ -1073,10 +1072,7 @@ def _range_indices(
     ]
 
     right_index = np.concatenate(right_index)
-    left_index = np.repeat(left_index, repeater)
-
-    if fastpath:
-        return left_index, right_index
+    left_index = left_index.repeat(repeater)
     # here we search for actual positions
     # where left_c is </<= right_c
     # safe to index the arrays, since we are picking the positions
diff --git a/janitor/functions/utils.py b/janitor/functions/utils.py
--- a/janitor/functions/utils.py
+++ b/janitor/functions/utils.py
@@ -817,7 +817,7 @@ def _less_than_indices(
         right = [arr.max() for arr in right]
         return left_index, right
     right = np.concatenate(right)
-    left = np.repeat(left_index, len_right - search_indices)
+    left = left_index.repeat(len_right - search_indices)
     return left, right
 
 
@@ -904,7 +904,7 @@ def _greater_than_indices(
         right = [arr.max() for arr in right]
         return left_index, right
     right = np.concatenate(right)
-    left = np.repeat(left_index, search_indices)
+    left = left_index.repeat(search_indices)
     return left, right
 
 


Test patch:
diff --git a/tests/functions/test_complete.py b/tests/functions/test_complete.py
--- a/tests/functions/test_complete.py
+++ b/tests/functions/test_complete.py
@@ -220,7 +220,7 @@ def test_all_strings_no_nulls(df):
     expected = (
         df.set_index(cols)  # noqa: PD013, PD010
         .unstack(cols[-1])  # noqa: PD010
-        .stack(dropna=False)  # noqa: PD013
+        .stack(future_stack=True)  # noqa: PD013
         .reset_index()
         .reindex(columns=columns)
     )
@@ -393,7 +393,7 @@ def test_tuple_column():
     expected = (
         df.set_index(columns)
         .unstack("group")
-        .stack(dropna=False)
+        .stack(future_stack=True)
         .reset_index()
         .reindex(columns=df.columns)
         .sort_values(columns, ignore_index=True)
@@ -441,7 +441,7 @@ def test_fill_value_scalar(taxonomy_df):
     expected = (
         taxonomy_df.set_index(["Year", "Taxon"])
         .unstack(fill_value=0)
-        .stack(dropna=False)
+        .stack(future_stack=True)
         .reset_index()
         .astype({"Taxon": "object"})
     )
@@ -468,7 +468,7 @@ def test_dict_tuple_callable(taxonomy_df):
         .unstack("Year")
         .droplevel(0, 1)
         .reindex(columns=range(1999, 2005))
-        .stack(dropna=False)
+        .stack(future_stack=True)
         .reset_index()
         .iloc[:, :-1]
         .reindex(columns=["Year", "Taxon", "Abundance"])
@@ -496,7 +496,7 @@ def test_dict_tuple(taxonomy_df):
         .unstack("Year")
         .droplevel(0, 1)
         .reindex(columns=range(1999, 2005))
-        .stack(dropna=False)
+        .stack(future_stack=True)
         .reset_index()
         .iloc[:, :-1]
         .reindex(columns=["Year", "Taxon", "Abundance"])
@@ -545,7 +545,7 @@ def test_explicit_scalar(fill_df):
     expected = (
         fill_df.set_index(columns)
         .unstack("group", fill_value=0)
-        .stack(dropna=False)
+        .stack(future_stack=True)
         .reset_index()
         .reindex(columns=fill_df.columns)
         .sort_values(columns, ignore_index=True)
@@ -570,7 +570,7 @@ def test_explicit_scalar_cat(fill_df):
     expected = (
         fill_df.set_index(columns)
         .unstack("group", fill_value=0)
-        .stack(dropna=False)
+        .stack(future_stack=True)
         .reset_index()
         .reindex(columns=fill_df.columns)
         .sort_values(columns, ignore_index=True)
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
@@ -3884,6 +3884,41 @@ def test_multiple_eqs(df, right):
     assert_frame_equal(expected, actual)
 
 
+@pytest.mark.turtle
+@settings(deadline=None, max_examples=10)
+@given(df=conditional_df(), right=conditional_right())
+def test_multiple_eqs_numba_range(df, right):
+    """Test output for multiple conditions."""
+
+    columns = ["B", "A", "E", "Floats", "Integers", "Dates"]
+    expected = (
+        df.merge(
+            right,
+            left_on=["A"],
+            right_on=["Integers"],
+            how="inner",
+            sort=False,
+        )
+        .loc[lambda df: df.E.lt(df.Dates) & df.B.gt(df.Floats), columns]
+        .sort_values(columns, ignore_index=True)
+    )
+    expected = expected.filter(columns)
+    actual = (
+        df[["B", "A", "E"]]
+        .conditional_join(
+            right[["Floats", "Integers", "Dates"]],
+            ("E", "Dates", "<"),
+            ("B", "Floats", ">"),
+            ("A", "Integers", "=="),
+            how="inner",
+            use_numba=True,
+        )
+        .sort_values(columns, ignore_index=True)
+    )
+
+    assert_frame_equal(expected, actual)
+
+
 @pytest.mark.turtle
 @settings(deadline=None, max_examples=10)
 @given(df=conditional_df(), right=conditional_right())
@@ -4309,3 +4344,35 @@ def test_timedelta_dtype():
     actual.index = range(len(actual))
 
     assert_frame_equal(expected, actual)
+
+
+# https://stackoverflow.com/q/61948103/7175713
+def test_numba_equi_extension_array():
+    """
+    Test output for equi join and numba
+    """
+    df1 = pd.DataFrame(
+        {"id": [1, 1, 1, 2, 2, 3], "value_1": [2, 5, 7, 1, 3, 4]}
+    )
+    df2 = pd.DataFrame(
+        {
+            "id": [1, 1, 1, 1, 2, 2, 2, 3],
+            "value_2A": [0, 3, 7, 12, 0, 2, 3, 1],
+            "value_2B": [1, 9, 5, 15, 1, 6, 4, 3],
+        }
+    )
+    df1["value_1"] = df1["value_1"].astype(pd.Int64Dtype())
+    df2["value_2A"] = df2["value_2A"].astype(pd.Int64Dtype())
+    df2["value_2B"] = df2["value_2B"].astype(pd.Int64Dtype())
+    expected = df1.merge(df2, on="id").query("value_2A < value_1 < value_2B")
+    expected.index = range(expected.index.size)
+    actual = df1.conditional_join(
+        df2,
+        ("id", "id", "=="),
+        ("value_1", "value_2A", ">"),
+        ("value_1", "value_2B", "<"),
+        right_columns="value*",
+        use_numba=True,
+    )
+
+    assert_frame_equal(expected, actual)