refactor(rust): Add equi joins to new streaming engine

crates/polars-arrow/src/compute/take/generic_binary.rs

@@ -73,7 +73,7 @@ pub(super) unsafe fn take_values_validity<O: Offset, I: Index, A: GenericBinaryA
         .map(|index| validity_values.get_bit_unchecked(index.to_usize()));
     let validity = Bitmap::from_trusted_len_iter(validity);
 
-    let mut length = O::default();
+    let mut total_length = O::default();
 
     let offsets = values.offsets();
     let values_values = values.values();
@@ -82,12 +82,13 @@ pub(super) unsafe fn take_values_validity<O: Offset, I: Index, A: GenericBinaryA
     let lengths = indices.iter().map(|index| {
         let index = index.to_usize();
         let start = *offsets.get_unchecked(index);
-        length += *offsets.get_unchecked(index + 1) - start;
+        let length = *offsets.get_unchecked(index + 1) - start;
+        total_length += length;
         starts.push_unchecked(start);
         length.to_usize()
     });
     let offsets = create_offsets(lengths, indices.len());
-    let buffer = take_values(length, starts.as_slice(), &offsets, values_values);
+    let buffer = take_values(total_length, starts.as_slice(), &offsets, values_values);
 
     (offsets, buffer, validity.into())
 }
@@ -98,26 +99,31 @@ pub(super) unsafe fn take_indices_validity<O: Offset, I: Index>(
     values: &[u8],
     indices: &PrimitiveArray<I>,
 ) -> (OffsetsBuffer<O>, Buffer<u8>, Option<Bitmap>) {
-    let mut length = O::default();
+    let mut total_length = O::default();
 
     let offsets = offsets.buffer();
 
     let mut starts = Vec::<O>::with_capacity(indices.len());
     let lengths = indices.values().iter().map(|index| {
         let index = index.to_usize();
+        let length;
         match offsets.get(index + 1) {
             Some(&next) => {
                 let start = *offsets.get_unchecked(index);
-                length += next - start;
+                length = next - start;
+                total_length += length;
                 starts.push_unchecked(start);
             },
-            None => starts.push_unchecked(O::default()),
+            None => {
+                length = O::zero();
+                starts.push_unchecked(O::default());
+            },
         };
         length.to_usize()
     });
     let offsets = create_offsets(lengths, indices.len());
 
-    let buffer = take_values(length, &starts, &offsets, values);
+    let buffer = take_values(total_length, &starts, &offsets, values);
 
     (offsets, buffer, indices.validity().cloned())
 }
@@ -127,7 +133,7 @@ pub(super) unsafe fn take_values_indices_validity<O: Offset, I: Index, A: Generi
     values: &A,
     indices: &PrimitiveArray<I>,
 ) -> (OffsetsBuffer<O>, Buffer<u8>, Option<Bitmap>) {
-    let mut length = O::default();
+    let mut total_length = O::default();
     let mut validity = MutableBitmap::with_capacity(indices.len());
 
     let values_validity = values.validity().unwrap();
@@ -136,28 +142,32 @@ pub(super) unsafe fn take_values_indices_validity<O: Offset, I: Index, A: Generi
 
     let mut starts = Vec::<O>::with_capacity(indices.len());
     let lengths = indices.iter().map(|index| {
+        let length;
         match index {
             Some(index) => {
                 let index = index.to_usize();
                 if values_validity.get_bit(index) {
                     validity.push(true);
-                    length += *offsets.get_unchecked(index + 1) - *offsets.get_unchecked(index);
+                    length = *offsets.get_unchecked(index + 1) - *offsets.get_unchecked(index);
                     starts.push_unchecked(*offsets.get_unchecked(index));
                 } else {
                     validity.push(false);
+                    length = O::zero();
                     starts.push_unchecked(O::default());
                 }
             },
             None => {
                 validity.push(false);
+                length = O::zero();
                 starts.push_unchecked(O::default());
             },
         };
+        total_length += length;
         length.to_usize()
     });
     let offsets = create_offsets(lengths, indices.len());
 
-    let buffer = take_values(length, &starts, &offsets, values_values);
+    let buffer = take_values(total_length, &starts, &offsets, values_values);
 
     (offsets, buffer, validity.into())
 }

crates/polars-core/src/datatypes/field.rs

@@ -97,6 +97,12 @@ impl Field {
         self.name = name;
     }
 
+    /// Returns this `Field`, renamed.
+    pub fn with_name(mut self, name: PlSmallStr) -> Self {
+        self.name = name;
+        self
+    }
+
     /// Converts the `Field` to an `arrow::datatypes::Field`.
     ///
     /// # Example

crates/polars-core/src/frame/mod.rs

@@ -380,6 +380,15 @@ impl DataFrame {
         unsafe { DataFrame::new_no_checks(0, cols) }
     }
 
