Tidy ups

Author: adamreeve

crates/polars-ops/src/frame/join/iejoin/filtered_bit_array.rs

@@ -2,8 +2,9 @@ use std::cmp::min;
 
 use arrow::bitmap::MutableBitmap;
 
-/// Bit array with a filter to speed up searching for set bits,
-/// based on section 4.1 in Khayyat et al.
+/// Bit array with a filter to speed up searching for set bits when sparse,
+/// based on section 4.1 from Khayyat et al. 2015,
+/// "Lightning Fast and Space Efficient Inequality Joins"
 pub struct FilteredBitArray {
     bit_array: MutableBitmap,
     filter: MutableBitmap,

crates/polars-ops/src/frame/join/iejoin/mod.rs

@@ -33,11 +33,11 @@ pub struct IEJoinOptions {
     pub operator2: InequalityOperator,
 }
 
-/// Inequality join. Matches rows from this DataFrame with rows from another DataFrame
-/// using two inequality operators (one of [<, <=, >, >=]).
+/// Inequality join. Matches rows between two DataFrames using two inequality operators
+/// (one of [<, <=, >, >=]).
 /// Based on Khayyat et al. 2015, "Lightning Fast and Space Efficient Inequality Joins"
 /// and extended to work with duplicate values.
-pub fn join_dataframes(
+pub fn iejoin(
     left: &DataFrame,
     right: &DataFrame,
     selected_left: Vec<Series>,
@@ -221,6 +221,8 @@ trait L1Array {
     fn mark_visited(&self, index: usize, bit_array: &mut FilteredBitArray);
 }
 
+/// Find the position in the L1 array where we should begin checking for matches,
+/// given the index in L1 corresponding to the current position in L2.
 fn find_search_start_index<T>(
     l1_array: &[L1Item<T>],
     index: usize,
@@ -284,6 +286,39 @@ where
     }
 }
 
+fn find_matches_in_l1<T>(
+    l1_array: &[L1Item<T>],
+    l1_index: usize,
+    row_index: i64,
+    bit_array: &FilteredBitArray,
+    op1: InequalityOperator,
+    left_row_ids: &mut PrimitiveChunkedBuilder<IdxType>,
+    right_row_ids: &mut PrimitiveChunkedBuilder<IdxType>,
+) -> i64
+where
+    T: NumericNative,
+    T: TotalOrd,
+{
+    debug_assert!(row_index > 0);
+    let mut match_count = 0;
+
+    // This entry comes from the left hand side DataFrame.
+    // Find all following entries in L1 (meaning they satisfy the first operator)
+    // that have already been visited (so satisfy the second operator).
+    // Because we use a stable sort for l2, we know that we won't find any
+    // matches for duplicate y values when traversing forwards in l1.
+    let start_index = find_search_start_index(l1_array, l1_index, op1);
+    bit_array.on_set_bits_from(start_index, |set_bit: usize| {
+        let right_row_index = l1_array[set_bit].row_index;
+        debug_assert!(right_row_index < 0);
+        left_row_ids.append_value((row_index - 1) as IdxSize);
+        right_row_ids.append_value((-right_row_index) as IdxSize - 1);
+        match_count += 1;
+    });
+
+    match_count
+}
+
 impl<T> L1Array for Vec<L1Item<T>>
 where
     T: NumericNative,
@@ -296,27 +331,22 @@ where
         left_row_ids: &mut PrimitiveChunkedBuilder<IdxType>,
         right_row_ids: &mut PrimitiveChunkedBuilder<IdxType>,
     ) -> i64 {
-        let mut match_count = 0;
         let row_index = self[l1_index].row_index;
         let from_lhs = row_index > 0;
         if from_lhs {
-            // This entry comes from the left hand side DataFrame.
-            // Find all following entries in L1 (meaning they satisfy the first operator)
-            // that have already been visited (so satisfy the second operator).
-            // Because we use a stable sort for l2, we know that we won't find any
-            // matches for duplicate y values when traversing forwards in l1.
-            let start_index = find_search_start_index(self, l1_index, op1);
-            bit_array.on_set_bits_from(start_index, |set_bit: usize| {
-                let right_row_index = self[set_bit].row_index;
-                debug_assert!(right_row_index < 0);
-                left_row_ids.append_value((row_index - 1) as IdxSize);
-                right_row_ids.append_value((-right_row_index) as IdxSize - 1);
-                match_count += 1;
-            });
+            find_matches_in_l1(
+                self,
+                l1_index,
+                row_index,
+                bit_array,
+                op1,
+                left_row_ids,
+                right_row_ids,
+            )
         } else {
             bit_array.set_bit(l1_index);
+            0
         }
-        match_count
     }
 
     fn process_lhs_entry(
@@ -327,30 +357,35 @@ where
         left_row_ids: &mut PrimitiveChunkedBuilder<IdxType>,
         right_row_ids: &mut PrimitiveChunkedBuilder<IdxType>,
     ) -> i64 {
-        let mut match_count = 0;
         let row_index = self[l1_index].row_index;
         let from_lhs = row_index > 0;
         if from_lhs {
-            let start_index = find_search_start_index(self, l1_index, op1);
-            bit_array.on_set_bits_from(start_index, |set_bit: usize| {
-                let right_row_index = self[set_bit].row_index;
-                debug_assert!(right_row_index < 0);
-                left_row_ids.append_value((row_index - 1) as IdxSize);
-                right_row_ids.append_value((-right_row_index) as IdxSize - 1);
-                match_count += 1;
-            });
+            find_matches_in_l1(
+                self,
+                l1_index,
+                row_index,
+                bit_array,
+                op1,
+                left_row_ids,
+                right_row_ids,
+            )
+        } else {
+            0
         }
-        match_count
     }
 
     fn mark_visited(&self, index: usize, bit_array: &mut FilteredBitArray) {
         let from_lhs = self[index].row_index > 0;
+        // We only mark RHS entries as visited,
+        // so that we don't try to match LHS entries with other LHS entries.
         if !from_lhs {
             bit_array.set_bit(index);
         }
     }
 }
 
