@@ -296,6 +296,17 @@ def _merge_chunk_array(
296296 is_complete_chunk : bool ,
297297 drop_axes : tuple [int , ...],
298298 ) -> 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
299310 if chunk_selection == () or is_scalar (
300311 value .as_ndarray_like (), chunk_spec .dtype .to_native_dtype ()
301312 ):
@@ -311,20 +322,6 @@ def _merge_chunk_array(
311322 for idx in range (chunk_spec .ndim )
312323 )
313324 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
328325 chunk_array [chunk_selection ] = chunk_value
329326 return chunk_array
330327
0 commit comments