set_coordinate_selection runs out-of-memory for large arrays

[sehoffmann]

Zarr version

v3.2.1

Numcodecs version

v0.16.5

Python Version

3.12.11

Operating System

Linux

Installation

uv .venv

Description

Cause is the np.bincount call which will allocate a very large array even if we only set a handful of indices:

Traceback (most recent call last):
  File "bug.py", line 29, in <module>
    arr.set_coordinate_selection((rows, cols), vals)  # set_coordinate_selection: MemoryError
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "site-packages/zarr/core/array.py", line 3584, in set_coordinate_selection
    indexer = CoordinateIndexer(selection, self.shape, self._chunk_grid)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "site-packages/zarr/core/indexing.py", line 1252, in __init__
    chunk_nitems = np.bincount(chunks_raveled_indices, minlength=nchunks)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
numpy._core._exceptions._ArrayMemoryError: Unable to allocate 1.01 TiB for an array with shape (138540000000,) and data type int64

It makes scaling O(chunks) as well

Steps to reproduce

# /// script
# requires-python = ">=3.12"
# dependencies = [
#   "zarr@git+https://github.com/zarr-developers/zarr-python.git@main",
# ]
# ///
#
# This script automatically imports the development branch of zarr to check for issues

import zarr
import numpy as np

# chunks = 10_000_000 * 13_854
arr = zarr.create_array(
    zarr.storage.MemoryStore(),
    shape=(10_000_000, 13_854),
    chunks=(1, 1),
    dtype='int32',
    fill_value=0,
)

# we write 3 points
rows = np.array([0, 1, 2])
cols = np.array([5, 5, 5])
vals = np.array([1, 2, 3], dtype='int32')

arr[0, 5] = 1                                     # set_basic_selection: OK
arr.set_orthogonal_selection((rows, 5), vals)     # set_orthogonal_selection: OK
arr.set_coordinate_selection((rows, cols), vals)  # set_coordinate_selection: MemoryError

Additional output

No response

[d-v-b]

thanks for this report, i'm sure there's something more efficient than bincount here. I'll look into getting a fix up shortly

[sehoffmann]

Thanks!

arr.get_coordinate_selection((rows, cols)) is also affected and results in the same issue.