Writing from one zarr.Array to another (maybe lazily)

[ilan-gold]

I believe that previously something like:

import zarr
import numpy as np

g = zarr.open("foo.zarr")
g["bar"] = np.arange(100)
g["zap"] = np.arange(100) # without this, I get a different, but also interesting error, about loops in loops
g["zap"][...] = g["bar"]

would lazily write data, or at least this appeared to used to be functional. Is there interest in resurrecting this? Am I missing something?

Error with setting first from numpy:

File ~/Projects/Theis/zarr-python/src/zarr/core/array.py:2472, in Array.__setitem__(self, selection, value)
   2470     self.vindex[cast(CoordinateSelection | MaskSelection, selection)] = value
   2471 elif is_pure_orthogonal_indexing(pure_selection, self.ndim):
-> 2472     self.set_orthogonal_selection(pure_selection, value, fields=fields)
   2473 else:
   2474     self.set_basic_selection(cast(BasicSelection, pure_selection), value, fields=fields)

File ~/Projects/Theis/zarr-python/src/zarr/_compat.py:43, in _deprecate_positional_args.<locals>._inner_deprecate_positional_args.<locals>.inner_f(*args, **kwargs)
     41 extra_args = len(args) - len(all_args)
     42 if extra_args <= 0:
---> 43     return f(*args, **kwargs)
     45 # extra_args > 0
     46 args_msg = [
     47     f"{name}={arg}"
     48     for name, arg in zip(kwonly_args[:extra_args], args[-extra_args:], strict=False)
     49 ]

File ~/Projects/Theis/zarr-python/src/zarr/core/array.py:2928, in Array.set_orthogonal_selection(self, selection, value, fields, prototype)
   2926     prototype = default_buffer_prototype()
   2927 indexer = OrthogonalIndexer(selection, self.shape, self.metadata.chunk_grid)
-> 2928 return sync(
   2929     self._async_array._set_selection(indexer, value, fields=fields, prototype=prototype)
   2930 )

File ~/Projects/Theis/zarr-python/src/zarr/core/sync.py:142, in sync(coro, loop, timeout)
    139 return_result = next(iter(finished)).result()
    141 if isinstance(return_result, BaseException):
--> 142     raise return_result
    143 else:
    144     return return_result

File ~/Projects/Theis/zarr-python/src/zarr/core/sync.py:98, in _runner(coro)
     93 """
     94 Await a coroutine and return the result of running it. If awaiting the coroutine raises an
     95 exception, the exception will be returned.
     96 """
     97 try:
---> 98     return await coro
     99 except Exception as ex:
    100     return ex

File ~/Projects/Theis/zarr-python/src/zarr/core/array.py:1362, in AsyncArray._set_selection(self, indexer, value, prototype, fields)
   1359     _config = replace(_config, order=self.metadata.order)
   1361 # merging with existing data and encoding chunks
-> 1362 await self.codec_pipeline.write(
   1363     [
   1364         (
   1365             self.store_path / self.metadata.encode_chunk_key(chunk_coords),
   1366             self.metadata.get_chunk_spec(chunk_coords, _config, prototype),
   1367             chunk_selection,
   1368             out_selection,
   1369         )
   1370         for chunk_coords, chunk_selection, out_selection in indexer
   1371     ],
   1372     value_buffer,
   1373     drop_axes=indexer.drop_axes,
   1374 )

File ~/Projects/Theis/zarr-python/src/zarr/core/codec_pipeline.py:468, in BatchedCodecPipeline.write(self, batch_info, value, drop_axes)
    462 async def write(
    463     self,
    464     batch_info: Iterable[tuple[ByteSetter, ArraySpec, SelectorTuple, SelectorTuple]],
    465     value: NDBuffer,
    466     drop_axes: tuple[int, ...] = (),
    467 ) -> None:
--> 468     await concurrent_map(
    469         [
    470             (single_batch_info, value, drop_axes)
    471             for single_batch_info in batched(batch_info, self.batch_size)
    472         ],
    473         self.write_batch,
    474         config.get("async.concurrency"),
    475     )

File ~/Projects/Theis/zarr-python/src/zarr/core/common.py:68, in concurrent_map(items, func, limit)
     65     async with sem:
     66         return await func(*item)
---> 68 return await asyncio.gather(*[asyncio.ensure_future(run(item)) for item in items])

