no bitmap for inner

Eventual-Inc · Oct 4, 2024 · e429fdd · e429fdd
1 parent 2f53dbc
commit e429fdd
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 14 deletions.
diff --git a/daft/daft/__init__.pyi b/daft/daft/__init__.pyi
@@ -1745,7 +1745,7 @@ class NativeExecutor:
         logical_plan_builder: LogicalPlanBuilder,
     ) -> NativeExecutor: ...
     def run(
-        self, psets: dict[str, list[PyMicroPartition]], cfg: PyDaftExecutionConfig, results_buffer_size: int | None
+        self, psets: dict[str, list[PartitionT]], cfg: PyDaftExecutionConfig, results_buffer_size: int | None
     ) -> Iterator[PyMicroPartition]: ...
 
 class PyDaftExecutionConfig:

diff --git a/src/daft-local-execution/src/intermediate_ops/hash_join_probe.rs b/src/daft-local-execution/src/intermediate_ops/hash_join_probe.rs
@@ -2,10 +2,7 @@ use std::sync::Arc;
 
 use common_error::DaftResult;
 use daft_core::{
-    prelude::{
-        bitmap::{Bitmap, MutableBitmap},
-        BooleanArray, SchemaRef,
-    },
+    prelude::{SchemaRef, UInt64Array},
     series::IntoSeries,
 };
 use daft_dsl::ExprRef;
@@ -108,7 +105,7 @@ impl HashJoinProbeOperator {
         let mut build_side_growable =
             GrowableTable::new(&tables.iter().collect::<Vec<_>>(), false, 20)?;
 
-        let mut input_idx_matches = MutableBitmap::from_len_zeroed(input.len());
+        let mut input_idx_matches = vec![];
 
         drop(_growables);
         {
@@ -124,15 +121,14 @@ impl HashJoinProbeOperator {
                             build_row_idx as usize,
                             1,
                         );
-                        input_idx_matches.set(probe_row_idx, true);
+                        input_idx_matches.push(probe_row_idx as u64);
                     }
                 }
             }
         }
         let build_side_table = build_side_growable.build()?;
-        let bitmap: Bitmap = input_idx_matches.into();
         let probe_side_table =
-            input.mask_filter(&BooleanArray::from(("bitmap", bitmap)).into_series())?;
+            input.take(&UInt64Array::from(("matches", input_idx_matches)).into_series())?;
 
         let (left_table, right_table) = if self.build_on_left {
             (build_side_table, probe_side_table)
@@ -161,7 +157,7 @@ impl HashJoinProbeOperator {
             tables.iter().map(|t| t.len()).sum(),
         )?;
 
-        let mut input_idx_matches = MutableBitmap::from_len_zeroed(input.len());
+        let mut input_idx_matches = Vec::with_capacity(input.len());
 
         drop(_growables);
         {
@@ -177,19 +173,18 @@ impl HashJoinProbeOperator {
                             build_row_idx as usize,
                             1,
                         );
-                        input_idx_matches.set(probe_row_idx, true);
+                        input_idx_matches.push(probe_row_idx as u64);
                     }
                 } else {
                     // if there's no match, we should still emit the probe side and fill the build side with nulls
                     build_side_growable.add_nulls(1);
-                    input_idx_matches.set(probe_row_idx, true);
+                    input_idx_matches.push(probe_row_idx as u64);
                 }
             }
         }
         let build_side_table = build_side_growable.build()?;
-        let bitmap: Bitmap = input_idx_matches.into();
         let probe_side_table =
-            input.mask_filter(&BooleanArray::from(("bitmap", bitmap)).into_series())?;
+            input.take(&UInt64Array::from(("matches", input_idx_matches)).into_series())?;
 
         let final_table = if self.join_type == JoinType::Left {
             let join_table = probe_side_table.get_columns(&self.common_join_keys)?;