more work

Signed-off-by: Adam Gutglick <adamgsal@gmail.com>

Author: AdamGS

datafusion/physical-plan/src/joins/hash_join/exec.rs

@@ -197,6 +197,8 @@ pub(super) struct JoinLeftData {
     values: Vec<ArrayRef>,
     /// Shared bitmap builder for visited left indices
     visited_indices_bitmap: SharedBitmapBuilder,
+    /// Shared bitmap builder for null marks
+    null_indices_bitmap: SharedBitmapBuilder,
     /// Counter of running probe-threads, potentially
     /// able to update `visited_indices_bitmap`
     probe_threads_counter: AtomicUsize,
@@ -2024,16 +2026,18 @@ async fn collect_left_input(
         };
 
     // Reserve additional memory for visited indices bitmap and create shared builder
-    let visited_indices_bitmap = if with_visited_indices_bitmap {
+    let (visited_indices_bitmap, null_indices_bitmap) = if with_visited_indices_bitmap {
         let bitmap_size = bit_util::ceil(batch.num_rows(), 8);
-        reservation.try_grow(bitmap_size)?;
-        metrics.build_mem_used.add(bitmap_size);
-
-        let mut bitmap_buffer = BooleanBufferBuilder::new(batch.num_rows());
-        bitmap_buffer.append_n(num_rows, false);
-        bitmap_buffer
+        reservation.try_grow(bitmap_size * 2)?;
+        metrics.build_mem_used.add(bitmap_size * 2);
+
+        let mut visited = BooleanBufferBuilder::new(batch.num_rows());
+        let mut nulls = BooleanBufferBuilder::new(batch.num_rows());
+        visited.append_n(num_rows, false);
+        nulls.append_n(num_rows, false);
+        (visited, nulls)
     } else {
-        BooleanBufferBuilder::new(0)
+        (BooleanBufferBuilder::new(0), BooleanBufferBuilder::new(0))
     };
 
     let map = Arc::new(join_hash_map);
@@ -2073,6 +2077,7 @@ async fn collect_left_input(
         batch,
         values: left_values,
         visited_indices_bitmap: Mutex::new(visited_indices_bitmap),
+        null_indices_bitmap: Mutex::new(null_indices_bitmap),
         probe_threads_counter: AtomicUsize::new(probe_threads_count),
         _reservation: reservation,
         bounds,

datafusion/sqllogictest/test_files/null_aware_mark_join.slt

@@ -33,6 +33,20 @@ CREATE TABLE inner_table_with_null(id INT) AS VALUES
 (2),
 (NULL);
 
+statement ok
+CREATE TABLE outer_corr_table(id INT, grp INT, value TEXT) AS VALUES
+(1, 1, 'a'),
+(3, 1, 'b'),
+(1, 2, 'c'),
+(NULL, 1, 'd'),
+(5, 3, 'e');
+
+statement ok
+CREATE TABLE inner_corr_table(id INT, grp INT) AS VALUES
+(2, 1),
+(NULL, 1),
+(1, 2);
+
 statement ok
 CREATE TABLE empty_table(id INT) AS
 SELECT *
@@ -121,9 +135,22 @@ c
 d
 e
 
-###################################
-## Sort-merge join null-aware mark
-###################################
+# Regression test for correlated NOT IN inside a mark join.
+query T rowsort
+SELECT value
+FROM outer_corr_table
+WHERE (id NOT IN (
+  SELECT id
+  FROM inner_corr_table
+  WHERE inner_corr_table.grp = outer_corr_table.grp
+)) IS NULL;
+----
+a
+b
+d
+
+# These queries would fall back to a sort-merge, but we force them
+# into hash joins to make sure we null awareness.
 
 statement ok
 set datafusion.execution.target_partitions = 2;
@@ -257,5 +284,11 @@ DROP TABLE inner_table_with_null;
 statement ok
 DROP TABLE inner_table_no_null;
 
+statement ok
+DROP TABLE inner_corr_table;
+
+statement ok
+DROP TABLE outer_corr_table;
+
 statement ok
 DROP TABLE outer_table;