fix: handle dictionary nulls in null-aware left mark joins

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

@@ -355,6 +355,11 @@ struct JoinBuildData {
     /// `build_side_has_null` were computed from the original keys in
     /// `collect_left_input`.
     values: Vec<ArrayRef>,
+    /// Logical nulls of the scalar key for uncorrelated null-aware `LeftMark`.
+    ///
+    /// Computed once while collecting the build side because dictionary
+    /// logical-null discovery can scan nested dictionary values.
+    left_mark_build_key_nulls: Option<NullBuffer>,
     /// Bounds computed from the build side; absent for an empty partition.
     bounds: Option<PartitionBounds>,
     /// IN-list values or a hash-table reference used for filter pushdown.
@@ -431,6 +436,10 @@ impl JoinLeftData {
         &self.build.values
     }
 
+    pub(super) fn left_mark_build_key_nulls(&self) -> Option<&NullBuffer> {
+        self.build.left_mark_build_key_nulls.as_ref()
+    }
+
     /// returns a reference to the visited indices bitmap
     pub(super) fn visited_indices_bitmap(&self) -> &SharedBitmapBuilder {
         &self.visited_indices_bitmap
@@ -3543,6 +3552,33 @@ async fn collect_left_input(
         && !left_values.is_empty()
         && left_values[0].logical_null_count() > 0;
 
+    let left_mark_build_key_nulls = if matches!(
+        null_aware,
+        Some(NullAwareMode::LeftMark {
+            correlated: false,
+            ..
+        })
+    ) {
+        left_values[0]
+            .logical_nulls()
+            .filter(|nulls| nulls.null_count() > 0)
+    } else {
+        None
+    };
+    if let Some(nulls) = &left_mark_build_key_nulls {
+        // `logical_nulls` reuses an array's physical null buffer when it can.
+        // Only charge newly materialized buffers, such as dictionary logical
+        // nulls, because the build batch's buffers were reserved above.
+        let reuses_input_buffer = left_values[0]
+            .nulls()
+            .is_some_and(|input| input.inner().ptr_eq(nulls.inner()));
+        if !reuses_input_buffer {
+            let retained_size = nulls.buffer().capacity();
+            reservation.try_grow(retained_size)?;
+            metrics.build_mem_used.add(retained_size);
+        }
+    }
+
     if prepared {
         drop(batches);
         // Prepared keys are direct columns. IN-list arrays share these batch
@@ -3561,6 +3597,7 @@ async fn collect_left_input(
             map,
             batch,
             values: normalized_values,
+            left_mark_build_key_nulls,
             bounds,
             membership,
             reservation,
@@ -9396,6 +9433,59 @@ mod tests {
         Ok(())
     }
 
+    /// The dictionary counterpart of
+    /// [`test_null_aware_left_mark_probe_null_in_other_partition`], where the
+    /// probe NULL exists in a dictionary value instead of the key bitmap.
+    #[apply(hash_join_exec_configs)]
+    #[tokio::test]
+    async fn test_null_aware_left_mark_probe_logical_null_in_other_partition(
+        batch_size: usize,
+    ) -> Result<()> {
+        let task_ctx = prepare_task_ctx(batch_size, false);
+
+        for order in [[0, 1], [1, 0]] {
+            let left = build_table_dict_key(
+                "c1",
+                vec![Some(1), Some(4)],
+                vec![0, 1],
+                "dummy",
+                vec![Some(10), Some(40)],
+            );
+            let right = build_two_partition_dict_probe_with_logical_null_in_partition_0();
+
+            let on = vec![(
+                Arc::new(Column::new_with_schema("c1", &left.schema())?) as _,
+                Arc::new(Column::new_with_schema("c2", &right.schema())?) as _,
+            )];
+
+            let join = HashJoinExec::try_new(
+                left,
+                right,
+                on,
+                None,
+                &JoinType::LeftMark,
+                None,
+                PartitionMode::CollectLeft,
+                NullEquality::NullEqualsNothing,
+                true,
+            )?;
+
+            let batches = collect_partitions_in_order(&join, &order, &task_ctx).await?;
+
+            allow_duplicates! {
+                assert_snapshot!(batches_to_sort_string(&batches), @r"
+                +----+-------+------+
+                | c1 | dummy | mark |
+                +----+-------+------+
+                | 1  | 10    | true |
+                | 4  | 40    |      |
+                +----+-------+------+
+                ");
+            }
+        }
+        Ok(())
+    }
+
     /// Test null-aware anti join when build side (left) contains NULL keys
     /// Expected: rows with NULL keys should not be output
     #[apply(hash_join_exec_configs)]
@@ -10316,6 +10406,145 @@ mod tests {
         Ok(())
     }
 
+    /// A dictionary key that points to a NULL dictionary value is a NULL build
+    /// key even when the dictionary key bitmap itself has no NULLs.
+    #[apply(hash_join_exec_configs)]
+    #[tokio::test]
+    async fn test_null_aware_left_mark_build_logical_null(
+        batch_size: usize,
+    ) -> Result<()> {
+        let task_ctx = prepare_task_ctx(batch_size, false);
+
+        let left = build_table_dict_key(
+            "c1",
+            vec![Some(1), None, Some(4)],
+            vec![0, 1, 2],
+            "dummy",
+            vec![Some(10), Some(0), Some(40)],
+        );
+        let left_key = left
+            .execute(0, Arc::clone(&task_ctx))?
+            .next()
+            .await
+            .unwrap()?;
+        assert_eq!(left_key.column(0).null_count(), 0);
+        assert_eq!(left_key.column(0).logical_null_count(), 1);
+
+        let right = build_table_dict_key(
+            "c2",
+            vec![Some(1), Some(2), Some(3)],
+            vec![0, 1, 2],
+            "dummy",
+            vec![Some(100), Some(200), Some(300)],
+        );
+
+        let on = vec![(
+            Arc::new(Column::new_with_schema("c1", &left.schema())?) as _,
+            Arc::new(Column::new_with_schema("c2", &right.schema())?) as _,
+        )];
+
+        let join = HashJoinExec::try_new(
+            left,
+            right,
+            on,
+            None,
+            &JoinType::LeftMark,
+            None,
+            PartitionMode::CollectLeft,
+            NullEquality::NullEqualsNothing,
+            true,
+        )?;
+
+        let stream = join.execute(0, task_ctx)?;
+        let batches = common::collect(stream).await?;
+
+        allow_duplicates! {
+            assert_snapshot!(batches_to_sort_string(&batches), @r"
+            +----+-------+-------+
+            | c1 | dummy | mark  |
+            +----+-------+-------+
+            |    | 0     |       |
+            | 1  | 10    | true  |
+            | 4  | 40    | false |
+            +----+-------+-------+
+            ");
+        }
+        Ok(())
+    }
+
+    fn build_single_column_table(name: &str, array: ArrayRef) -> Arc<dyn ExecutionPlan> {
+        let schema = Arc::new(Schema::new(vec![Field::new(
+            name,
+            array.data_type().clone(),
+            true,
+        )]));
+        let batch = RecordBatch::try_new(Arc::clone(&schema), vec![array]).unwrap();
+        TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap()
+    }
+
+    fn nested_dictionary_array(
+        leaf_values: Vec<Option<i32>>,
+        inner_keys: Vec<i32>,
+        outer_keys: Vec<i32>,
+    ) -> ArrayRef {
+        let leaf_values: ArrayRef = Arc::new(Int32Array::from(leaf_values));
+        let inner_keys = Int32Array::from(inner_keys);
+        let inner_values: ArrayRef =
+            Arc::new(DictionaryArray::<Int32Type>::new(inner_keys, leaf_values));
+        Arc::new(DictionaryArray::<Int32Type>::new(
+            Int32Array::from(outer_keys),
+            inner_values,
+        ))
+    }
+
+    /// A nested dictionary's logical null bitmap must be computed once for the
+    /// complete build key and remain indexed by the original build row.
+    #[apply(hash_join_exec_configs)]
+    #[tokio::test]
+    async fn test_null_aware_left_mark_nested_dictionary_logical_null(
+        batch_size: usize,
+    ) -> Result<()> {
+        let task_ctx = prepare_task_ctx(batch_size, false);
+        let left_key = nested_dictionary_array(
+            vec![Some(1), None, Some(4)],
+            vec![0, 1, 2],
+            vec![0, 1, 2],
+        );
+        assert_eq!(left_key.logical_null_count(), 1);
+        let left = build_single_column_table("c1", left_key);
+        let right = build_single_column_table(
+            "c2",
+            nested_dictionary_array(
+                vec![Some(1), Some(2), Some(3)],
+                vec![0, 1, 2],
+                vec![0, 1, 2],
+            ),
+        );
+        let on = vec![(
+            Arc::new(Column::new_with_schema("c1", &left.schema())?) as _,
+            Arc::new(Column::new_with_schema("c2", &right.schema())?) as _,
+        )];
+        let join = HashJoinExec::try_new(
+            left,
+            right,
+            on,
+            None,
+            &JoinType::LeftMark,
+            None,
+            PartitionMode::CollectLeft,
+            NullEquality::NullEqualsNothing,
+            true,
+        )?;
+
+        let batches = common::collect(join.execute(0, task_ctx)?).await?;
+        let marks = batches
+            .iter()
+            .flat_map(|batch| batch.column(batch.num_columns() - 1).as_boolean().iter())
+            .collect::<Vec<_>>();
+        assert_eq!(marks, vec![Some(true), None, Some(false)]);
+        Ok(())
+    }
+
     /// Scalar correlated null-aware left mark join where only the build side's
     /// value key contains NULL.
     ///

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

@@ -1289,18 +1289,19 @@ fn null_aware_skip_probe_batch(
         | NullAwareMode::LeftMark {
             correlated: false, ..
         } => {
-            // `on[0]` is the `NOT IN` value key for both modes.
-            let probe_key_column = &state.values[0];
             let is_anti = matches!(mode, NullAwareMode::LeftAnti { .. });
-            let probe_has_null = if is_anti {
-                probe_key_column.logical_null_count() > 0
-            } else {
-                probe_key_column.null_count() > 0
-            };
-            // Only batches with rows count: `NULL NOT IN (empty)` is TRUE.
-            left_data.record_probe_batch(state.batch.num_rows() > 0, probe_has_null);
+            if state.offset == (0, None) {
+                // `on[0]` is the `NOT IN` value key for both modes. A probe
+                // batch may resume lookup many times, but its summary is
+                // invariant and only needs to be computed once.
+                let probe_has_null = state.values[0].logical_null_count() > 0;
+                // Only batches with rows count: `NULL NOT IN (empty)` is TRUE.
+                left_data.record_probe_batch(state.batch.num_rows() > 0, probe_has_null);
+            }
             // Best-effort early exit; the final stage re-checks the flag
-            // through `report_probe_completed`.
+            // through `report_probe_completed`. Keep this outside the first-
+            // chunk guard so a sibling partition can stop an anti join between
+            // resumed lookups.
             is_anti && left_data.probe_side_has_null_hint()
         }
     }
@@ -1388,13 +1389,19 @@ fn null_aware_left_mark_column(
     left_side: &UInt64Array,
     right_side: &UInt32Array,
 ) -> ArrayRef {
-    let build_key_column = &left_data.values()[0];
     // Correlated joins precomputed the UNKNOWN decision per build row.
     let null_indices_bitmap = correlated.then(|| left_data.null_indices_bitmap().lock());
+    // Uncorrelated joins use the logical-null bitmap computed once with the
+    // build side. Correlated joins carry the answer in `null_indices_bitmap`.
+    let build_key_nulls = if correlated {
+        None
+    } else {
+        left_data.left_mark_build_key_nulls()
+    };
     build_null_aware_left_mark_column(
         left_side,
         right_side,
-        build_key_column.as_ref(),
+        build_key_nulls,
         null_indices_bitmap.as_deref(),
         probe_summary.has_null,
         probe_summary.non_empty,

datafusion/physical-plan/src/joins/utils.rs

@@ -1382,12 +1382,14 @@ pub(crate) fn build_batch_from_indices(
 /// For correlated scalar `NOT IN`, `null_indices_bitmap` carries the same UNKNOWN
 /// decision per build row, scoped by the correlated equality keys.
 ///
+/// When present, `build_key_nulls` is indexed by the original build row number.
+///
 /// This is the helper equivalent of the paper's "null bucket" and `hadNull` handling.
 /// It is intentionally scoped to scalar null-aware mark joins.
 pub(crate) fn build_null_aware_left_mark_column(
     build_indices: &UInt64Array,
     probe_indices: &UInt32Array,
-    build_key_column: &dyn Array,
+    build_key_nulls: Option<&NullBuffer>,
     null_indices_bitmap: Option<&BooleanBufferBuilder>,
     probe_side_has_null: bool,
     probe_side_non_empty: bool,
@@ -1397,7 +1399,9 @@ pub(crate) fn build_null_aware_left_mark_column(
     // uncorrelated rules are cases 1-4 in the doc above.
     let unmatched_mark_is_null = |build_idx: usize| match null_indices_bitmap {
         Some(bitmap) => bitmap.get_bit(build_idx),
-        None if build_key_column.is_null(build_idx) => probe_side_non_empty,
+        None if build_key_nulls.is_some_and(|nulls| nulls.is_null(build_idx)) => {
+            probe_side_non_empty
+        }
         None => probe_side_has_null,
     };