diff --git a/src/spikeinterface/core/baserecording.py b/src/spikeinterface/core/baserecording.py index d4bb36f2ab..c4f689e9c6 100644 --- a/src/spikeinterface/core/baserecording.py +++ b/src/spikeinterface/core/baserecording.py @@ -876,6 +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. + - 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. 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 4c4abc3492..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,15 +451,27 @@ 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 - - # 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 + 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. + 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) 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, )