+/// Create a vector of L1 items from the array of LHS x values concatenated with RHS x values
+/// and their ordering.
 fn build_l1_array<T>(
     ca: &ChunkedArray<T>,
     order: &IdxCa,
@@ -366,15 +401,20 @@ where
             // Nulls should have been skipped over
             .ok_or_else(|| polars_err!(ComputeError: "Unexpected null value in IEJoin data"))?;
         let row_index = if index < right_df_offset {
+            // Row from LHS
             index as i64 + 1
         } else {
+            // Row from RHS
             -((index - right_df_offset) as i64) - 1
         };
         array.push(L1Item { row_index, value });
     }
     Ok(Box::new(array))
 }
 
+/// Create a vector of L2 items from the array of y values ordered according to the L1 order,
+/// and their ordering. We don't need to store actual y values but only track whether we're at
+/// the end of a run of equal values.
 fn build_l2_array<T>(ca: &ChunkedArray<T>, order: &IdxCa) -> PolarsResult<Vec<L2Item>>
 where
     T: PolarsNumericType,

crates/polars-ops/src/frame/join/mod.rs

@@ -200,7 +200,7 @@ pub trait DataFrameJoinOps: IntoDf {
         }
 
         if let JoinType::IEJoin(options) = args.how {
-            return iejoin::join_dataframes(
+            return iejoin::iejoin(
                 left_df,
                 other,
                 selected_left,

crates/polars-python/src/lazyframe/general.rs

@@ -969,7 +969,7 @@ impl PyLazyFrame {
             .into())
     }
 
-    fn ie_join(&self, other: Self, on: Vec<PyExpr>, suffix: String) -> PyResult<Self> {
+    fn inequality_join(&self, other: Self, on: Vec<PyExpr>, suffix: String) -> PyResult<Self> {
         let ldf = self.ldf.clone();
         let other = other.ldf;
         let (left_on, operators, right_on) = parse_ie_join_expressions(on)?;

py-polars/polars/dataframe/frame.py

@@ -7084,41 +7084,39 @@ def join(
             .collect(_eager=True)
         )
 
-    def ie_join(
+    def inequality_join(
         self,
         other: DataFrame,
         *,
         on: Sequence[Expr],
         suffix: str = "_right",
     ) -> DataFrame:
         """
-        Perform a join using inequality operations.
+        Perform a join using two inequality expressions.
 
         Parameters
         ----------
         other
             DataFrame to join with.
         on
-            Inequality expressions to join with,
+            A sequence of two inequality expressions to join on, where each expression
+            is in the form `left_hand_side_expr op right_hand_side_expr` and op
+            is one of <, <=, >, >=.
             for example [pl.col("a") < pl.col("b"), pl.col("c") > pl.col("d")]
         suffix
             Suffix to append to columns with a duplicate name.
 
         Returns
         -------
         DataFrame
-
-        See Also
-        --------
-        join_asof
         """
         if not isinstance(other, DataFrame):
             msg = f"expected `other` join table to be a DataFrame, got {type(other).__name__!r}"
             raise TypeError(msg)
 
         return (
             self.lazy()
-            .ie_join(
+            .inequality_join(
                 other.lazy(),
                 on=on,
                 suffix=suffix,

py-polars/polars/lazyframe/frame.py

@@ -4341,33 +4341,31 @@ def join_asof(
             )
         )
 
-    def ie_join(
+    def inequality_join(
         self,
         other: LazyFrame,
         *,
         on: Sequence[Expr],
         suffix: str = "_right",
     ) -> LazyFrame:
         """
-        Perform a join using inequality operations.
+        Perform a join using two inequality expressions.
 
         Parameters
         ----------
         other
             LazyFrame to join with.
         on
-            Inequality expressions to join with,
+            A sequence of two inequality expressions to join on, where each expression
+            is in the form `left_hand_side_expr op right_hand_side_expr` and op
+            is one of <, <=, >, >=.
             for example [pl.col("a") < pl.col("b"), pl.col("c") > pl.col("d")]
         suffix
             Suffix to append to columns with a duplicate name.
 
         Returns
         -------
         LazyFrame
-
-        See Also
-        --------
-        join_asof
         """
         if not isinstance(other, LazyFrame):
             msg = f"expected `other` join table to be a LazyFrame, not a {type(other).__name__!r}"
@@ -4382,7 +4380,7 @@ def ie_join(
             raise ValueError(msg)
 
         return self._from_pyldf(
-            self._ldf.ie_join(
+            self._ldf.inequality_join(
                 other._ldf,
                 [expr._pyexpr for expr in on],
                 suffix,

py-polars/tests/unit/operations/test_ie_join.py py-polars/tests/unit/operations/test_inequality_join.py

@@ -23,7 +23,7 @@ def test_self_join() -> None:
         }
     )
 
-    actual = west.ie_join(
+    actual = west.inequality_join(
         west, on=[pl.col("time") > pl.col("time"), pl.col("cost") < pl.col("cost")]
     )
 
@@ -60,7 +60,7 @@ def test_basic_ie_join() -> None:
         }
     )
 
-    actual = east.ie_join(
+    actual = east.inequality_join(
         west, on=[pl.col("dur") < pl.col("time"), pl.col("rev") > pl.col("cost")]
     )
 
@@ -102,7 +102,7 @@ def test_ie_join_with_slice(offset: int, length: int) -> None:
     ).lazy()
 
     actual = (
-        east.ie_join(
+        east.inequality_join(
             west, on=[pl.col("dur") < pl.col("time"), pl.col("rev") < pl.col("cost")]
         )
         .slice(offset, length)
@@ -144,7 +144,7 @@ def test_ie_join_with_expressions() -> None:
         }
     )
 
-    actual = east.ie_join(
+    actual = east.inequality_join(
         west,
         on=[
             (pl.col("dur") * 2) < pl.col("time"),
@@ -272,7 +272,7 @@ def test_ie_join(east: pl.DataFrame, west: pl.DataFrame, op1: str, op2: str) ->
     expr0 = _inequality_expression("dur", op1, "time")
     expr1 = _inequality_expression("rev", op2, "cost")
 
-    actual = east.ie_join(west, on=[expr0, expr1])
+    actual = east.inequality_join(west, on=[expr0, expr1])
 
     expected = east.join(west, how="cross").filter(expr0 & expr1)
     assert_frame_equal(actual, expected, check_row_order=False, check_exact=True)
@@ -290,7 +290,7 @@ def test_ie_join_with_nulls(
     expr0 = _inequality_expression("dur", op1, "time")
     expr1 = _inequality_expression("rev", op2, "cost")
 
-    actual = east.ie_join(west, on=[expr0, expr1])
+    actual = east.inequality_join(west, on=[expr0, expr1])
 
     expected = east.join(west, how="cross").filter(expr0 & expr1)
     assert_frame_equal(actual, expected, check_row_order=False, check_exact=True)
@@ -308,7 +308,7 @@ def test_ie_join_with_floats(
     expr0 = _inequality_expression("dur", op1, "time")
     expr1 = _inequality_expression("rev", op2, "cost")
 
-    actual = east.ie_join(west, on=[expr0, expr1])
+    actual = east.inequality_join(west, on=[expr0, expr1])
 
     expected = east.join(west, how="cross").filter(expr0 & expr1)
     assert_frame_equal(actual, expected, check_row_order=False, check_exact=True)