@@ -1244,22 +1244,31 @@ where
12441244 /// is marked.
12451245 ///
12461246 /// When the pairs decide the build rows, the unmarked build rows are
1247- /// paired in chunks of `batch_size` rows, and `probe_rows` is called
1248- /// again for each chunk. Neither side is then held whole: a large build
1249- /// side, as when this pairing stands in for a scope map that the memory
1250- /// pool declined, or a probe batch larger than the batch size. A chunk is
1251- /// filled with unmarked rows only, so a build side with few unmarked rows
1252- /// left pairs them all at once, as a single chunk would.
1247+ /// paired in chunks of `batch_size` rows. Neither side is then held
1248+ /// whole: a large build side, as when this pairing stands in for a scope
1249+ /// map that the memory pool declined, or a probe batch larger than the
1250+ /// batch size. Up to `batch_size` probe rows are kept for every chunk.
1251+ /// More are produced again for each chunk by calling `probe_rows`, which
1252+ /// costs less than pairing them with the chunk. A chunk is filled with
1253+ /// unmarked rows only, so a build side with few unmarked rows left pairs
1254+ /// them all at once, as a single chunk would.
12531255 fn cross_product < P : Iterator < Item = u32 > > (
12541256 & mut self ,
12551257 build_rows : impl Iterator < Item = u64 > ,
12561258 probe_rows : impl Fn ( ) -> P ,
12571259 ) -> Result < ( ) > {
12581260 match self . marks {
12591261 UnknownMarks :: Build ( bitmap) => {
1260- if probe_rows ( ) . next ( ) . is_none ( ) {
1262+ // A few probe rows can be spread over a large probe batch, so
1263+ // keep them rather than walk the batch for every chunk.
1264+ let mut kept: Vec < u32 > = probe_rows ( ) . take ( self . batch_size + 1 ) . collect ( ) ;
1265+ if kept. is_empty ( ) {
12611266 return Ok ( ( ) ) ;
12621267 }
1268+ let keep_probe_rows = kept. len ( ) <= self . batch_size ;
1269+ if !keep_probe_rows {
1270+ kept = vec ! [ ] ;
1271+ }
12631272 let mut build_rows = build_rows. peekable ( ) ;
12641273 let mut chunk = Vec :: with_capacity ( self . batch_size ) ;
12651274 while build_rows. peek ( ) . is_some ( ) {
@@ -1275,7 +1284,13 @@ where
12751284 }
12761285 // The chunk holds no row marked so far.
12771286 self . needs_refresh = false ;
1278- self . pair_unmarked ( & mut chunk, probe_rows ( ) , |b, p| ( b, p) ) ?;
1287+ if keep_probe_rows {
1288+ self . pair_unmarked ( & mut chunk, kept. iter ( ) . copied ( ) , |b, p| {
1289+ ( b, p)
1290+ } ) ?;
1291+ } else {
1292+ self . pair_unmarked ( & mut chunk, probe_rows ( ) , |b, p| ( b, p) ) ?;
1293+ }
12791294 }
12801295 Ok ( ( ) )
12811296 }
@@ -1391,6 +1406,8 @@ where
13911406 }
13921407 let mut groups = hash_groups ( & rest) ;
13931408 if rest. is_empty ( ) || !saves_work ( rest. len ( ) , groups. len ( ) ) {
1409+ // The groups are found again from where the lookups stop.
1410+ drop ( groups) ;
13941411 self . wasted_matches = 0 ;
13951412 let Some ( ( offset, _) ) = self . lookup_each_row (
13961413 scope_map,
@@ -1916,6 +1933,7 @@ mod tests {
19161933 Ok ( ( ) )
19171934 } ;
19181935
1936+ let probe_walks = std:: cell:: Cell :: new ( 0 ) ;
19191937 let mut pairs = UnmarkedPairs :: new ( marks, self . batch_size , mark) ;
19201938 if self . through_scope_map {
19211939 let mut scope_map = JoinHashMapU64 :: with_capacity ( self . build_rows . len ( ) ) ;
@@ -1937,6 +1955,7 @@ mod tests {
19371955 ) ?;
19381956 } else {
19391957 pairs. cross_product ( self . build_rows . iter ( ) . copied ( ) , || {
1958+ probe_walks. set ( probe_walks. get ( ) + 1 ) ;
19401959 self . probe_rows . iter ( ) . copied ( )
19411960 } ) ?;
19421961 }
@@ -1958,6 +1977,7 @@ mod tests {
19581977 max_decided_rows,
19591978 sends,
19601979 sent_pairs,
1980+ probe_walks : probe_walks. get ( ) ,
19611981 foreign_marks,
19621982 } ,
19631983 ) )
@@ -1975,6 +1995,8 @@ mod tests {
19751995 /// The chunks of pairs sent, and the pairs in them.
19761996 sends : usize ,
19771997 sent_pairs : usize ,
1998+ /// The calls of the `probe_rows` of a cross product.
1999+ probe_walks : usize ,
19782000 /// The rows that `foreign_marks` marked.
19792001 foreign_marks : Vec < usize > ,
19802002 }
@@ -2163,8 +2185,9 @@ mod tests {
21632185 }
21642186
21652187 /// A cross product with marks on the build rows holds at most a batch of
2166- /// build rows at once, however many build rows there are. No pair passes,
2167- /// so no row is marked and every build row is paired.
2188+ /// build rows at once, however many build rows there are, and at most a
2189+ /// batch of probe rows. No pair passes, so no row is marked and every
2190+ /// build row is paired.
21682191 #[ test]
21692192 fn cross_product_chunks_the_build_rows ( ) -> Result < ( ) > {
21702193 let ( num_build, batch_size) = ( 10_000 , 64 ) ;
@@ -2185,6 +2208,17 @@ mod tests {
21852208 let ( marks, stats) = case. run ( ) ?;
21862209 assert ! ( !marks. contains( & true ) ) ;
21872210 assert_eq ! ( stats. max_decided_rows, batch_size) ;
2211+ // The three probe rows are kept for all 157 chunks.
2212+ assert_eq ! ( stats. probe_walks, 1 ) ;
2213+ // More probe rows than a batch are produced again for each chunk.
2214+ let ( marks, stats) = Case {
2215+ probe_keys : vec ! [ 0 ; batch_size + 1 ] ,
2216+ probe_rows : ( 0 ..=batch_size as u32 ) . collect ( ) ,
2217+ ..case
2218+ }
2219+ . run ( ) ?;
2220+ assert ! ( !marks. contains( & true ) ) ;
2221+ assert_eq ! ( stats. probe_walks, 1 + num_build. div_ceil( batch_size) ) ;
21882222 Ok ( ( ) )
21892223 }
21902224
0 commit comments