File ~/Projects/Theis/zarr-python/src/zarr/core/common.py:66, in concurrent_map.<locals>.run(item)
     64 async def run(item: tuple[Any]) -> V:
     65     async with sem:
---> 66         return await func(*item)

File ~/Projects/Theis/zarr-python/src/zarr/core/codec_pipeline.py:403, in BatchedCodecPipeline.write_batch(self, batch_info, value, drop_axes)
    400         else:
    401             chunk_array_batch.append(chunk_array)
--> 403 chunk_bytes_batch = await self.encode_batch(
    404     [
    405         (chunk_array, chunk_spec)
    406         for chunk_array, (_, chunk_spec, _, _) in zip(
    407             chunk_array_batch, batch_info, strict=False
    408         )
    409     ],
    410 )
    412 async def _write_key(byte_setter: ByteSetter, chunk_bytes: Buffer | None) -> None:
    413     if chunk_bytes is None:

File ~/Projects/Theis/zarr-python/src/zarr/core/codec_pipeline.py:210, in BatchedCodecPipeline.encode_batch(self, chunk_arrays_and_specs)
    205     chunk_array_batch = await aa_codec.encode(
    206         zip(chunk_array_batch, chunk_specs, strict=False)
    207     )
    208     chunk_specs = resolve_batched(aa_codec, chunk_specs)
--> 210 chunk_bytes_batch = await self.array_bytes_codec.encode(
    211     zip(chunk_array_batch, chunk_specs, strict=False)
    212 )
    213 chunk_specs = resolve_batched(self.array_bytes_codec, chunk_specs)
    215 for bb_codec in self.bytes_bytes_codecs:

File ~/Projects/Theis/zarr-python/src/zarr/abc/codec.py:152, in BaseCodec.encode(self, chunks_and_specs)
    136 async def encode(
    137     self,
    138     chunks_and_specs: Iterable[tuple[CodecInput | None, ArraySpec]],
    139 ) -> Iterable[CodecOutput | None]:
    140     """Encodes a batch of chunks.
    141     Chunks can be None in which case they are ignored by the codec.
    142 
   (...)
    150     Iterable[CodecOutput | None]
    151     """
--> 152     return await _batching_helper(self._encode_single, chunks_and_specs)

File ~/Projects/Theis/zarr-python/src/zarr/abc/codec.py:407, in _batching_helper(func, batch_info)
    403 async def _batching_helper(
    404     func: Callable[[CodecInput, ArraySpec], Awaitable[CodecOutput | None]],
    405     batch_info: Iterable[tuple[CodecInput | None, ArraySpec]],
    406 ) -> list[CodecOutput | None]:
--> 407     return await concurrent_map(
    408         list(batch_info),
    409         _noop_for_none(func),
    410         config.get("async.concurrency"),
    411     )

File ~/Projects/Theis/zarr-python/src/zarr/core/common.py:68, in concurrent_map(items, func, limit)
     65     async with sem:
     66         return await func(*item)
---> 68 return await asyncio.gather(*[asyncio.ensure_future(run(item)) for item in items])

File ~/Projects/Theis/zarr-python/src/zarr/core/common.py:66, in concurrent_map.<locals>.run(item)
     64 async def run(item: tuple[Any]) -> V:
     65     async with sem:
---> 66         return await func(*item)

File ~/Projects/Theis/zarr-python/src/zarr/abc/codec.py:420, in _noop_for_none.<locals>.wrap(chunk, chunk_spec)
    418 if chunk is None:
    419     return None
--> 420 return await func(chunk, chunk_spec)

File ~/Projects/Theis/zarr-python/src/zarr/codecs/_v2.py:83, in V2Codec._encode_single(self, chunk_array, chunk_spec)
     80 chunk = chunk_array.as_ndarray_like()
     82 # ensure contiguous and correct order
---> 83 chunk = chunk.astype(chunk_spec.dtype, order=chunk_spec.order, copy=False)
     85 # apply filters
     86 if self.filters:

AttributeError: 'Array' object has no attribute 'astype'

and without

---------------------------------------------------------------------------
SyncError                                 Traceback (most recent call last)
Cell In[7], line 1
----> 1 g["foo"][...] = g["bar"]

