@@ -301,6 +301,22 @@ def _merge_chunk_array(
301301is_complete_chunk :bool ,
302302drop_axes :tuple [int , ...],
303303 )-> NDBuffer :
304+ if (
305+ is_complete_chunk
306+ and value .shape == chunk_spec .shape
307+ # Guard that this is not a partial chunk at the end with is_complete_chunk=True
308+ and value [out_selection ].shape == chunk_spec .shape
309+ ):
310+ return value
311+ if existing_chunk_array is None :
312+ chunk_array = chunk_spec .prototype .nd_buffer .create (
313+ shape = chunk_spec .shape ,
314+ dtype = chunk_spec .dtype .to_native_dtype (),
315+ order = chunk_spec .order ,
316+ fill_value = fill_value_or_default (chunk_spec ),
317+ )
318+ else :
319+ chunk_array = existing_chunk_array .copy ()# make a writable copy
304320if chunk_selection == ()or is_scalar (
305321value .as_ndarray_like (),chunk_spec .dtype .to_native_dtype ()
306322 ):
@@ -316,20 +332,6 @@ def _merge_chunk_array(
316332for idx in range (chunk_spec .ndim )
317333 )
318334chunk_value = chunk_value [item ]
319- if is_complete_chunk and chunk_value .shape == chunk_spec .shape :
320- # TODO: For the last chunk, we could have is_complete_chunk=True
321- # that is smaller than the chunk_spec.shape but this throws
322- # an error in the _decode_single
323- return chunk_value
324- if existing_chunk_array is None :
325- chunk_array = chunk_spec .prototype .nd_buffer .create (
326- shape = chunk_spec .shape ,
327- dtype = chunk_spec .dtype .to_native_dtype (),
328- order = chunk_spec .order ,
329- fill_value = fill_value_or_default (chunk_spec ),
330- )
331- else :
332- chunk_array = existing_chunk_array .copy ()# make a writable copy
333335chunk_array [chunk_selection ]= chunk_value
334336return chunk_array
335337