sunchao commented on code in PR #26214:
URL: https://github.com/apache/datafusion/pull/26214#discussion_r4436475959


##########
docs/source/library-user-guide/upgrading/56.0.0.md:
##########
@@ -934,9 +934,15 @@ Wire compatibility is directional:
 ### Null-aware joins record their `NOT IN` value key count
 
 Multi-column `(a, b) NOT IN (SELECT x, y ...)` is now planned as a null-aware
-join. Its first equi-join keys are the tuple elements, followed by any
-correlation keys, so a null-aware join now records how many of its leading keys
-are `NOT IN` value keys:
+join. Its first equi-join keys are the tuple elements that can be NULL. A
+nullable element cannot be NULL here when a correlated filter of the subquery
+rejects NULL for it, such as `y > o.c` for `y`. The elements that cannot be
+NULL and any equality correlations follow as further equi-join keys when their
+types can be hashed, and otherwise stay in the join filter. `Null`,
+dictionary, run-end encoded and union elements stay among the first keys,

Review Comment:
   “Stay among the first keys” suggests a run-end encoded element can appear in 
a planned join. The regression case in `null_aware_multi_column.slt` instead 
expects planning to fail even when that element is declared `NOT NULL`, because 
the value key is unhashable. Could the guide say these types must be treated as 
value keys and that the planner rejects an unhashable value key, as the PR 
description explains?



##########
datafusion/physical-plan/src/joins/hash_join/not_in.rs:
##########
@@ -1245,33 +1379,46 @@ where
                 self.batch_size,
             )
         };
-        let rest_hashes = &hashes[group_from..];
         // Estimate the groups from a sample first, because a sort of all the
         // rest that finds no group worth it costs as much as the wasted
         // matches did.
