Fix buffered slice writer overwriting data under certain conditions (#973)

* fix that buffered slice writer could overwrite data when writing less slices than buffer_size at unaligned offset

* format

* update changelog

Author: philippotto

webknossos/Changelog.md

@@ -19,7 +19,7 @@ For upgrade instructions, please check the respective _Breaking Changes_ section
 ### Changed
 
 ### Fixed
-
+- Fixes that the buffered slice writer could overwrite data when writing less slices than buffer_size at an offset that is not aligned. [#973](https://github.com/scalableminds/webknossos-libs/pull/973)
 
 ## [0.14.11](https://github.com/scalableminds/webknossos-libs/releases/tag/v0.14.11) - 2023-12-06
 [Commits](https://github.com/scalableminds/webknossos-libs/compare/v0.14.10...v0.14.11)

webknossos/tests/dataset/test_buffered_slice_utils.py

@@ -164,6 +164,49 @@ def test_basic_buffered_slice_writer(tmp_path: Path) -> None:
     assert np.all(data == written_data)
 
 
+def test_buffered_slice_writer_unaligned(
+    tmp_path: Path,
+) -> None:
+    # Create DS
+    dataset = Dataset(tmp_path, voxel_size=(1, 1, 1))
+    layer = dataset.add_layer(
+        layer_name="color", category="color", dtype_per_channel="uint8", num_channels=1
+    )
+    mag1 = layer.add_mag("1", chunk_shape=(32, 32, 32), chunks_per_shard=(8, 8, 8))
+
+    # Write some data to z=32. We will check that this
+    # data is left untouched by the buffered slice writer.
+    ones_at_z32 = np.ones((512, 512, 4), dtype=np.uint8)
+    ones_offset = (0, 0, 32)
+    mag1.write(ones_at_z32, absolute_offset=ones_offset)
+
+    # Allocate some data (~ 8 MB). Note that this will write
+    # from z=1 to z=31 (i.e., 31 slices instead of 32 which
+    # is the buffer_size with which we configure the BufferedSliceWriter).
+    offset = (1, 1, 1)
+    shape = (512, 512, 31)
+    data = np.random.randint(0, 255, shape, dtype=np.uint8)
+
+    with warnings.catch_warnings():
+        warnings.filterwarnings("default", module="webknossos", message=r"\[WARNING\]")
+        with mag1.get_buffered_slice_writer(
+            absolute_offset=offset, buffer_size=32
+        ) as writer:
+            for z in range(0, shape[2]):
+                section = data[:, :, z]
+                writer.send(section)
+
+    written_data = mag1.read(absolute_offset=offset, size=shape)
+    assert np.all(
+        data == written_data
+    ), "Read data is not equal to the data that was just written."
+
+    data_at_z32 = mag1.read(absolute_offset=ones_offset, size=ones_at_z32.shape)
+    assert np.all(
+        ones_at_z32 == data_at_z32
+    ), "The BufferedSliceWriter seems to have overwritten older data."
+
+
 def test_buffered_slice_writer_should_warn_about_unaligned_usage(
     tmp_path: Path,
 ) -> None:

webknossos/webknossos/dataset/_utils/buffered_slice_writer.py

@@ -130,23 +130,22 @@ def _flush_buffer(self) -> None:
             max_height = max(section.shape[-1] for section in self.slices_to_write)
             channel_count = self.slices_to_write[0].shape[0]
 
-            buffer_bbox = BoundingBox(
-                (0, 0, 0), (max_width, max_height, self.buffer_size)
-            )
+            buffer_depth = min(self.buffer_size, len(self.slices_to_write))
+            buffer_bbox = BoundingBox((0, 0, 0), (max_width, max_height, buffer_depth))
 
             shard_dimensions = self.view._get_file_dimensions().moveaxis(
                 -1, self.dimension
             )
             chunk_size = Vec3Int(
                 min(shard_dimensions[0], max_width),
                 min(shard_dimensions[1], max_height),
-                self.buffer_size,
+                buffer_depth,
             )
             for chunk_bbox in buffer_bbox.chunk(chunk_size):
                 info(f"Writing chunk {chunk_bbox}")
-                width, height, _ = chunk_bbox.size
+                width, height, depth = chunk_bbox.size
                 data = np.zeros(
-                    (channel_count, width, height, self.buffer_size),
+                    (channel_count, width, height, depth),
                     dtype=self.slices_to_write[0].dtype,
                 )