+    /// Create a new `DataFrame` with the given schema, only containing nulls.
+    pub fn full_null(schema: &Schema, height: usize) -> Self {
+        let columns = schema
+            .iter_fields()
+            .map(|f| Column::full_null(f.name.clone(), height, f.dtype()))
+            .collect();
+        DataFrame { height, columns }
+    }
+
     /// Removes the last `Series` from the `DataFrame` and returns it, or [`None`] if it is empty.
     ///
     /// # Example
@@ -713,7 +722,7 @@ impl DataFrame {
     /// - The length of each appended column matches the height of the [`DataFrame`]. For
     ///   `DataFrame`]s with no columns (ZCDFs), it is important that the height is set afterwards
     ///   with [`DataFrame::set_height`].
-    pub unsafe fn column_extend_unchecked(&mut self, iter: impl Iterator<Item = Column>) {
+    pub unsafe fn column_extend_unchecked(&mut self, iter: impl IntoIterator<Item = Column>) {
         unsafe { self.get_columns_mut() }.extend(iter)
     }
 
@@ -1894,11 +1903,15 @@ impl DataFrame {
         unsafe { DataFrame::new_no_checks(idx.len(), cols) }
     }
 
-    pub(crate) unsafe fn take_slice_unchecked(&self, idx: &[IdxSize]) -> Self {
+    /// # Safety
+    /// The indices must be in-bounds.
+    pub unsafe fn take_slice_unchecked(&self, idx: &[IdxSize]) -> Self {
         self.take_slice_unchecked_impl(idx, true)
     }
 
-    unsafe fn take_slice_unchecked_impl(&self, idx: &[IdxSize], allow_threads: bool) -> Self {
+    /// # Safety
+    /// The indices must be in-bounds.
+    pub unsafe fn take_slice_unchecked_impl(&self, idx: &[IdxSize], allow_threads: bool) -> Self {
         let cols = if allow_threads {
             POOL.install(|| self._apply_columns_par(&|s| s.take_slice_unchecked(idx)))
         } else {

crates/polars-expr/src/chunked_idx_table/mod.rs

@@ -0,0 +1,65 @@
+use std::any::Any;
+
+use polars_core::prelude::*;
+use polars_utils::index::ChunkId;
+use polars_utils::IdxSize;
+
+use crate::hash_keys::HashKeys;
+
+mod row_encoded;
+
+pub trait ChunkedIdxTable: Any + Send + Sync {
+    /// Creates a new empty ChunkedIdxTable similar to this one.
+    fn new_empty(&self) -> Box<dyn ChunkedIdxTable>;
+
+    /// Reserves space for the given number additional keys.
+    fn reserve(&mut self, additional: usize);
+
+    /// Returns the number of unique keys in this ChunkedIdxTable.
+    fn num_keys(&self) -> IdxSize;
+
+    /// Inserts the given key chunk into this ChunkedIdxTable.
+    fn insert_key_chunk(&mut self, keys: HashKeys, track_unmatchable: bool);
+
+    /// Probe the table, updating table_match and probe_match with
+    /// (ChunkId, IdxSize) pairs for each match. Will stop processing new keys
+    /// once limit matches have been generated, returning the number of keys
+    /// processed.
+    ///
+    /// If mark_matches is true, matches are marked in the table as such.
+    ///
+    /// If emit_unmatched is true, for keys that do not have a match we emit a
+    /// match with ChunkId::null() on the table match.
+    fn probe(
+        &self,
+        hash_keys: &HashKeys,
+        table_match: &mut Vec<ChunkId<32>>,
+        probe_match: &mut Vec<IdxSize>,
+        mark_matches: bool,
+        emit_unmatched: bool,
+        limit: IdxSize,
+    ) -> IdxSize;
+
+    /// The same as probe, except it will only apply to the specified subset of keys.
+    /// # Safety
+    /// The provided subset indices must be in-bounds.
+    #[allow(clippy::too_many_arguments)]
+    unsafe fn probe_subset(
+        &self,
+        hash_keys: &HashKeys,
+        subset: &[IdxSize],
+        table_match: &mut Vec<ChunkId<32>>,
+        probe_match: &mut Vec<IdxSize>,
+        mark_matches: bool,
+        emit_unmatched: bool,
+        limit: IdxSize,
+    ) -> IdxSize;
+
+    /// Get the ChunkIds for each key which was never marked during probing.
+    fn unmarked_keys(&self, out: &mut Vec<ChunkId<32>>, offset: IdxSize, limit: IdxSize)
+        -> IdxSize;
+}
+
+pub fn new_chunked_idx_table(_key_schema: Arc<Schema>) -> Box<dyn ChunkedIdxTable> {
+    Box::new(row_encoded::RowEncodedChunkedIdxTable::new())
+}