@@ -124,19 +124,13 @@ fn unsliced_hash_len(data: &ArrayData) -> usize {
124124/// about 1/16 to the hashing work.
125125const SHARED_VALUES_SLICE_FACTOR : usize = 16 ;
126126
127- /// The headroom from [`SHARED_VALUES_SLICE_FACTOR`] stops at this many rows
128- /// (8 MiB of hashes), so a large merged dictionary cannot grow the hashes
129- /// buffer and scratch space to the whole build side. Past it only the
130- /// rehashing overhead grows, and it stays linear in the build rows.
131- const MAX_HASH_CHUNK_ROWS : usize = 1 << 20 ;
132-
133127/// The number of rows in each slice the hash map is filled from.
134128///
135129/// A slice is never shorter than the largest build batch or the longest
136130/// shared values under any key, which keeps hashing linear in the build rows.
137131/// The overhead ratio is per key column, so the longest key counts, not the sum.
138132/// Beyond that, a slice is [`SHARED_VALUES_SLICE_FACTOR`] times the shared
139- /// values, up to [`MAX_HASH_CHUNK_ROWS`], and never longer than the build side.
133+ /// values, and never longer than the build side.
140134fn hash_chunk_rows (
141135 left_values : & [ ArrayRef ] ,
142136 max_batch_rows : usize ,
@@ -147,9 +141,9 @@ fn hash_chunk_rows(
147141 . map ( |values| unsliced_hash_len ( & values. to_data ( ) ) )
148142 . max ( )
149143 . unwrap_or ( 0 ) ;
150- let base = max_batch_rows. max ( shared ) ;
151- let amortized = base . max ( shared. saturating_mul ( SHARED_VALUES_SLICE_FACTOR ) ) ;
152- amortized . min ( base . max ( MAX_HASH_CHUNK_ROWS ) ) . min ( num_rows)
144+ max_batch_rows
145+ . max ( shared. saturating_mul ( SHARED_VALUES_SLICE_FACTOR ) )
146+ . min ( num_rows)
153147}
154148
155149/// The longest array anywhere in `data`, including `data` itself.
@@ -739,9 +733,8 @@ impl From<&HashJoinExec> for HashJoinExecBuilder {
739733/// concatenated in their input order. The hash table is then filled from slices of the
740734/// concatenated batch, as long as the largest build-side batch. When a key has values
741735/// under a dictionary or list view, which every slice hashes in full, slices grow to 16
742- /// times the length of those values, capped at 2^20 rows unless the values alone are
743- /// longer. Filling starts from the last slice, and each slice is inserted with the offset
744- /// of its first row in the concatenated batch.
736+ /// times the length of those values. Filling starts from the last slice, and each slice
737+ /// is inserted with the offset of its first row in the concatenated batch.
745738///
746739/// Example of build phase for 3 record batches, filled from slices of 3 rows:
747740///
@@ -3278,8 +3271,7 @@ fn build_in_memory(
32783271 // Every slice shares everything beneath a dictionary or list view,
32793272 // at any depth of a key, which `create_hashes` hashes on each call:
32803273 // slices `SHARED_VALUES_SLICE_FACTOR` times longer than those values
3281- // (up to `MAX_HASH_CHUNK_ROWS`, never shorter than the values or the
3282- // largest batch) keep that cost small.
3274+ // (never shorter than the largest batch) keep that cost small.
32833275 let chunk_rows = hash_chunk_rows ( & left_values, max_batch_rows, num_rows) ;
32843276
32853277 let scratch_reservation = reservation. new_empty ( ) ;
@@ -3289,6 +3281,9 @@ fn build_in_memory(
32893281 on_left. len ( ) ,
32903282 null_equality,
32913283 ) ?) ?;
3284+ } else if chunk_rows > max_batch_rows {
3285+ // A hashes buffer longer than a batch grows with the build side
3286+ scratch_reservation. try_grow ( chunk_rows * size_of :: < u64 > ( ) ) ?;
32923287 }
32933288 // The maximum is known: avoid geometric growth and its excess capacity.
32943289 let mut hashes_buffer = Vec :: with_capacity ( chunk_rows) ;
@@ -5265,8 +5260,8 @@ mod tests {
52655260 }
52665261
52675262 /// A key with shared values gets slices `SHARED_VALUES_SLICE_FACTOR` times
5268- /// longer than those values, up to `MAX_HASH_CHUNK_ROWS` unless the values
5269- /// alone are longer, and never longer than the build side.
5263+ /// longer than those values, never shorter than the largest batch and
5264+ /// never longer than the build side.
52705265 #[ test]
52715266 fn hash_chunk_rows_amortizes_shared_values ( ) -> Result < ( ) > {
52725267 fn dict ( values : i32 ) -> Result < ArrayRef > {
@@ -5287,10 +5282,10 @@ mod tests {
52875282 // The largest batch is longer than 16 times the dictionary
52885283 ( vec ! [ dict( 100 ) ?] , 8192 , 2_097_152 , 8192 ) ,
52895284 ( vec ! [ Arc :: clone( & plain) , dict( 8192 ) ?] , 8192 , 50_000 , 50_000 ) ,
5290- // Values longer than the cap are the slice length, not 16 times them
5291- ( vec ! [ dict ( 2_097_152 ) ? ] , 8192 , 4_194_304 , 2_097_152 ) ,
5292- ( vec ! [ dict( 100_000 ) ?] , 8192 , 4_194_304 , MAX_HASH_CHUNK_ROWS ) ,
5293- // A batch longer than the cap is still the slice length
5285+ ( vec ! [ dict ( 100_000 ) ? ] , 8192 , 4_194_304 , 1_600_000 ) ,
5286+ // The build side caps the slice
5287+ ( vec ! [ dict( 2_097_152 ) ?] , 8192 , 4_194_304 , 4_194_304 ) ,
5288+ // A batch longer than 16 times the values is the slice length
52945289 ( vec ! [ dict( 1000 ) ?] , 2_097_152 , 4_194_304 , 2_097_152 ) ,
52955290 ( vec ! [ ] , 8192 , 0 , 0 ) ,
52965291 ( vec ! [ dict( 8 ) ?] , 8192 , 0 , 0 ) ,
@@ -8757,6 +8752,101 @@ mod tests {
87578752 Ok ( ( ) )
87588753 }
87598754
8755+ /// Slices longer than a build batch need a hashes buffer that grows with
8756+ /// the build side, which must be reserved.
8757+ #[ tokio:: test]
8758+ async fn join_build_hashes_buffer_is_reserved ( ) -> Result < ( ) > {
8759+ let dict_type =
8760+ DataType :: Dictionary ( Box :: new ( DataType :: Int32 ) , Box :: new ( DataType :: Int32 ) ) ;
8761+ let schema = Arc :: new ( Schema :: new ( vec ! [ Field :: new( "b1" , dict_type, false ) ] ) ) ;
8762+ // 8 batches of 1024 rows share 512 dictionary values, so the whole
8763+ // build side is hashed as one slice of 8192 rows
8764+ let values: ArrayRef = Arc :: new ( Int32Array :: from_iter_values ( 0 ..512 ) ) ;
8765+ let batches = ( 0 ..8 )
8766+ . map ( |_| {
8767+ let keys = Int32Array :: from_iter_values ( ( 0 ..1024 ) . map ( |i| i % 512 ) ) ;
8768+ let b1 =
8769+ DictionaryArray :: < Int32Type > :: try_new ( keys, Arc :: clone ( & values) ) ?;
8770+ Ok ( RecordBatch :: try_new (
8771+ Arc :: clone ( & schema) ,
8772+ vec ! [ Arc :: new( b1) ] ,
8773+ ) ?)
8774+ } )
8775+ . collect :: < Result < Vec < _ > > > ( ) ?;
8776+ let build_rows = 8 * 1024 ;
8777+ assert_eq ! (
8778+ hash_chunk_rows( & [ Arc :: clone( batches[ 0 ] . column( 0 ) ) ] , 1024 , build_rows) ,
8779+ build_rows
8780+ ) ;
8781+ let right = RecordBatch :: try_new (
8782+ Arc :: clone ( & schema) ,
8783+ vec ! [ Arc :: new( DictionaryArray :: <Int32Type >:: try_new(
8784+ Int32Array :: from( vec![ 0 ] ) ,
8785+ Arc :: clone( & values) ,
8786+ ) ?) ] ,
8787+ ) ?;
8788+ let on = vec ! [ (
8789+ Arc :: new( Column :: new_with_schema( "b1" , & schema) ?) as _,
8790+ Arc :: new( Column :: new_with_schema( "b1" , & schema) ?) as _,
8791+ ) ] ;
8792+
8793+ // The concatenated copy and the map are held while the map is filled
8794+ let copy = get_record_batch_memory_size ( & concat_batches ( & schema, & batches) ?) ;
8795+ let map =
8796+ estimate_memory_size :: < ( u32 , u64 ) > ( build_rows, size_of :: < JoinHashMapU32 > ( ) ) ?
8797+ + build_rows * size_of :: < u32 > ( ) ;
8798+ let hashes = build_rows * size_of :: < u64 > ( ) ;
8799+
8800+ for ( limit, fits) in [
8801+ // Room for the copy and the map, not for the hashes too
8802+ ( copy + map + hashes / 2 , false ) ,
8803+ ( copy + map + hashes * 3 / 2 , true ) ,
8804+ ] {
8805+ let left = TestMemoryExec :: try_new_exec (
8806+ std:: slice:: from_ref ( & batches) ,
8807+ Arc :: clone ( & schema) ,
8808+ None ,
8809+ ) ?;
8810+ let right = TestMemoryExec :: try_new_exec (
8811+ & [ vec ! [ right. clone( ) ] ] ,
8812+ Arc :: clone ( & schema) ,
8813+ None ,
8814+ ) ?;
8815+ let join = HashJoinExec :: try_new (
8816+ left,
8817+ right,
8818+ on. clone ( ) ,
8819+ None ,
8820+ & JoinType :: Inner ,
8821+ None ,
8822+ PartitionMode :: CollectLeft ,
8823+ NullEquality :: NullEqualsNothing ,
8824+ false ,
8825+ ) ?;
8826+
8827+ let runtime = RuntimeEnvBuilder :: new ( )
8828+ . with_memory_limit ( limit, 1.0 )
8829+ . build_arc ( ) ?;
8830+ let task_ctx = prepare_task_ctx ( 8192 , false ) ;
8831+ let task_ctx = Arc :: new (
8832+ TaskContext :: default ( )
8833+ . with_session_config ( task_ctx. session_config ( ) . clone ( ) )
8834+ . with_runtime ( runtime) ,
8835+ ) ;
8836+
8837+ let result = common:: collect ( join. execute ( 0 , task_ctx) ?) . await ;
8838+ if fits {
8839+ assert_eq ! ( result?. iter( ) . map( |b| b. num_rows( ) ) . sum:: <usize >( ) , 16 ) ;
8840+ } else {
8841+ assert_contains ! (
8842+ result. unwrap_err( ) . to_string( ) ,
8843+ "Resources exhausted: Additional allocation failed for HashJoinInput"
8844+ ) ;
8845+ }
8846+ }
8847+ Ok ( ( ) )
8848+ }
8849+
87608850 fn build_table_struct (
87618851 struct_name : & str ,
87628852 field_name_and_values : ( & str , & Vec < Option < i32 > > ) ,
0 commit comments