File ~/Projects/Theis/zarr-python/src/zarr/core/array.py:2474, in Array.__setitem__(self, selection, value)
   2472     self.set_orthogonal_selection(pure_selection, value, fields=fields)
   2473 else:
-> 2474     self.set_basic_selection(cast(BasicSelection, pure_selection), value, fields=fields)

File ~/Projects/Theis/zarr-python/src/zarr/_compat.py:43, in _deprecate_positional_args.<locals>._inner_deprecate_positional_args.<locals>.inner_f(*args, **kwargs)
     41 extra_args = len(args) - len(all_args)
     42 if extra_args <= 0:
---> 43     return f(*args, **kwargs)
     45 # extra_args > 0
     46 args_msg = [
     47     f"{name}={arg}"
     48     for name, arg in zip(kwonly_args[:extra_args], args[-extra_args:], strict=False)
     49 ]

File ~/Projects/Theis/zarr-python/src/zarr/core/array.py:2694, in Array.set_basic_selection(self, selection, value, fields, prototype)
   2692     prototype = default_buffer_prototype()
   2693 indexer = BasicIndexer(selection, self.shape, self.metadata.chunk_grid)
-> 2694 sync(self._async_array._set_selection(indexer, value, fields=fields, prototype=prototype))

File ~/Projects/Theis/zarr-python/src/zarr/core/sync.py:142, in sync(coro, loop, timeout)
    139 return_result = next(iter(finished)).result()
    141 if isinstance(return_result, BaseException):
--> 142     raise return_result
    143 else:
    144     return return_result

File ~/Projects/Theis/zarr-python/src/zarr/core/sync.py:98, in _runner(coro)
     93 """
     94 Await a coroutine and return the result of running it. If awaiting the coroutine raises an
     95 exception, the exception will be returned.
     96 """
     97 try:
---> 98     return await coro
     99 except Exception as ex:
    100     return ex

File ~/Projects/Theis/zarr-python/src/zarr/core/array.py:1362, in AsyncArray._set_selection(self, indexer, value, prototype, fields)
   1359     _config = replace(_config, order=self.metadata.order)
   1361 # merging with existing data and encoding chunks
-> 1362 await self.codec_pipeline.write(
   1363     [
   1364         (
   1365             self.store_path / self.metadata.encode_chunk_key(chunk_coords),
   1366             self.metadata.get_chunk_spec(chunk_coords, _config, prototype),
   1367             chunk_selection,
   1368             out_selection,
   1369         )
   1370         for chunk_coords, chunk_selection, out_selection in indexer
   1371     ],
   1372     value_buffer,
   1373     drop_axes=indexer.drop_axes,
   1374 )

File ~/Projects/Theis/zarr-python/src/zarr/core/codec_pipeline.py:468, in BatchedCodecPipeline.write(self, batch_info, value, drop_axes)
    462 async def write(
    463     self,
    464     batch_info: Iterable[tuple[ByteSetter, ArraySpec, SelectorTuple, SelectorTuple]],
    465     value: NDBuffer,
    466     drop_axes: tuple[int, ...] = (),
    467 ) -> None:
--> 468     await concurrent_map(
    469         [
    470             (single_batch_info, value, drop_axes)
    471             for single_batch_info in batched(batch_info, self.batch_size)
    472         ],
    473         self.write_batch,
    474         config.get("async.concurrency"),
    475     )

File ~/Projects/Theis/zarr-python/src/zarr/core/common.py:68, in concurrent_map(items, func, limit)
     65     async with sem:
     66         return await func(*item)
---> 68 return await asyncio.gather(*[asyncio.ensure_future(run(item)) for item in items])

File ~/Projects/Theis/zarr-python/src/zarr/core/common.py:66, in concurrent_map.<locals>.run(item)
     64 async def run(item: tuple[Any]) -> V:
     65     async with sem:
---> 66         return await func(*item)

File ~/Projects/Theis/zarr-python/src/zarr/core/codec_pipeline.py:396, in BatchedCodecPipeline.write_batch(self, batch_info, value, drop_axes)
    394     chunk_array_batch.append(None)  # type: ignore[unreachable]
    395 else:
--> 396     if not chunk_spec.config.write_empty_chunks and chunk_array.all_equal(
    397         chunk_spec.fill_value
    398     ):
    399         chunk_array_batch.append(None)
    400     else:

File ~/Projects/Theis/zarr-python/src/zarr/core/buffer/core.py:471, in NDBuffer.all_equal(self, other, equal_nan)
    467     return False
    468 # use array_equal to obtain equal_nan=True functionality
    469 # Since fill-value is a scalar, isn't there a faster path than allocating a new array for fill value
    470 # every single time we have to write data?
--> 471 _data, other = np.broadcast_arrays(self._data, other)
    472 return np.array_equal(
    473     self._data,
    474     other,
    475     equal_nan=equal_nan if self._data.dtype.kind not in "USTOV" else False,
    476 )

File ~/Projects/Theis/anndata/venv/lib/python3.12/site-packages/numpy/lib/_stride_tricks_impl.py:549, in broadcast_arrays(subok, *args)
    491 """
    492 Broadcast any number of arrays against each other.
    493 
   (...)
    542 
    543 """
    544 # nditer is not used here to avoid the limit of 32 arrays.
    545 # Otherwise, something like the following one-liner would suffice:
    546 # return np.nditer(args, flags=['multi_index', 'zerosize_ok'],
    547 #                  order='C').itviews
--> 549 args = tuple(np.array(_m, copy=None, subok=subok) for _m in args)
    551 shape = _broadcast_shape(*args)
    553 if all(array.shape == shape for array in args):
    554     # Common case where nothing needs to be broadcasted.

File ~/Projects/Theis/anndata/venv/lib/python3.12/site-packages/numpy/lib/_stride_tricks_impl.py:549, in <genexpr>(.0)
    491 """
    492 Broadcast any number of arrays against each other.
    493 
   (...)
    542 
    543 """
    544 # nditer is not used here to avoid the limit of 32 arrays.
    545 # Otherwise, something like the following one-liner would suffice:
    546 # return np.nditer(args, flags=['multi_index', 'zerosize_ok'],
    547 #                  order='C').itviews
--> 549 args = tuple(np.array(_m, copy=None, subok=subok) for _m in args)
    551 shape = _broadcast_shape(*args)
    553 if all(array.shape == shape for array in args):
    554     # Common case where nothing needs to be broadcasted.

File ~/Projects/Theis/zarr-python/src/zarr/core/array.py:2219, in Array.__array__(self, dtype, copy)
   2216     msg = "`copy=False` is not supported. This method always creates a copy."
   2217     raise ValueError(msg)
-> 2219 arr_np = self[...]
   2221 if dtype is not None:
   2222     arr_np = arr_np.astype(dtype)

File ~/Projects/Theis/zarr-python/src/zarr/core/array.py:2375, in Array.__getitem__(self, selection)
   2373     return self.get_orthogonal_selection(pure_selection, fields=fields)
   2374 else:
-> 2375     return self.get_basic_selection(cast(BasicSelection, pure_selection), fields=fields)

File ~/Projects/Theis/zarr-python/src/zarr/_compat.py:43, in _deprecate_positional_args.<locals>._inner_deprecate_positional_args.<locals>.inner_f(*args, **kwargs)
     41 extra_args = len(args) - len(all_args)
     42 if extra_args <= 0:
---> 43     return f(*args, **kwargs)
     45 # extra_args > 0
     46 args_msg = [
     47     f"{name}={arg}"
     48     for name, arg in zip(kwonly_args[:extra_args], args[-extra_args:], strict=False)
     49 ]

File ~/Projects/Theis/zarr-python/src/zarr/core/array.py:2591, in Array.get_basic_selection(self, selection, out, prototype, fields)
   2589 if prototype is None:
   2590     prototype = default_buffer_prototype()
-> 2591 return sync(
   2592     self._async_array._get_selection(
   2593         BasicIndexer(selection, self.shape, self.metadata.chunk_grid),
   2594         out=out,
   2595         fields=fields,
   2596         prototype=prototype,
   2597     )
   2598 )

File ~/Projects/Theis/zarr-python/src/zarr/core/sync.py:129, in sync(coro, loop, timeout)
    127     loop0 = asyncio.events.get_running_loop()
    128     if loop0 is loop:
--> 129         raise SyncError("Calling sync() from within a running loop")
    130 except RuntimeError:
    131     pass

SyncError: Calling sync() from within a running loop

[jhamman]

@ilan-gold - @brokkoli71 is working on bringing this functionality back in #2622