@@ -212,6 +212,7 @@ impl NewTransformPartitionBucket {
212
212
#[ allow( unused_assignments) ]
213
213
fn add_bucket ( & mut self , mut data_block : DataBlock ) -> Result < ( isize , usize ) > {
214
214
let ( mut bucket, mut partition_count) = ( 0 , 0 ) ;
215
+ let mut is_empty_block = false ;
215
216
if let Some ( block_meta) = data_block. get_meta ( ) {
216
217
if let Some ( block_meta) = AggregateMeta :: downcast_ref_from ( block_meta) {
217
218
( bucket, partition_count) = match block_meta {
@@ -250,7 +251,11 @@ impl NewTransformPartitionBucket {
250
251
if let Some ( AggregateMeta :: Spilled ( buckets_payload) ) =
251
252
AggregateMeta :: downcast_from ( meta)
252
253
{
253
- let partition_count = buckets_payload[ 0 ] . max_partition_count ;
254
+ let partition_count = if !buckets_payload. is_empty ( ) {
255
+ buckets_payload[ 0 ] . max_partition_count
256
+ } else {
257
+ MAX_PARTITION_COUNT
258
+ } ;
254
259
self . max_partition_count =
255
260
self . max_partition_count . max ( partition_count) ;
256
261
@@ -274,12 +279,14 @@ impl NewTransformPartitionBucket {
274
279
unreachable ! ( )
275
280
}
276
281
AggregateMeta :: Serialized ( payload) => {
282
+ is_empty_block = payload. data_block . is_empty ( ) ;
277
283
self . max_partition_count =
278
284
self . max_partition_count . max ( payload. max_partition_count ) ;
279
285
280
286
( payload. bucket , payload. max_partition_count )
281
287
}
282
288
AggregateMeta :: AggregatePayload ( payload) => {
289
+ is_empty_block = payload. payload . len ( ) == 0 ;
283
290
self . max_partition_count =
284
291
self . max_partition_count . max ( payload. max_partition_count ) ;
285
292
@@ -298,23 +305,25 @@ impl NewTransformPartitionBucket {
298
305
) ) ;
299
306
}
300
307
301
- if self . all_inputs_init {
302
- if partition_count != self . max_partition_count {
303
- return Err ( ErrorCode :: Internal (
308
+ if !is_empty_block {
309
+ if self . all_inputs_init {
310
+ if partition_count != self . max_partition_count {
311
+ return Err ( ErrorCode :: Internal (
304
312
"Internal, the partition count does not equal the max partition count on TransformPartitionBucket.
305
313
" ,
306
314
) ) ;
307
- }
308
- match self . buckets_blocks . entry ( bucket) {
309
- Entry :: Vacant ( v) => {
310
- v. insert ( vec ! [ data_block] ) ;
311
- }
312
- Entry :: Occupied ( mut v) => {
313
- v. get_mut ( ) . push ( data_block) ;
314
315
}
315
- } ;
316
- } else {
317
- self . unpartitioned_blocks . push ( data_block) ;
316
+ match self . buckets_blocks . entry ( bucket) {
317
+ Entry :: Vacant ( v) => {
318
+ v. insert ( vec ! [ data_block] ) ;
319
+ }
320
+ Entry :: Occupied ( mut v) => {
321
+ v. get_mut ( ) . push ( data_block) ;
322
+ }
323
+ } ;
324
+ } else {
325
+ self . unpartitioned_blocks . push ( data_block) ;
326
+ }
318
327
}
319
328
320
329
Ok ( ( bucket, partition_count) )
0 commit comments