Iterative Shard Writing ala TensorStore Transactions / TensorStore `chunk_layout.write_chunk`

[srivarra]

Describe the new feature you'd like

Hi, I saw a comment comment on the ossci zulip by @mkitti where he raised the idea for creating a simplified version of TensorStore's transactions specifically with shards:

with shard.write_context() as shard_write:
    for chunk in shard_write.inner_chunks():
       shard_write[chunk] = calculation(chunk)

I've found TensorStore's transactions and iterating over their chunk_layout useful for memory bounded workloads. For example when downsampling datasets.

The following plot displays downsampling a fish of shape TCZYX (24, 2, 42, ~3.8K, ~13.1K) where we vary the "batch size" with ts.Transaction.

Here's some sample code I've used for this in TensorStore:

In this case, the batch would be the number of chunks we write at once in a single transaction.

downsampled = ts.downsample(
    source_ts, downsample_factors=downsample_factors, method=method
)

step = target_ts.chunk_layout.write_chunk.shape[0]

for start in range(0, downsampled.shape[0], step):
    with ts.Transaction() as txn:
        target_with_txn = target_ts.with_transaction(txn)
        downsampled_with_txn = downsampled.with_transaction(txn)
        stop = min(start + step, downsampled.shape[0])
        target_with_txn[start:stop].write(downsampled_with_txn[start:stop]).result()

As an aside I find this notation, and all of the result() calls, and indexing and ranging to be repulsive, but oh well. Maybe if there's a nice way to handle that under the hood within the context manager that'd be very cool.

---

I see there is this "regular grid" for chunk grids, but I have no idea how to use this or if it's even iterable like ts.chunk_layout.write_chunk or related to this concept.

Let me know your thoughts, maybe there's stuff already in the zarr-python codebase where I can do something like this manually, tyty.

✨Le Fishe✨

[mkitti]

I think the basic idea is there. You do a bunch of operarions with the write context, but no I/O happens until the context closes. When the context closes there is an opportunity to coalesce I/O operations. For example, write the entire shard all at once rather than individual write operations for individual chunks.

Just a few asides:

1. If I'm iterating in tensorstore, I might use an IndexDomain. It would be nice to have chunk domain iterator here as well.

2. The alternative to using .result() is to use await and write asynchronous functions. zarr-python also has this facility. That said I do not like how asynchronous functions are colored in Python in that they cannot be easily mixed with non-asynchronous functions.

[mkitti]

I think the basic idea is there. You do a bunch of operarions with the write context, but no I/O happens until the context closes. When the context closes there is an opportunity to coalesce I/O operations. For example, write the entire shard all at once rather than individual write operations for individual chunks.

Just a few asides:

1. If I'm iterating in tensorstore, I might use an IndexDomain. It would be nice to have chunk domain iterator here as well.

2. The alternative to using .result() is to use await and write asynchronous functions. zarr-python also has this facility. That said I do not like how asynchronous functions are colored in Python in that they cannot be easily mixed with non-asynchronous functions.

[d-v-b]

i don't really understand why you need transactions here? aren't you just iterating over chunk or shard-sized regions? We do have some functions for this kind of iteration. They are currently private but I think there's interest in making these public:

zarr-python/src/zarr/core/array.py