-        let worth_sorting = saves_work(rest_hashes.len(), 
estimate_groups(rest_hashes));
-        let mut rest: Vec<(u64, u32)> = vec![];
-        let mut groups = vec![];
-        if worth_sorting {
-            rest = rest_hashes
-                .iter()
-                .copied()
-                .zip(probe_rows[group_from..].iter().copied())
-                .collect();
-            rest.sort_unstable_by_key(|(hash, _)| *hash);
-            groups = rest.chunk_by(|a, b| a.0 == b.0).collect();
+        let mut rest = vec![];
+        if saves_work(
+            probe_rows.len() - group_from,
+            estimate_groups(&hashes[group_from..]),
+        ) {
+            rest = sort_by_hash(&hashes[group_from..], 
&probe_rows[group_from..]);
         }
-        if !worth_sorting || !saves_work(rest.len(), groups.len()) {
-            self.lookup_each_row(
+        let mut groups = hash_groups(&rest);
+        if rest.is_empty() || !saves_work(rest.len(), groups.len()) {
+            self.wasted_matches = 0;
+            let Some((offset, _)) = self.lookup_each_row(
                 scope_map,
                 hashes,
                 (!all_rows).then_some(probe_rows.as_slice()),
                 offset,
                 &to_build_row,
-                usize::MAX,
+                max_wasted_matches,
                 lookup,
-            )?;
-            return Ok(());
+            )?
+            else {
+                return Ok(());
+            };
+            let group_from = match offset {
+                AtRow { row } | MidChain { row, .. } => row,
+            };
+            if rest.is_empty() {
+                rest = sort_by_hash(&hashes[group_from..], 
&probe_rows[group_from..]);
+            } else {
+                // `probe_rows` is increasing, so the rows from `group_from` on
+                // are those not below its row, and keeping them keeps the
+                // order by hash.
+                match probe_rows.get(group_from) {
+                    Some(&first) => rest.retain(|(_, row)| *row >= first),
+                    None => rest.clear(),
+                }
+            }
+            groups = hash_groups(&rest);

Review Comment:
   The first `groups` vector stays allocated throughout the retry. This 
assignment builds its replacement before dropping it. I reproduced the overlap 
with `SELECT a NOT IN (SELECT x FROM i WHERE i.y = o.b) FROM o`, using 1,024 
outer `(NULL,0)` rows and a streaming probe starting with `(1,0)`, followed by 
1,048,576 rows with `x=1`. The first 20,000 share `y=0` and the rest have 
distinct keys. At batch size 8,192, the allocation trace shows two live 16 MiB 
group descriptor vectors here. The merge base only collects one.
   
   Could we `drop(groups)` when entering the retry at line 1393? Its 
descriptors are no longer used once the first grouping is rejected. In an 
isolated run of the production grouping code, that reduced peak allocations 
from 52.0 MiB to 43.8 MiB.



##########
datafusion/physical-plan/src/joins/hash_join/not_in.rs:
##########
@@ -1136,25 +1233,56 @@ where
             build_lookups: 0,
             #[cfg(test)]
             mark_checks: 0,
+            #[cfg(test)]
+            max_decided_rows: 0,
             f,
         }
     }
 
-    /// Pairs each of `probe_rows` with each of `build_rows`. The pairing stops
-    /// when every row that the pairs decide is marked.
-    fn cross_product(
+    /// Pairs each of the rows that `probe_rows` yields with each of
+    /// `build_rows`. The pairing stops when every row that the pairs decide
+    /// is marked.
+    ///
+    /// When the pairs decide the build rows, the unmarked build rows are
+    /// paired in chunks of `batch_size` rows, and `probe_rows` is called
+    /// again for each chunk. Neither side is then held whole: a large build
+    /// side, as when this pairing stands in for a scope map that the memory
+    /// pool declined, or a probe batch larger than the batch size. A chunk is
+    /// filled with unmarked rows only, so a build side with few unmarked rows
+    /// left pairs them all at once, as a single chunk would.
+    fn cross_product<P: Iterator<Item = u32>>(
         &mut self,
         build_rows: impl Iterator<Item = u64>,
-        probe_rows: impl Iterator<Item = u32>,
+        probe_rows: impl Fn() -> P,
     ) -> Result<()> {
-        // Drop the rows that are already marked before the first pair.
-        self.needs_refresh = true;
         match self.marks {
-            UnknownMarks::Build(_) => {
-                self.pair_unmarked(&mut build_rows.collect(), probe_rows, |b, 
p| (b, p))
+            UnknownMarks::Build(bitmap) => {
+                if probe_rows().next().is_none() {
+                    return Ok(());
+                }
+                let mut build_rows = build_rows.peekable();
+                let mut chunk = Vec::with_capacity(self.batch_size);
+                while build_rows.peek().is_some() {
+                    chunk.clear();
+                    {
+                        let bitmap = bitmap.lock();
+                        chunk.extend(
+                            build_rows
+                                .by_ref()
+                                .filter(|row| !bitmap.get_bit(*row as usize))
+                                .take(self.batch_size),
+                        );
+                    }
+                    // The chunk holds no row marked so far.
+                    self.needs_refresh = false;
+                    self.pair_unmarked(&mut chunk, probe_rows(), |b, p| (b, 
p))?;

Review Comment:
   With one NULL near the end of a large probe batch, `probe_rows()` walks the 
NULL bitmap from the beginning for every build chunk. I reproduced this with 
`(a,b) NOT IN (SELECT x,y FROM i)`, 4,194,304 build rows `(a,7)` with `a >= 
10`, and one 8,388,608-row `StreamingTable` batch of `(3,7)` ending in 
`(NULL,7)`. At the default batch size of 8,192, this walks the same bitmap 513 
times including the emptiness check. The merge base walks it once. Both 
evaluate the same candidate pairs and return the same result, but median query 
time over three local debug runs rose from 4.16 s to 5.69 s.
   
   Could we cache the eligible probe indices when they fit in one `batch_size` 
buffer and replay the iterator only when they exceed that bound? This case 
would retain one index and preserve the memory bound for dense probes.



-- 
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