jayzhan211 commented on code in PR #25339:
URL: https://github.com/apache/datafusion/pull/25339#discussion_r4027554354


##########
datafusion/physical-plan/src/joins/hash_join/exec.rs:
##########
@@ -239,24 +244,27 @@ impl NullAwareMode {
         num_keys: usize,
         has_filter: bool,
     ) -> Result<Self> {
+        let correlated = num_keys > 1 || has_filter;

Review Comment:
   `correlated = num_keys > 1 || has_filter` assumes `on[0]` is the NOT IN 
value key. When the value has no outer columns, `1 = i.id` is pushed into the 
subquery, so `on[0]` becomes the correlation key `o.g = i.g`, and a NULL `o.g` 
is marked UNKNOWN:
   
   ```sql
   CREATE TABLE o(id INT, g INT, z INT) AS VALUES (1,1,10),(2,NULL,10),(3,2,10);
   CREATE TABLE i(id INT, g INT, z INT) AS VALUES (1,1,5),(5,2,5),(NULL,3,5);
   SELECT id FROM o WHERE 1 NOT IN (SELECT i.id FROM i WHERE i.g = o.g AND i.z 
< o.z);
   -- expected 2, 3; returns 3
   ```
   
   The form without `AND i.z < o.z` is also wrong and doesn't go through the 
new code, so this predates the PR. Fine as a follow-up: in `build_join`, only 
set `null_aware` when the in-predicate's outer side references a left column.



##########
datafusion/physical-plan/src/joins/hash_join/stream.rs:
##########
@@ -1344,117 +1360,196 @@ fn null_aware_left_mark_column(
     )
 }
 
-/// Records which build rows of a correlated null-aware `LeftMark` join are
-/// UNKNOWN candidates for this probe batch.
+/// Records which build rows of a correlated null-aware join are UNKNOWN
+/// candidates for this probe batch.
 ///
-/// Key layout: `on[0]` is the `NOT IN` value key, `on[1..]` the correlation
-/// scope keys (see `HashJoinExec::null_aware`). A build row's mark must be
-/// NULL (SQL UNKNOWN) instead of FALSE when it is unmatched and either:
-/// 1. its value key is NULL and any probe row shares its correlation scope, or
-/// 2. some probe row in its correlation scope has a NULL value key.
+/// Key layout: `on[0]` is the `NOT IN` value key, `on[1..]` the (possibly
+/// empty) correlation scope keys (see `HashJoinExec::null_aware`). An
+/// unmatched build row's `NOT IN` is UNKNOWN instead of TRUE (its mark is NULL
+/// instead of FALSE) when either:
+/// 1. its value key is NULL and any probe row in its correlation scope passes
+///    the join filter, or
+/// 2. some probe row in its correlation scope with a NULL value key passes the
+///    join filter.
 ///
