@@ -296,6 +296,17 @@ def _merge_chunk_array(
296
296
is_complete_chunk : bool ,
297
297
drop_axes : tuple [int , ...],
298
298
) -> NDBuffer :
299
+ if is_complete_chunk and value .shape == chunk_spec .shape and value [out_selection ].shape == chunk_spec .shape :
300
+ return value
301
+ if existing_chunk_array is None :
302
+ chunk_array = chunk_spec .prototype .nd_buffer .create (
303
+ shape = chunk_spec .shape ,
304
+ dtype = chunk_spec .dtype .to_native_dtype (),
305
+ order = chunk_spec .order ,
306
+ fill_value = fill_value_or_default (chunk_spec ),
307
+ )
308
+ else :
309
+ chunk_array = existing_chunk_array .copy () # make a writable copy
299
310
if chunk_selection == () or is_scalar (
300
311
value .as_ndarray_like (), chunk_spec .dtype .to_native_dtype ()
301
312
):
@@ -311,20 +322,6 @@ def _merge_chunk_array(
311
322
for idx in range (chunk_spec .ndim )
312
323
)
313
324
chunk_value = chunk_value [item ]
314
- if is_complete_chunk and chunk_value .shape == chunk_spec .shape :
315
- # TODO: For the last chunk, we could have is_complete_chunk=True
316
- # that is smaller than the chunk_spec.shape but this throws
317
- # an error in the _decode_single
318
- return chunk_value .copy () # make a writable copy
319
- if existing_chunk_array is None :
320
- chunk_array = chunk_spec .prototype .nd_buffer .create (
321
- shape = chunk_spec .shape ,
322
- dtype = chunk_spec .dtype .to_native_dtype (),
323
- order = chunk_spec .order ,
324
- fill_value = fill_value_or_default (chunk_spec ),
325
- )
326
- else :
327
- chunk_array = existing_chunk_array .copy () # make a writable copy
328
325
chunk_array [chunk_selection ] = chunk_value
329
326
return chunk_array
330
327
0 commit comments