Lines 5340 to 5494 in 94d543c

	def _iter_chunk_coords(
	array: AnyArray | AnyAsyncArray,
	*,
	origin: Sequence[int] | None = None,
	selection_shape: Sequence[int] | None = None,
	) -> Iterator[tuple[int, ...]]:
	"""
	Create an iterator over the coordinates of chunks in chunk grid space. If the `origin`
	keyword is used, iteration will start at the chunk index specified by `origin`.
	The default behavior is to start at the origin of the grid coordinate space.
	If the `selection_shape` keyword is used, iteration will be bounded over a contiguous region
	ranging from `[origin, origin selection_shape]`, where the upper bound is exclusive as
	per python indexing conventions.
	Parameters
	----------
	array : Array | AsyncArray
	The array to iterate over.
	origin : Sequence[int] | None, default=None
	The origin of the selection in grid coordinates.
	selection_shape : Sequence[int] | None, default=None
	The shape of the selection in grid coordinates.
	Yields
	------
	chunk_coords: tuple[int, ...]
	The coordinates of each chunk in the selection.
	"""
	return _iter_grid(array._chunk_grid_shape, origin=origin, selection_shape=selection_shape)
	def _iter_shard_coords(
	array: AnyArray | AnyAsyncArray,
	*,
	origin: Sequence[int] | None = None,
	selection_shape: Sequence[int] | None = None,
	) -> Iterator[tuple[int, ...]]:
	"""
	Create an iterator over the coordinates of shards in shard grid space. If the `origin`
	keyword is used, iteration will start at the shard index specified by `origin`.
	The default behavior is to start at the origin of the grid coordinate space.
	If the `selection_shape` keyword is used, iteration will be bounded over a contiguous region
	ranging from `[origin, origin selection_shape]`, where the upper bound is exclusive as
	per python indexing conventions.
	Parameters
	----------
	array : Array | AsyncArray
	The array to iterate over.
	origin : Sequence[int] | None, default=None
	The origin of the selection in grid coordinates.
	selection_shape : Sequence[int] | None, default=None
	The shape of the selection in grid coordinates.
	Yields
	------
	chunk_coords: tuple[int, ...]
	The coordinates of each shard in the selection.
	"""
	return _iter_grid(array._shard_grid_shape, origin=origin, selection_shape=selection_shape)
	def _iter_shard_keys(
	array: AnyArray | AnyAsyncArray,
	*,
	origin: Sequence[int] | None = None,
	selection_shape: Sequence[int] | None = None,
	) -> Iterator[str]:
	"""
	Iterate over the storage keys of each shard, relative to an optional origin, and optionally
	limited to a contiguous region in shard grid coordinates.
	Parameters
	----------
	array : Array | AsyncArray
	The array to iterate over.
	origin : Sequence[int] | None, default=None
	The origin of the selection in grid coordinates.
	selection_shape : Sequence[int] | None, default=None
	The shape of the selection in grid coordinates.
	Yields
	------
	key: str
	The storage key of each chunk in the selection.
	"""
	# Iterate over the coordinates of chunks in chunk grid space.
	_iter = _iter_grid(array._shard_grid_shape, origin=origin, selection_shape=selection_shape)
	return (array.metadata.encode_chunk_key(k) for k in _iter)
	def _iter_shard_regions(
	array: AnyArray | AnyAsyncArray,
	*,
	origin: Sequence[int] | None = None,
	selection_shape: Sequence[int] | None = None,
	) -> Iterator[tuple[slice, ...]]:
	"""
	Iterate over the regions spanned by each shard.
	These are the smallest regions of the array that are safe to write concurrently.
	Parameters
	----------
	array : Array | AsyncArray
	The array to iterate over.
	origin : Sequence[int] | None, default=None
	The origin of the selection relative to the array's shard grid.
	selection_shape : Sequence[int] | None, default=None
	The shape of the selection in shard grid coordinates.
	Yields
	------
	region: tuple[slice, ...]
	A tuple of slice objects representing the region spanned by each shard in the selection.
	"""
	if array.shards is None:
	shard_shape = array.chunks
	else:
	shard_shape = array.shards
	return _iter_regions(
	array.shape, shard_shape, origin=origin, selection_shape=selection_shape, trim_excess=True
	)
	def _iter_chunk_regions(
	array: AnyArray | AnyAsyncArray,
	*,
	origin: Sequence[int] | None = None,
	selection_shape: Sequence[int] | None = None,
	) -> Iterator[tuple[slice, ...]]:
	"""
	Iterate over the regions spanned by each shard.
	These are the smallest regions of the array that are efficient to read concurrently.
	Parameters
	----------
	array : Array | AsyncArray
	The array to iterate over.
	origin : Sequence[int] | None, default=None
	The origin of the selection in grid coordinates.
	selection_shape : Sequence[int] | None, default=None
	The shape of the selection in grid coordinates.
	Returns
	-------
	region: tuple[slice, ...]
	A tuple of slice objects representing the region spanned by each shard in the selection.
	"""
	return _iter_regions(
	array.shape, array.chunks, origin=origin, selection_shape=selection_shape, trim_excess=True
	)

[mkitti]

i don't really understand why you need transactions here? aren't you just iterating over chunk or shard-sized regions?

Let's say you want to iterate over 32^3 shaped chunks but have 1024^3 shaped shards. You do want to write the shard until you have processed all 32,786 chunks.

What is the API to not write a shard until you finished processing all the chunks?