-/// Case 1 probes the build-side NULL-value scope map with all probe rows;
-/// case 2 probes the full scope map with only the NULL-valued probe rows.
+/// Case 1 pairs the NULL-valued build rows with all probe rows; case 2 pairs
+/// all build rows with the NULL-valued probe rows. Scope keys narrow these
+/// pairs through a hash lookup; without scope keys every pair is a candidate.
+/// The join filter, if any, then decides which candidates count.
+#[expect(clippy::too_many_arguments)]
 fn mark_null_candidates_for_probe_batch(
     build_side: &BuildSideReadyState,
     state: &ProcessProbeBatchState,
+    filter: Option<&JoinFilter>,
+    join_type: JoinType,
     random_state: &RandomState,
     batch_size: usize,
     hashes_buffer: &mut Vec<u64>,
     probe_indices_buffer: &mut Vec<u32>,
     build_indices_buffer: &mut Vec<u64>,
 ) -> Result<()> {
-    let Some(scope_map) = build_side.left_data.null_aware_mark_scope_map() 
else {
+    let left_data = &build_side.left_data;
+    let null_value_build_rows = left_data.null_value_build_rows();
+    let probe_value_key = &state.values[0];
+    let probe_has_null_values = probe_value_key.logical_null_count() > 0;
+    if null_value_build_rows.is_none() && !probe_has_null_values {
         return Ok(());
-    };
+    }
 
     debug_assert_eq!(
-        build_side.left_data.values().len(),
+        left_data.values().len(),
         state.values.len(),
         "build/probe key counts must match"
     );
-    debug_assert!(state.values.len() > 1, "keys must be [value, scope..]");
-
-    let probe_value_key = &state.values[0];
-    let build_scope_values = &build_side.left_data.values()[1..];
+    let build_scope_values = &left_data.values()[1..];
     let probe_scope_values = &state.values[1..];
 
-    let null_value_scope_map = build_side.left_data.null_value_scope_map();
-    let probe_has_null_values = probe_value_key.null_count() > 0;
-    if null_value_scope_map.is_none() && !probe_has_null_values {
-        return Ok(());
-    }
+    // Keeps the candidate pairs that pass the join filter and marks their
+    // build rows as UNKNOWN.
+    let mut mark = |build_indices: UInt64Array, probe_indices: UInt32Array| {
+        let build_indices = match filter {
+            Some(filter) => {
+                apply_join_filter_to_indices(
+                    left_data.batch(),
+                    &state.batch,
+                    build_indices,
+                    probe_indices,
+                    filter,
+                    JoinSide::Left,
+                    None,
+                    join_type,
+                )?
+                .0
+            }
+            None => build_indices,
+        };
+        if !build_indices.is_empty() {
+            let mut null_bitmap = left_data.null_indices_bitmap().lock();
+            for build_idx in build_indices.values() {
+                null_bitmap.set_bit(*build_idx as usize, true);
+            }
+        }
+        Ok(())
+    };
 
     // Case 1: build rows with a NULL value key are UNKNOWN as soon as any
-    // probe row shares their correlation scope.
-    if let Some(null_value_scope_map) = null_value_scope_map {
-        hashes_buffer.clear();
-        hashes_buffer.resize(state.batch.num_rows(), 0);
-        create_hashes(probe_scope_values, random_state, hashes_buffer)?;
-
-        scan_scope_matches_into_bitmap(
-            null_value_scope_map.map.as_ref(),
-            &null_value_scope_map.scope_values,
-            probe_scope_values,
-            hashes_buffer,
-            batch_size,
-            probe_indices_buffer,
-            build_indices_buffer,
-            // The map indexes only the NULL-valued build rows; translate its
-            // positions back to row indices in the full build batch.
-            |position| null_value_scope_map.build_indices.value(position as 
usize),
-            build_side.left_data.null_indices_bitmap(),
-        )?;
+    // probe row in their correlation scope passes the filter.
+    if let Some(null_rows) = null_value_build_rows {
+        match &null_rows.scope_map {
+            Some(scope_map) => {
+                hashes_buffer.clear();
+                hashes_buffer.resize(state.batch.num_rows(), 0);
+                create_hashes(probe_scope_values, random_state, 
hashes_buffer)?;
+
+                for_each_scope_match(
+                    scope_map.as_ref(),
+                    &null_rows.scope_values,
+                    probe_scope_values,
+                    hashes_buffer,
+                    batch_size,
+                    probe_indices_buffer,
+                    build_indices_buffer,
+                    |positions, probe_indices| {
+                        // The map indexes only the NULL-valued build rows;
+                        // translate its positions back to build row indices.
+                        let build_indices = UInt64Array::from_iter_values(
+                            positions
+                                .values()
+                                .iter()
+                                .map(|p| null_rows.build_indices.value(*p as 
usize)),
+                        );
+                        mark(build_indices, probe_indices)
+                    },
+                )?;
+            }
+            None => {
+                let probe_rows =
+                    UInt32Array::from_iter_values(0..state.batch.num_rows() as 
u32);
+                for_each_cross_product(
+                    &null_rows.build_indices,
+                    &probe_rows,
+                    batch_size,
+                    &mut mark,
+                )?;
+            }
+        }
     }
 
     // Case 2: NULL-valued probe rows make every build row in their correlation
-    // scope an UNKNOWN candidate.
+    // scope that passes the filter an UNKNOWN candidate.
     if probe_has_null_values {
         let null_mask = arrow::compute::is_null(probe_value_key.as_ref())?;
-        let probe_null_scope_values = probe_scope_values
-            .iter()
-            .map(|values| Ok(arrow::compute::filter(values.as_ref(), 
&null_mask)?))
-            .collect::<Result<Vec<_>>>()?;
-
-        hashes_buffer.clear();
-        hashes_buffer.resize(null_mask.true_count(), 0);
-        create_hashes(&probe_null_scope_values, random_state, hashes_buffer)?;
+        let null_probe_rows = UInt32Array::from_iter_values(
+            null_mask.values().set_indices().map(|i| i as u32),
+        );
 
-        scan_scope_matches_into_bitmap(
-            scope_map,
-            build_scope_values,
-            &probe_null_scope_values,
-            hashes_buffer,
-            batch_size,
-            probe_indices_buffer,
-            build_indices_buffer,
-            |position| position,
-            build_side.left_data.null_indices_bitmap(),
-        )?;
+        match left_data.null_aware_scope_map() {
+            Some(scope_map) => {
+                let probe_null_scope_values = probe_scope_values
+                    .iter()
+                    .map(|values| {
+                        Ok(arrow::compute::filter(values.as_ref(), 
&null_mask)?)
+                    })
+                    .collect::<Result<Vec<_>>>()?;
+
+                hashes_buffer.clear();
+                hashes_buffer.resize(null_probe_rows.len(), 0);
+                create_hashes(&probe_null_scope_values, random_state, 
hashes_buffer)?;
+
+                for_each_scope_match(
+                    scope_map,
+                    build_scope_values,
+                    &probe_null_scope_values,
+                    hashes_buffer,
+                    batch_size,
+                    probe_indices_buffer,
+                    build_indices_buffer,
+                    |build_indices, positions| {
+                        // The lookup ran over only the NULL-valued probe rows;
+                        // translate its positions back to probe row indices.
+                        let probe_indices = UInt32Array::from_iter_values(
+                            positions
+                                .values()
+                                .iter()
+                                .map(|p| null_probe_rows.value(*p as usize)),
+                        );
+                        mark(build_indices, probe_indices)
+                    },
+                )?;
+            }
+            None => {
+                let build_rows =
+                    
UInt64Array::from_iter_values(0..left_data.batch().num_rows() as u64);
+                for_each_cross_product(

Review Comment:
   **The case with no scope keys re-checks build rows that are already marked 
UNKNOWN**
   
   Without scope keys, case 2 evaluates the filter for every (build row × NULL 
probe row) pair in every probe batch, including build rows already set in 
`null_indices_bitmap`. Those bits never clear, so that work is wasted. 20K 
outer × 10K NULL inner with `i.z < o.z` spends 8.85s in `join_time` (debug 
build). Skip marked rows and stop once none are left (same for case 1 at 
`:1466`):
   
   ```rs
   None => {
       let num_build_rows = left_data.batch().num_rows();
       for probe_rows in null_probe_rows.values().chunks(batch_size.max(1)) {
           let build_rows = {
               let bitmap = left_data.null_indices_bitmap().lock();
               UInt64Array::from_iter_values(
                   (0..num_build_rows)
                       .filter(|i| !bitmap.get_bit(*i))
                       .map(|i| i as u64),
               )
           };
           if build_rows.is_empty() {
               break;
           }
           let probe_rows = UInt32Array::from(probe_rows.to_vec());
           for_each_cross_product(&build_rows, &probe_rows, batch_size, &mut 
mark)?;
       }
   }
   ```
   
   Fine to handle in a follow-up



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to