WIP: Rework to use PATCH /array/full

Author: danielballan

tiled/adapters/zarr.py

@@ -174,41 +174,40 @@ async def write_block(
         -------
 
         """
-        if slice is not ...:
-            raise NotImplementedError
         block_slice, shape = slice_and_shape_from_block_and_chunks(
             block, self.structure().chunks
         )
         self._array[block_slice] = data
 
-    async def append_block(
+    async def patch(
         self,
         data: NDArray[Any],
-        axis: int,
-    ) -> List[int]:
-        """
-
-        Parameters
-        ----------
-        data :
-        block :
-        slice :
-
-        Returns
-        -------
-
-        """
-
-        new_shape = list(self._array.shape)
-        new_shape[axis] += list(data.shape)[axis]  # Extend along axis
-
-        # Resize the Zarr array to accommodate new data
-        self._array.resize(tuple(new_shape))
-
-        # Append the new data to the resized array
-        # Slicing to place data at the end
-        self._array[-data.shape[0] :] = data  # noqa: E203
-        return new_shape
+        slice: NDSlice,
+        grow: bool = False,
+    ) -> Tuple[int, ...]:
+        """
+        Write data into a slice of the array, maybe growing it.
+
+        If the specified slice does not fit into the array, and grow=True, the
+        array will be resize (grown, never shrunk) to fit it. The new shape is
+        returned.
+        """
+        current_shape = self._array.shape
+        new_shape = list(current_shape)
+        for i, (s, dim) in enumerate(zip(slice, current_shape)):
+            if isinstance(s, int):
+                new_shape[i] = max(new_shape[i], s)
+            elif isinstance(s, builtins.slice) and isinstance(s.stop, int):
+                new_shape[i] = max(new_shape[i], s.stop)
+        new_shape_tuple = tuple(new_shape)
+        if new_shape_tuple != current_shape:
+            if grow:
+                # Resize the Zarr array to accommodate new data
+                self._array.resize(new_shape_tuple)
+            else:
+                raise ValueError(f"Slice does not fit into array shape {current_shape}")
+        self._array[slice] = data
+        return new_shape_tuple
 
 
 if sys.version_info < (3, 9):

tiled/catalog/adapter.py

@@ -1041,15 +1041,13 @@ async def write_block(self, *args, **kwargs):
             (await self.get_adapter()).write_block, *args, **kwargs
         )
 
-    async def append_block(self, *args, **kwargs):
+    async def patch(self, *args, **kwargs):
         # assumes a single DataSource (currently only supporting zarr)
         async with self.context.session() as db:
             try:
                 new_shape = await ensure_awaitable(
-                    (await self.get_adapter()).append_block, *args, **kwargs
+                    (await self.get_adapter()).patch, *args, **kwargs
                 )
-                if new_shape is None:
-                    raise ValueError("No new shape returned from append_block.")
                 node = await db.get(orm.Node, self.node.id)
                 data_source = node.data_sources[0]
                 structure_row = await db.get(orm.Structure, data_source.structure_id)

tiled/client/array.py

@@ -3,7 +3,9 @@
 import dask
 import dask.array
 import numpy
+from numpy.typing import NDArray
 
+from ..type_aliases import NDSlice
 from .base import BaseClient
 from .utils import export_util, handle_error, params_from_slice
 
@@ -167,40 +169,39 @@ def write(self, array):
             )
         )
 
-    def write_block(self, array, block):
+    def write_block(self, array, block, slice=...):
         handle_error(
             self.context.http_client.put(
                 self.item["links"]["block"].format(*block),
                 content=array.tobytes(),
                 headers={"Content-Type": "application/octet-stream"},
+                params=params_from_slice(slice),
             )
         )
 
-    def append_block(self, array: numpy.nd, axis):
+    def patch(self, array: NDArray, slice: NDSlice, grow=False):
         """
-        Append a block to the array along the given axis. The block must have
-        the same shape as the existing blocks along that axis.
-        This method differs from `write_block` as it increases the size
-        of the array along the given axis, while `write_block` overwrites
-        the data in the block at the given index. This is useful for
-        cases where you do not know ahead of time how many blocks you will
-        eventually receive.
+        Write data
 
         Parameters
         ----------
         array : array-like
-            The block to append.
-        axis : int
-            The axis along which to append the block.
-
-
+            The data to write
+        slice : NDSlice
+            Where to place this data in the array
+        grow : bool
+            Grow the array shape to fit the new slice, if necessary
         """
-        formatted_shape = ",".join(f"{value}" for value in array.shape)
+        array_ = numpy.ascontiguousarray(array)
+        params = params_from_slice(slice)
+        params["shape"] = ",".join(map(str, array_.shape))
+        params["grow"] = bool(grow)
         handle_error(
             self.context.http_client.patch(
-                self.item["links"]["append"].format(formatted_shape, axis),
-                content=array.tobytes(),
+                self.item["links"]["full"],
+                content=array_.tobytes(),
                 headers={"Content-Type": "application/octet-stream"},
+                params=params,
             )
         )
 

tiled/server/dependencies.py

@@ -165,6 +165,13 @@ def expected_shape(
     return tuple(map(int, expected_shape.split(",")))
 
 
+def shape_param(
+    shape: str = Query(..., min_length=1, pattern="^[0-9]+(,[0-9]+)*$|^scalar$"),
+):
+    "Specify and parse a shape parameter."
+    return tuple(map(int, shape.split(",")))
+
+
 def np_style_slicer(indices: tuple):
     return indices[0] if len(indices) == 1 else slice_func(*indices)
 

tiled/server/links.py

@@ -20,7 +20,6 @@ def links_for_array(structure_family, structure, base_url, path_str):
     block_template = ",".join(f"{{{index}}}" for index in range(len(structure.shape)))
     links["block"] = f"{base_url}/array/block/{path_str}?block={block_template}"
     links["full"] = f"{base_url}/array/full/{path_str}"
-    links["append"] = f"{base_url}/array/append/{path_str}?shape={{0}}&axis={{1}}"
     return links
 
 

tiled/server/router.py

@@ -6,7 +6,7 @@
 from datetime import datetime, timedelta
 from functools import partial
 from pathlib import Path
-from typing import Any, List, Optional, Tuple
+from typing import Any, List, Optional
 
 import anyio
 from fastapi import APIRouter, Body, Depends, HTTPException, Query, Request, Security
@@ -55,6 +55,7 @@
     get_query_registry,
     get_serialization_registry,
     get_validation_registry,
+    shape_param,
     slice_,
 )
 from .file_response_with_range import FileResponseWithRange
@@ -1290,39 +1291,32 @@ async def put_array_block(
     return json_or_msgpack(request, None)
 
 
-@router.patch("/array/append/{path:path}")
-async def patch_array_block(
+@router.patch("/array/full/{path:path}")
+async def patch_array_full(
     request: Request,
-    shape: str,
-    axis: int,
+    slice=Depends(slice_),
+    shape=Depends(shape_param),
+    grow: bool = False,
     entry=SecureEntry(
         scopes=["write:data"],
-        structure_families={StructureFamily.array, StructureFamily.sparse},
+        structure_families={StructureFamily.array},
     ),
     deserialization_registry=Depends(get_deserialization_registry),
 ):
-    if not hasattr(entry, "append_block"):
+    if slice is None:
+        slice = ...
+    if not hasattr(entry, "patch"):
         raise HTTPException(
             status_code=HTTP_405_METHOD_NOT_ALLOWED,
             detail="This node cannot accept array data.",
         )
 
     dtype = entry.structure().data_type.to_numpy_dtype()
-    shape_tuple: Tuple[int, ...] = tuple(map(int, shape.split(",")))
     body = await request.body()
     media_type = request.headers["content-type"]
-    if entry.structure_family == "array":
-        # dtype = entry.structure().data_type.to_numpy_dtype()
-        # _, shape = slice_and_shape_from_block_and_chunks(
-        #     block, entry.structure().chunks
-        # )
-        deserializer = deserialization_registry.dispatch("array", media_type)
-        data = await ensure_awaitable(deserializer, body, dtype, shape_tuple)
-    elif entry.structure_family == "sparse":  # TODO: Handle sparse
-        raise NotImplementedError(entry.structure_family)
-    else:
-        raise NotImplementedError(entry.structure_family)
-    await ensure_awaitable(entry.append_block, data, axis)
+    deserializer = deserialization_registry.dispatch("array", media_type)
+    data = await ensure_awaitable(deserializer, body, dtype, shape)
+    await ensure_awaitable(entry.patch, data, slice, grow)
     return json_or_msgpack(request, None)
 
 

tiled/server/schemas.py

@@ -210,7 +210,6 @@ class ArrayLinks(pydantic.BaseModel):
     self: str
     full: str
     block: str
-    append: str
 
 
 class AwkwardLinks(pydantic.BaseModel):