From c4a7956e28eb0603475c5c04d03b071bd534559d Mon Sep 17 00:00:00 2001 From: Heberto Mayorquin Date: Wed, 4 Nov 2026 10:39:51 -0600 Subject: [PATCH 1/3] Compute zarr shards one chunk at a time to reduce peak memory --- src/spikeinterface/core/time_series_tools.py | 28 ++++++++++++++------ 1 file changed, 20 insertions(+), 8 deletions(-) diff --git a/src/spikeinterface/core/time_series_tools.py b/src/spikeinterface/core/time_series_tools.py index 4c4abc3492..2381377711 100644 --- a/src/spikeinterface/core/time_series_tools.py +++ b/src/spikeinterface/core/time_series_tools.py @@ -448,14 +448,26 @@ def _write_zarr_chunk(segment_index, start_frame, end_frame, worker_ctx): else: zarr_timestamps_dataset = None - # apply function - data = time_series.get_data( - start_frame=start_frame, - end_frame=end_frame, - segment_index=segment_index, - ) - data = data.astype(dtype) - zarr_dataset[start_frame:end_frame, :] = data + # With shards, a job covers a whole shard so only one worker writes to it. The data is computed one + # zarr chunk at a time to keep the memory of the chain small, and written as many whole chunks at a time + # as fit in write_buffer_megabytes because zarr reads and rewrites the whole shard on every write. + write_buffer_megabytes = 500 + num_samples_per_read = zarr_dataset.chunks[0] + num_values_per_sample = int(np.prod(zarr_dataset.shape[1:])) + bytes_per_read = num_samples_per_read * num_values_per_sample * dtype.itemsize + num_reads_per_write = max(1, write_buffer_megabytes * 1_000_000 // bytes_per_read) + num_samples_per_write = num_reads_per_write * num_samples_per_read + for write_start in range(start_frame, end_frame, num_samples_per_write): + write_end = min(write_start + num_samples_per_write, end_frame) + data = np.empty((write_end - write_start,) + zarr_dataset.shape[1:], dtype=dtype) + for read_start in range(write_start, write_end, num_samples_per_read): + read_end = min(read_start + num_samples_per_read, write_end) + data[read_start - write_start : read_end - write_start] = time_series.get_data( + start_frame=read_start, + end_frame=read_end, + segment_index=segment_index, + ) + zarr_dataset[write_start:write_end, :] = data if zarr_timestamps_dataset is not None: timestamps = time_series.get_times(start_frame=start_frame, end_frame=end_frame, segment_index=segment_index) From 325c837bd580bb00afe559c45f90424b2059e377 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Mon, 9 Nov 2026 11:19:27 +0100 Subject: [PATCH 2/3] feat: expose buffer_megabytes argument in zarr save --- src/spikeinterface/core/baserecording.py | 3 ++- src/spikeinterface/core/recording_tools.py | 4 ++++ src/spikeinterface/core/time_series_tools.py | 14 +++++++++----- src/spikeinterface/core/zarrextractors.py | 2 ++ 4 files changed, 17 insertions(+), 6 deletions(-) diff --git a/src/spikeinterface/core/baserecording.py b/src/spikeinterface/core/baserecording.py index d4bb36f2ab..5ce43b64cf 100644 --- a/src/spikeinterface/core/baserecording.py +++ b/src/spikeinterface/core/baserecording.py @@ -876,7 +876,8 @@ def save(self, format="binary", verbose: bool = False, **save_kwargs): - shards: tuple | None, default: None Number of shard size. If None, no sharding is done. Note that shards dimensions need to be larger than chunk dimensions (if chunks is not None) and that sharding is only done on the first dimension. - - compressors: BytesBytesCodec | ArrayBytesCodec | list of them, or None, default: None + - buffer_megabytes: int, default: 500 + Buffer size in megabytes for writing data to zarr. - compressors: BytesBytesCodec | ArrayBytesCodec | list of them, or None, default: None Global compressors. Only one ArrayBytesCodec (serializer) is supported. If None, Blosc-zstd, level 5, with bit shuffle is used. - filters: ArrayArrayCodec | list[ArrayArrayCodec] or None, default: None diff --git a/src/spikeinterface/core/recording_tools.py b/src/spikeinterface/core/recording_tools.py index 1bce36e05d..fb37cb44b4 100644 --- a/src/spikeinterface/core/recording_tools.py +++ b/src/spikeinterface/core/recording_tools.py @@ -122,6 +122,7 @@ def write_recording_to_zarr( dataset_timestamps_paths=None, extra_chunks=None, dtype=None, + buffer_megabytes=500, compressor_data=None, filters_data=None, compressor_times=None, @@ -147,6 +148,8 @@ def write_recording_to_zarr( controlled by the job_kwargs. Useful to chunk by channel, with ``extra_chunks=(channel_chunk_size,)``. dtype : dtype, default: None Type of the saved data. + buffer_megabytes : int, default: 500 + Buffer size in megabytes for writing data to zarr. compressor_data : zarr compressor or None, default: None Zarr compressor for data. filters_data : list, default: None @@ -166,6 +169,7 @@ def write_recording_to_zarr( dataset_timestamps_paths=dataset_timestamps_paths, extra_chunks=extra_chunks, dtype=dtype, + buffer_megabytes=buffer_megabytes, compressor_data=compressor_data, filters_data=filters_data, compressor_times=compressor_times, diff --git a/src/spikeinterface/core/time_series_tools.py b/src/spikeinterface/core/time_series_tools.py index 2381377711..d31424ecde 100644 --- a/src/spikeinterface/core/time_series_tools.py +++ b/src/spikeinterface/core/time_series_tools.py @@ -299,6 +299,7 @@ def _write_time_series_to_zarr( dtype=None, chunks=None, shards=None, + buffer_megabytes=500, compressor_data=None, filters_data=None, compressor_times=None, @@ -320,12 +321,14 @@ def _write_time_series_to_zarr( List of paths to traces datasets in the zarr group dataset_timestamps_paths : list or None, default: None List of paths to timestamps datasets in the zarr group. If None, timestamps are not saved. + dtype : dtype, default: None + Type of the saved data chunks : tuple or None, default: None Chunking dimensions to use for the zarr dataset. shards : tuple or None, default: None Sharding configuration for the zarr dataset. - dtype : dtype, default: None - Type of the saved data + buffer_megabytes : int, default: 500 + Buffer size in megabytes for writing data to zarr compressor_data : zarr compressor or None, default: None Zarr compressor for data filters_data : list, default: None @@ -403,7 +406,7 @@ def _write_time_series_to_zarr( # use executor (loop or workers) func = _write_zarr_chunk init_func = _init_zarr_worker - init_args = (time_series, zarr_datasets, dtype, zarr_timestamps_datasets) + init_args = (time_series, zarr_datasets, dtype, buffer_megabytes, zarr_timestamps_datasets) executor = TimeSeriesChunkExecutor( time_series, func, init_func, init_args, verbose=verbose, job_name="write_zarr", **job_kwargs ) @@ -420,7 +423,7 @@ def _write_time_series_to_zarr( zarr_group.create_array(name="t_starts", data=t_starts, compressors=None) -def _init_zarr_worker(time_series, zarr_datasets, dtype, zarr_timestamps_datasets=None): +def _init_zarr_worker(time_series, zarr_datasets, dtype, buffer_megabytes, zarr_timestamps_datasets=None): import zarr # create a local dict per worker @@ -432,6 +435,7 @@ def _init_zarr_worker(time_series, zarr_datasets, dtype, zarr_timestamps_dataset else: worker_ctx["zarr_timestamps_datasets"] = None worker_ctx["dtype"] = np.dtype(dtype) + worker_ctx["buffer_megabytes"] = buffer_megabytes return worker_ctx @@ -447,11 +451,11 @@ def _write_zarr_chunk(segment_index, start_frame, end_frame, worker_ctx): zarr_timestamps_dataset = worker_ctx["zarr_timestamps_datasets"][segment_index] else: zarr_timestamps_dataset = None + write_buffer_megabytes = worker_ctx["buffer_megabytes"] # With shards, a job covers a whole shard so only one worker writes to it. The data is computed one # zarr chunk at a time to keep the memory of the chain small, and written as many whole chunks at a time # as fit in write_buffer_megabytes because zarr reads and rewrites the whole shard on every write. - write_buffer_megabytes = 500 num_samples_per_read = zarr_dataset.chunks[0] num_values_per_sample = int(np.prod(zarr_dataset.shape[1:])) bytes_per_read = num_samples_per_read * num_values_per_sample * dtype.itemsize diff --git a/src/spikeinterface/core/zarrextractors.py b/src/spikeinterface/core/zarrextractors.py index bfe70691ae..5b0cf5db16 100644 --- a/src/spikeinterface/core/zarrextractors.py +++ b/src/spikeinterface/core/zarrextractors.py @@ -977,6 +977,7 @@ def add_recording_to_zarr_group( job_kwargs=job_kwargs, time_series=recording, ) + buffer_megabytes = zarr_kwargs.get("buffer_megabytes", 500) _write_time_series_to_zarr( time_series=recording, @@ -991,6 +992,7 @@ def add_recording_to_zarr_group( compressor_times=compressor_times, filters_times=filters_times, verbose=verbose, + buffer_megabytes=buffer_megabytes, **job_kwargs, ) From c3e5b7a8aa769ac46945d683c595fde7ee0ebcc2 Mon Sep 17 00:00:00 2001 From: Alessio Buccino Date: Mon, 9 Nov 2026 11:20:08 +0100 Subject: [PATCH 3/3] feat: expose buffer_megabytes argument in zarr save 2 --- src/spikeinterface/core/baserecording.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/spikeinterface/core/baserecording.py b/src/spikeinterface/core/baserecording.py index 5ce43b64cf..c4f689e9c6 100644 --- a/src/spikeinterface/core/baserecording.py +++ b/src/spikeinterface/core/baserecording.py @@ -877,7 +877,8 @@ def save(self, format="binary", verbose: bool = False, **save_kwargs): Number of shard size. If None, no sharding is done. Note that shards dimensions need to be larger than chunk dimensions (if chunks is not None) and that sharding is only done on the first dimension. - buffer_megabytes: int, default: 500 - Buffer size in megabytes for writing data to zarr. - compressors: BytesBytesCodec | ArrayBytesCodec | list of them, or None, default: None + Buffer size in megabytes for writing data to zarr. + - compressors: BytesBytesCodec | ArrayBytesCodec | list of them, or None, default: None Global compressors. Only one ArrayBytesCodec (serializer) is supported. If None, Blosc-zstd, level 5, with bit shuffle is used. - filters: ArrayArrayCodec | list[ArrayArrayCodec] or None, default: None