Copy files, if possible (#1362)

* use fs-based copy

* move to Layer.add_mag_as_copy

* deprecations and changelog

* fix tests

* rm test

* pr feedback

Author: normanrz

webknossos/Changelog.md

@@ -19,6 +19,7 @@ For upgrade instructions, please check the respective _Breaking Changes_ section
 ### Added
 
 ### Changed
+- `Layer.add_mag_as_copy` now automatically uses file-based copy, if possible. Therefore, `Dataset.fs_copy_dataset`, `Dataset.add_fs_copy_layer` and `Layer.add_fs_copy_mag` are deprecated. [#1362](https://github.com/scalableminds/webknossos-libs/pull/1362)
 - The API methods now use `v10` API version of Webknossos. Datasets are now referenced by ID instead of dataset name or directory name. [#1363](https://github.com/scalableminds/webknossos-libs/pull/1363)
 
 ### Fixed

webknossos/tests/dataset/test_attachments.py

@@ -190,47 +190,6 @@ def test_copy_layer(tmp_upath: UPath) -> None:
     )
 
 
-def test_fs_copy_layer(tmp_upath: UPath) -> None:
-    dataset, seg_layer = make_dataset(tmp_upath)
-
-    mesh_path = dataset.path / "seg" / "meshes" / "meshfile.hdf5"
-    mesh_path.parent.mkdir(parents=True, exist_ok=True)
-    mesh_path.write_text("test")
-
-    agglomerate_path = (tmp_upath / "agglomerate_view_15").resolve()
-    agglomerate_path.mkdir(parents=True, exist_ok=True)
-    (agglomerate_path / "zarr.json").write_text("test")
-
-    seg_layer.attachments.add_mesh(
-        mesh_path,
-        name="meshfile",
-        data_format=AttachmentDataFormat.HDF5,
-    )
-
-    seg_layer.attachments.add_agglomerate(
-        Path("../agglomerate_view_15"),
-        name="agglomerate_view_15",
-        data_format=AttachmentDataFormat.Zarr3,
-    )
-
-    copy_dataset = Dataset(tmp_upath / "test_copy", voxel_size=(10, 10, 10))
-    copy_layer = copy_dataset.add_fs_copy_layer(seg_layer).as_segmentation_layer()
-
-    # has been copied
-    assert (
-        copy_layer.attachments.meshes[0].path
-        == copy_dataset.path / "seg" / "meshes" / "meshfile.hdf5"
-    )
-    assert (copy_dataset.path / "seg" / "meshes" / "meshfile.hdf5").exists()
-
-    # has not been copied
-    assert copy_layer.attachments.agglomerates[0].path == agglomerate_path
-    assert (
-        copy_layer.attachments.agglomerates[0]._properties.path
-        == agglomerate_path.as_posix()
-    )
-
-
 def test_symlink_layer(tmp_upath: UPath) -> None:
     dataset, seg_layer = make_dataset(tmp_upath)
 

webknossos/tests/dataset/test_dataset.py

@@ -5,6 +5,7 @@
 from collections.abc import Iterator
 from pathlib import Path
 from typing import cast
+from unittest import mock
 
 import numpy as np
 import pytest
@@ -2113,7 +2114,15 @@ def test_add_fs_copy_mag(data_format: DataFormat, output_path: UPath) -> None:
     copy_layer = copy_ds.add_layer(
         "color", COLOR_CATEGORY, dtype_per_channel="uint8", data_format=data_format
     )
-    copy_mag = copy_layer.add_fs_copy_mag(original_mag, extend_layer_bounding_box=True)
+
+    with mock.patch.object(
+        copy_layer, "_add_fs_copy_mag", wraps=copy_layer._add_fs_copy_mag
+    ) as mocked_method:
+        copy_mag = copy_layer.add_mag_as_copy(
+            original_mag, extend_layer_bounding_box=True
+        )
+        mocked_method.assert_called_once()
+
     assert not copy_layer.read_only
     assert not copy_mag.read_only
 
@@ -3533,9 +3542,7 @@ def test_zarr_copy_to_remote_dataset(data_format: DataFormat) -> None:
 
 @pytest.mark.parametrize("input_path", OUTPUT_PATHS)
 @pytest.mark.parametrize("output_path", OUTPUT_PATHS)
-def test_fs_copy_dataset_with_attachments(
-    input_path: UPath, output_path: UPath
-) -> None:
+def test_copy_dataset_with_attachments(input_path: UPath, output_path: UPath) -> None:
     ds_path = copy_simple_dataset(DEFAULT_DATA_FORMAT, input_path)
     new_ds_path = prepare_dataset_path(DEFAULT_DATA_FORMAT, output_path, "copied")
 
@@ -3562,7 +3569,7 @@ def test_fs_copy_dataset_with_attachments(
     )
 
     # Copy
-    copy_ds = ds.fs_copy_dataset(new_ds_path)
+    copy_ds = ds.copy_dataset(new_ds_path)
 
     assert (
         copy_ds.default_view_configuration

webknossos/webknossos/dataset/_array.py

@@ -133,6 +133,13 @@ def from_array_info(cls, array_info: ArrayInfo) -> "Zarr3ArrayInfo":
                 axis_order=array_info.axis_order,
             )
 
+    @property
+    def zarr3_config(self) -> "Zarr3Config":
+        return Zarr3Config(
+            codecs=self.codecs,
+            chunk_key_encoding=self.chunk_key_encoding,
+        )
+
 
 @dataclass(frozen=True)
 class Zarr3Config:

webknossos/webknossos/dataset/dataset.py

@@ -39,7 +39,7 @@
 )
 from ..geometry.mag import MagLike
 from ..geometry.nd_bounding_box import derive_nd_bounding_box_from_shape
-from ._array import ArrayException, ArrayInfo, BaseArray
+from ._array import ArrayException, ArrayInfo, BaseArray, Zarr3Config
 from ._metadata import DatasetMetadata
 from ._utils import pims_images
 from .defaults import (
@@ -2248,7 +2248,7 @@ def add_layer_as_copy(
         shard_shape: Vec3IntLike | int | None = None,
         chunks_per_shard: Vec3IntLike | int | None = None,
         data_format: str | DataFormat | None = None,
-        compress: bool | None = None,
+        compress: bool | Zarr3Config | None = None,
         exists_ok: bool = False,
         executor: Executor | None = None,
         with_attachments: bool = True,
@@ -2544,13 +2544,15 @@ def add_fs_copy_layer(
         foreign_layer: str | Path | Layer,
         new_layer_name: str | None = None,
     ) -> Layer:
-        """
+        """Deprecated. File-based copy is automatically used in `Dataset.add_layer_as_copy`.
+
         Copies the files at `foreign_layer` which belongs to another dataset
         to the current dataset via the filesystem. Additionally, the relevant
         information from the `datasource-properties.json` of the other dataset
         are copied too. If new_layer_name is None, the name of the foreign
         layer is used.
         """
+        warn_deprecated("add_fs_copy_layer", "add_layer_as_copy")
         self._ensure_writable()
         foreign_layer = Layer._ensure_layer(foreign_layer)
 
@@ -2725,7 +2727,8 @@ def fs_copy_dataset(
         exists_ok: bool = False,
         layers_to_ignore: Iterable[str] | None = None,
     ) -> "Dataset":
-        """
+        """Deprecated. File-based copy is automatically used by `Dataset.copy_dataset`.
+
         Creates an independent copy of the dataset with all layers at a new location.
 
         This method copies the files of the dataset as is and, therefore, might be faster than Dataset.copy_dataset, which decodes and encodes all the data.
@@ -2751,12 +2754,13 @@ def fs_copy_dataset(
         Note:
             WKW layers can only be copied to datasets on local file systems.
         """
+        warn_deprecated("fs_copy_dataset", "copy_dataset")
 
         new_dataset_path = UPath(new_dataset_path)
 
         if any(layer.data_format == DataFormat.WKW for layer in self.layers.values()):
             assert is_fs_path(new_dataset_path), (
-                "Cannot create WKW layers in remote datasets. Use explicit `data_format='zarr3'`."
+                "Cannot create WKW layers in remote datasets. Use `Dataset.copy_dataset` with `data_format='zarr3'`."
             )
 
         new_dataset = Dataset(

webknossos/webknossos/dataset/layer.py

@@ -2,6 +2,7 @@
 import operator
 import re
 import warnings
+from inspect import getframeinfo, stack
 from os import PathLike
 from os.path import relpath
 from pathlib import Path
@@ -14,7 +15,7 @@
 
 from ..geometry import Mag, NDBoundingBox, Vec3Int, Vec3IntLike
 from ..geometry.mag import MagLike
-from ._array import ArrayException, TensorStoreArray, Zarr3Config
+from ._array import ArrayException, TensorStoreArray, Zarr3ArrayInfo, Zarr3Config
 from ._downsampling_utils import (
     calculate_default_coarsest_mag,
     calculate_mags_to_downsample,
@@ -60,6 +61,8 @@
     DEFAULT_SHARD_SHAPE,
 )
 
+logger = logging.getLogger(__name__)
+
 
 def _is_int(s: str) -> bool:
     try:
@@ -870,6 +873,43 @@ def add_mag_as_copy(
         self._ensure_writable()
         foreign_mag_view = MagView._ensure_mag_view(foreign_mag_view_or_path)
 
+        has_same_shapes = (
+            (
+                chunk_shape is None
+                or Vec3Int.from_vec_or_int(chunk_shape)
+                == foreign_mag_view.info.chunk_shape
+            )
+            and (
+                shard_shape is None
+                or Vec3Int.from_vec_or_int(shard_shape)
+                == foreign_mag_view.info.shard_shape
+            )
+            and (
+                chunks_per_shard is None
+                or Vec3Int.from_vec_or_int(chunks_per_shard)
+                == foreign_mag_view.info.chunks_per_shard
+            )
+        )
+        has_same_format = self.data_format == foreign_mag_view.info.data_format
+        if compress is None:
+            has_same_compression = True
+        elif isinstance(compress, Zarr3Config) and isinstance(
+            foreign_mag_view.info, Zarr3ArrayInfo
+        ):
+            has_same_compression = compress == foreign_mag_view.info.zarr3_config
+        else:
+            has_same_compression = compress == foreign_mag_view.info.compression_mode
+
+        if has_same_shapes and has_same_format and has_same_compression:
+            logger.debug(
+                f"Optimization: Copying files from {foreign_mag_view.path} to {self.path}/{foreign_mag_view.mag} directly without re-encoding."
+            )
+            return self._add_fs_copy_mag(
+                foreign_mag_view,
+                extend_layer_bounding_box=extend_layer_bounding_box,
+                exists_ok=exists_ok,
+            )
+
         chunk_shape = Vec3Int.from_vec_or_int(
             chunk_shape or foreign_mag_view.info.chunk_shape
         )
@@ -1036,21 +1076,26 @@ def add_mag_as_ref(
 
         return self.get_mag(foreign_mag_view.mag)
 
-    def add_fs_copy_mag(
+    def _add_fs_copy_mag(
         self,
         foreign_mag_view_or_path: PathLike | str | MagView,
         *,
         extend_layer_bounding_box: bool = True,
+        exists_ok: bool = False,
     ) -> MagView:
-        """
-        Copies the data at `foreign_mag_view_or_path` which belongs to another dataset to the current dataset via the filesystem.
-        Additionally, the relevant information from the `datasource-properties.json` of the other dataset are copied, too.
-        """
         self._ensure_writable()
         foreign_mag_view = MagView._ensure_mag_view(foreign_mag_view_or_path)
         self._assert_mag_does_not_exist_yet(foreign_mag_view.mag)
 
+        assert foreign_mag_view.info.data_format == self.data_format, (
+            f"Cannot use file-based copy, because the foreign data format {foreign_mag_view.info.data_format} does not match the layer's data format {self.data_format}."
+        )
+
         mag_path = self.path / str(foreign_mag_view.mag)
+        if not exists_ok and mag_path.exists():
+            raise FileExistsError(
+                f"Cannot copy {foreign_mag_view.path} to {mag_path} because it already exists."
+            )
         copytree(
             foreign_mag_view.path,
             mag_path,
@@ -1067,6 +1112,30 @@ def add_fs_copy_mag(
 
         return mag
 
+    def add_fs_copy_mag(
+        self,
+        foreign_mag_view_or_path: PathLike | str | MagView,
+        *,
+        extend_layer_bounding_box: bool = True,
+        exists_ok: bool = False,
+    ) -> MagView:
+        """Deprecated. File-copy is automatically selected when using `Layer.add_mag_as_copy`.
+
+        Copies the data at `foreign_mag_view_or_path` which belongs to another dataset to the current dataset via the filesystem.
+        Additionally, the relevant information from the `datasource-properties.json` of the other dataset are copied, too.
+        """
+        caller = getframeinfo(stack()[2][0])
+        warnings.warn(
+            f"[DEPRECATION] Direct use of `Layer.add_fs_copy_mag` is deprecated, please use `Layer.add_mag_as_copy` instead (see {caller.filename}:{caller.lineno}), which automatically uses file-based copy if available.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return self._add_fs_copy_mag(
+            foreign_mag_view_or_path,
+            extend_layer_bounding_box=extend_layer_bounding_box,
+            exists_ok=exists_ok,
+        )
+
     def add_mag_from_zarrarray(
         self,
         mag: MagLike,