@@ -54,8 +54,8 @@ use datafusion_physical_expr_common::datum::apply_cmp;
5454use crate :: joins:: PartitionMode ;
5555use crate :: joins:: SharedBitmapBuilder ;
5656use crate :: joins:: hash_join:: exec:: { JoinLeftData , new_join_hashmap} ;
57- use crate :: joins:: join_hash_map:: JoinHashMapType ;
58- use crate :: joins:: join_hash_map:: ProbeOffset :: AtRow ;
57+ use crate :: joins:: join_hash_map:: ProbeOffset :: { AtRow , MidChain } ;
58+ use crate :: joins:: join_hash_map:: { JoinHashMapType , ProbeOffset } ;
5959use crate :: joins:: probe_completion:: ProbeSideSummary ;
6060use crate :: joins:: utils:: {
6161 BuildProbeJoinMetrics , JoinFilter , apply_join_filter_to_indices,
@@ -991,7 +991,8 @@ pub(super) struct LookupBuffers {
991991/// chunks of at most `batch_size` pairs.
992992///
993993/// `f` marks rows (see [`UnknownMarks`]), and a marked row cannot become
994- /// unmarked. So a cross product drops the marked rows again after each chunk.
994+ /// unmarked. So the pairing drops the pairs that decide a marked row before
995+ /// it sends them, and drops the marked rows again after each chunk.
995996struct UnmarkedPairs < ' a , F > {
996997 marks : UnknownMarks < ' a > ,
997998 batch_size : usize ,
@@ -1000,6 +1001,9 @@ struct UnmarkedPairs<'a, F> {
10001001 /// Whether `f` was called since the rows that the pairs decide were last
10011002 /// refreshed.
10021003 marks_since_refresh : bool ,
1004+ /// The matches that [`Self::scope_matches`] dropped because the row they
1005+ /// decide was marked.
1006+ skipped_matches : usize ,
10031007 f : F ,
10041008}
10051009
@@ -1014,6 +1018,7 @@ where
10141018 build_chunk : Vec :: with_capacity ( batch_size) ,
10151019 probe_chunk : Vec :: with_capacity ( batch_size) ,
10161020 marks_since_refresh : false ,
1021+ skipped_matches : 0 ,
10171022 f,
10181023 }
10191024 }
@@ -1044,6 +1049,14 @@ where
10441049 /// hashes can have different scope keys, so `f` must still check that the
10451050 /// scope keys are equal.
10461051 ///
1052+ /// The pairing starts with a lookup per probe row, like the join itself.
1053+ /// When the marks are on the build rows, each probe row still finds all
1054+ /// the build rows in its scope, marked or not. So a marked build row costs
1055+ /// a little for every later probe row in its scope. When those costs add
1056+ /// up to more than a sort of the probe rows, the rest of the probe rows
1057+ /// are grouped by hash. Each hash is then looked up once, and its group
1058+ /// is paired like a cross product that drops the build rows as they are
1059+ /// marked.
10471060 fn scope_matches (
10481061 & mut self ,
10491062 scope_map : & dyn JoinHashMapType ,
@@ -1057,20 +1070,82 @@ where
10571070 . iter ( )
10581071 . map ( |row| probe_hashes[ * row as usize ] )
10591072 . collect ( ) ;
1073+
1074+ // A sort costs about as much as this many skipped matches per probe
1075+ // row. A probe row is looked up once, so the marks on probe rows
1076+ // never make grouping worth it.
1077+ const SKIPPED_MATCHES_PER_SORTED_ROW : usize = 8 ;
1078+ let max_skipped_matches = match self . marks {
1079+ UnknownMarks :: Build ( _) => SKIPPED_MATCHES_PER_SORTED_ROW * probe_rows. len ( ) ,
1080+ UnknownMarks :: Probe ( _) => usize:: MAX ,
1081+ } ;
1082+ self . skipped_matches = 0 ;
10601083 let mut offset = AtRow { row : 0 } ;
1061- loop {
1062- let next_offset = scope_map. get_matched_indices_with_limit_offset (
1084+ let group_from = loop {
1085+ let next_offset = self . lookup_unmarked (
1086+ scope_map,
10631087 & hashes,
1064- None ,
1065- self . batch_size ,
1088+ |position| probe_rows[ position as usize ] ,
1089+ offset,
1090+ & to_build_row,
1091+ lookup,
1092+ ) ;
1093+ for ( probe_row, build_row) in
1094+ lookup. probe_indices . iter ( ) . zip ( lookup. build_indices . iter ( ) )
1095+ {
1096+ self . push ( * build_row, * probe_row) ?;
1097+ }
1098+ let Some ( next_offset) = next_offset else {
1099+ return Ok ( ( ) ) ;
1100+ } ;
1101+ offset = next_offset;
1102+ if self . skipped_matches > max_skipped_matches {
1103+ break match next_offset {
1104+ AtRow { row } | MidChain { row, .. } => row,
1105+ } ;
1106+ }
1107+ } ;
1108+
1109+ // Send the pairs that are left, so the lookups below see their marks.
1110+ self . send ( ) ?;
1111+
1112+ // Group the probe rows from the one being looked up. Its build rows
1113+ // that were already paired are paired again, which changes no mark.
1114+ let mut rest: Vec < ( u64 , u32 ) > = hashes[ group_from..]
1115+ . iter ( )
1116+ . copied ( )
1117+ . zip ( probe_rows[ group_from..] . iter ( ) . copied ( ) )
1118+ . collect ( ) ;
1119+ rest. sort_unstable_by_key ( |( hash, _) | * hash) ;
1120+ let groups: Vec < _ > = rest. chunk_by ( |a, b| a. 0 == b. 0 ) . collect ( ) ;
1121+ let group_hashes: Vec < u64 > = groups. iter ( ) . map ( |group| group[ 0 ] . 0 ) . collect ( ) ;
1122+
1123+ let mut build_rows = vec ! [ ] ;
1124+ let mut offset = AtRow { row : 0 } ;
1125+ loop {
1126+ // A lookup chunk holds the build rows of one or more groups, in
1127+ // group order.
1128+ let next_offset = self . lookup_unmarked (
1129+ scope_map,
1130+ & group_hashes,
1131+ |group| group,
10661132 offset,
1067- & mut lookup . probe_indices ,
1068- & mut lookup. build_indices ,
1133+ & to_build_row ,
1134+ lookup,
10691135 ) ;
1070- for i in 0 ..lookup. build_indices . len ( ) {
1071- self . push (
1072- to_build_row ( lookup. build_indices [ i] ) ,
1073- probe_rows[ lookup. probe_indices [ i] as usize ] ,
1136+ let mut matches = lookup. build_indices . as_slice ( ) ;
1137+ for group in lookup. probe_indices . chunk_by ( |a, b| a == b) {
1138+ let group_matches;
1139+ ( group_matches, matches) = matches. split_at ( group. len ( ) ) ;
1140+ build_rows. clear ( ) ;
1141+ build_rows. extend_from_slice ( group_matches) ;
1142+ // A build row belongs to one group only, so the marks made
1143+ // for the earlier groups do not concern these build rows.
1144+ self . marks_since_refresh = false ;
1145+ self . pair_unmarked (
1146+ & mut build_rows,
1147+ groups[ group[ 0 ] as usize ] . iter ( ) . map ( |( _, row) | * row) ,
1148+ |b, p| ( b, p) ,
10741149 ) ?;
10751150 }
10761151 let Some ( next_offset) = next_offset else {
@@ -1080,6 +1155,49 @@ where
10801155 }
10811156 }
10821157
1158+ /// Looks up the next chunk of matches of `hashes` in `scope_map`, from
1159+ /// `offset`, into the buffers: the position in `hashes` translated with
1160+ /// `to_probe_row`, and the build row translated with `to_build_row`. Drops the
1161+ /// matches that decide a marked row, and adds their number to
1162+ /// `skipped_matches`.
1163+ fn lookup_unmarked (
1164+ & mut self ,
1165+ scope_map : & dyn JoinHashMapType ,
1166+ hashes : & [ u64 ] ,
1167+ to_probe_row : impl Fn ( u32 ) -> u32 ,
1168+ offset : ProbeOffset ,
1169+ to_build_row : impl Fn ( u64 ) -> u64 ,
1170+ lookup : & mut LookupBuffers ,
1171+ ) -> Option < ProbeOffset > {
1172+ let LookupBuffers {
1173+ probe_indices,
1174+ build_indices,
1175+ } = lookup;
1176+ let next_offset = scope_map. get_matched_indices_with_limit_offset (
1177+ hashes,
1178+ None ,
1179+ self . batch_size ,
1180+ offset,
1181+ probe_indices,
1182+ build_indices,
1183+ ) ;
1184+ let bitmap = self . marks . bitmap ( ) . lock ( ) ;
1185+ let mut kept = 0 ;
1186+ for i in 0 ..build_indices. len ( ) {
1187+ let build = to_build_row ( build_indices[ i] ) ;
1188+ let probe = to_probe_row ( probe_indices[ i] ) ;
1189+ if !bitmap. get_bit ( self . marks . unknown_row ( build, probe) ) {
1190+ build_indices[ kept] = build;
1191+ probe_indices[ kept] = probe;
1192+ kept += 1 ;
1193+ }
1194+ }
1195+ self . skipped_matches += build_indices. len ( ) - kept;
1196+ build_indices. truncate ( kept) ;
1197+ probe_indices. truncate ( kept) ;
1198+ next_offset
1199+ }
1200+
10831201 /// Sends the pairs that are left.
10841202 fn finish ( mut self ) -> Result < ( ) > {
10851203 self . send ( )
@@ -1213,12 +1331,13 @@ mod tests {
12131331 Ok ( pairs. into_iter ( ) . max ( ) . unwrap_or ( 0 ) )
12141332 }
12151333
1216- /// A correlation scope with many build rows and many probe rows. A lookup
1217- /// per probe row pairs every build row once per probe row in its scope,
1218- /// even after the build row is marked. So the pairs made for one build
1219- /// row grow with the number of probe rows.
1334+ /// A correlation scope with many build rows and many probe rows: once a
1335+ /// build row is marked, the later probe rows in its scope must not pair
1336+ /// with it again. So the pairs made for one build row must not grow with
1337+ /// the number of probe rows. A lookup per probe row pairs every build row
1338+ /// once per probe row in its scope.
12201339 #[ test]
1221- fn scope_lookup_pairs_grow_with_probe_rows ( ) -> Result < ( ) > {
1340+ fn unmarked_pairs_do_not_grow_with_probe_rows ( ) -> Result < ( ) > {
12221341 // Build rows 0..60 are in the scope with hash 1, rows 60..100 in the
12231342 // scope with hash 2. The probe rows alternate between both scopes.
12241343 let build_hashes: Vec < u64 > = ( 0 ..100 ) . map ( |row| 1 + ( row >= 60 ) as u64 ) . collect ( ) ;
@@ -1229,8 +1348,10 @@ mod tests {
12291348 max_pairs_per_build_row ( & build_hashes, & probe_hashes, batch_size)
12301349 } ) ;
12311350 let ( few, many) = ( few?, many?) ;
1232- assert_eq ! ( few, 500 , "batch_size {batch_size}" ) ;
1233- assert_eq ! ( many, 5_000 , "batch_size {batch_size}" ) ;
1351+ assert ! (
1352+ many <= few,
1353+ "batch_size {batch_size}: {few} pairs for 1000 probe rows, {many} for 10000"
1354+ ) ;
12341355 }
12351356 Ok ( ( ) )
12361357 }
0 commit comments