Further tests and delete chunk data after resize (#49)

* add tests

* test resize and reopen array

* add resizeMetadataOnly argument for resize

* reformat

* default to parallel read, write

* resize default to parallel and resizeMetadataOnly

* update tests with long

Author: brokkoli71

src/main/java/dev/zarr/zarrjava/core/Array.java

@@ -3,7 +3,6 @@
 import dev.zarr.zarrjava.ZarrException;
 import dev.zarr.zarrjava.core.codec.CodecPipeline;
 import dev.zarr.zarrjava.store.FilesystemStore;
-import dev.zarr.zarrjava.store.Store;
 import dev.zarr.zarrjava.store.StoreHandle;
 import dev.zarr.zarrjava.utils.IndexingUtils;
 import dev.zarr.zarrjava.utils.MultiArrayUtils;
@@ -17,14 +16,12 @@
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.Arrays;
-import java.util.List;
-import java.util.Set;
-import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 public abstract class Array extends AbstractNode {
 
     protected CodecPipeline codecPipeline;
+    public static final boolean DEFAULT_PARALLELISM = true;
 
     protected Array(StoreHandle storeHandle) throws ZarrException {
         super(storeHandle);
@@ -184,6 +181,110 @@ public ucar.ma2.Array readChunk(long[] chunkCoords) throws ZarrException {
         return codecPipeline.decode(chunkBytes);
     }
 
+    /**
+     * Deletes chunks that are completely outside the new shape and trims boundary chunks.
+     *
+     * @param newShape the new shape of the array
+     * @param parallel utilizes parallelism if true
+     */
+    protected void cleanupChunksForResize(long[] newShape, boolean parallel) {
+        ArrayMetadata metadata = metadata();
+        final int[] chunkShape = metadata.chunkShape();
+        final int ndim = metadata.ndim();
+        final dev.zarr.zarrjava.core.chunkkeyencoding.ChunkKeyEncoding chunkKeyEncoding = metadata.chunkKeyEncoding();
+
+        // Calculate max valid chunk coordinates for the new shape
+        long[] newMaxChunkCoords = new long[ndim];
+        for (int i = 0; i < ndim; i++) {
+            newMaxChunkCoords[i] = (newShape[i] + chunkShape[i] - 1) / chunkShape[i];
+        }
+
+        // Iterate over all possible chunk coordinates in the old shape
+        long[][] allOldChunkCoords = IndexingUtils.computeChunkCoords(metadata.shape, chunkShape);
+
+        Stream<long[]> chunkStream = Arrays.stream(allOldChunkCoords);
+        if (parallel) {
+            chunkStream = chunkStream.parallel();
+        }
+
+        chunkStream.forEach(chunkCoords -> {
+            boolean isOutsideBounds = false;
+            boolean isOnBoundary = false;
+
+            for (int dimIdx = 0; dimIdx < ndim; dimIdx++) {
+                if (chunkCoords[dimIdx] >= newMaxChunkCoords[dimIdx]) {
+                    isOutsideBounds = true;
+                    break;
+                }
+                // Check if this chunk is on the boundary (partially outside new shape)
+                long chunkEnd = (chunkCoords[dimIdx] + 1) * chunkShape[dimIdx];
+                if (chunkEnd > newShape[dimIdx]) {
+                    isOnBoundary = true;
+                }
+            }
+
+            String[] chunkKeys = chunkKeyEncoding.encodeChunkKey(chunkCoords);
+            StoreHandle chunkHandle = storeHandle.resolve(chunkKeys);
+
+            if (isOutsideBounds) {
+                // Delete chunk that is completely outside
+                chunkHandle.delete();
+            } else if (isOnBoundary) {
+                // Trim boundary chunk - read, clear out-of-bounds data, write back
+                try {
+                    trimBoundaryChunk(chunkCoords, newShape, chunkShape);
+                } catch (ZarrException e) {
+                    throw new RuntimeException(e);
+                }
+            }
+        });
+    }
+
+    /**
+     * Trims a boundary chunk by reading it, clearing the out-of-bounds portion, and writing it back.
+     *
+     * @param chunkCoords the coordinates of the chunk to trim
+     * @param newShape    the new shape of the array
+     * @param chunkShape  the shape of the chunks
+     * @throws ZarrException if reading or writing the chunk fails
+     */
+    protected void trimBoundaryChunk(long[] chunkCoords, long[] newShape, int[] chunkShape) throws ZarrException {
+        ArrayMetadata metadata = metadata();
+        final int ndim = metadata.ndim();
+
+        // Calculate the valid region within this chunk
+        int[] validShape = new int[ndim];
+        boolean needsTrimming = false;
+        for (int dimIdx = 0; dimIdx < ndim; dimIdx++) {
+            long chunkStart = chunkCoords[dimIdx] * chunkShape[dimIdx];
+            long chunkEnd = chunkStart + chunkShape[dimIdx];
+            if (chunkEnd > newShape[dimIdx]) {
+                validShape[dimIdx] = (int) (newShape[dimIdx] - chunkStart);
+                needsTrimming = true;
+            } else {
+                validShape[dimIdx] = chunkShape[dimIdx];
+            }
+        }
+
+        if (!needsTrimming) {
+            return;
+        }
+
+        // Read the existing chunk
+        ucar.ma2.Array chunkData = readChunk(chunkCoords);
+
+        // Create a new chunk filled with fill value
+        ucar.ma2.Array newChunkData = metadata.allocateFillValueChunk();
+
+        // Copy only the valid region
+        MultiArrayUtils.copyRegion(
+                chunkData, new int[ndim], newChunkData, new int[ndim], validShape
+        );
+
+        // Write the trimmed chunk back
+        writeChunk(chunkCoords, newChunkData);
+    }
+
 
     /**
      * Writes a ucar.ma2.Array into the Zarr array at the beginning of the Zarr array. The shape of
@@ -205,7 +306,7 @@ public void write(ucar.ma2.Array array) {
      * @param array  the data to write
      */
     public void write(long[] offset, ucar.ma2.Array array) {
-        write(offset, array, false);
+        write(offset, array, DEFAULT_PARALLELISM);
     }
 
     /**
@@ -240,7 +341,7 @@ public ucar.ma2.Array read() throws ZarrException {
      */
     @Nonnull
     public ucar.ma2.Array read(final long[] offset, final long[] shape) throws ZarrException {
-        return read(offset, shape, false);
+        return read(offset, shape, DEFAULT_PARALLELISM);
     }
 
     /**
@@ -339,6 +440,46 @@ public ucar.ma2.Array read(final long[] offset, final long[] shape, final boolea
         return outputArray;
     }
 
+    /**
+     * Sets a new shape for the Zarr array. Only the metadata is updated by default.
+     * This method returns a new instance of the Zarr array class and the old instance
+     * becomes invalid.
+     *
+     * @param newShape the new shape of the Zarr array
+     * @throws ZarrException if the new metadata is invalid
+     * @throws IOException   throws IOException if the new metadata cannot be serialized
+     */
+    public Array resize(long[] newShape) throws ZarrException, IOException {
+        return resize(newShape, true);
+    }
+
+    /**
+     * Sets a new shape for the Zarr array. This method returns a new instance of the Zarr array class
+     * and the old instance becomes invalid.
+     *
+     * @param newShape           the new shape of the Zarr array
+     * @param resizeMetadataOnly if true, only the metadata is updated; if false, chunks outside the new
+     *                           bounds are deleted and boundary chunks are trimmed
+     * @throws ZarrException if the new metadata is invalid
+     * @throws IOException   throws IOException if the new metadata cannot be serialized
+     */
+    public Array resize(long[] newShape, boolean resizeMetadataOnly) throws ZarrException, IOException {
+        return resize(newShape, resizeMetadataOnly, DEFAULT_PARALLELISM);
+    }
+
+    /**
+     * Sets a new shape for the Zarr array. This method returns a new instance of the Zarr array class
+     * and the old instance becomes invalid.
+     *
+     * @param newShape           the new shape of the Zarr array
+     * @param resizeMetadataOnly if true, only the metadata is updated; if false, chunks outside the new
+     *                           bounds are deleted and boundary chunks are trimmed
+     * @param parallel           utilizes parallelism if true when cleaning up chunks
+     * @throws ZarrException if the new metadata is invalid
+     * @throws IOException   throws IOException if the new metadata cannot be serialized
+     */
+    public abstract Array resize(long[] newShape, boolean resizeMetadataOnly, boolean parallel) throws ZarrException, IOException;
+
     public ArrayAccessor access() {
         return new ArrayAccessor(this);
     }

src/main/java/dev/zarr/zarrjava/v2/Array.java

@@ -201,26 +201,63 @@ private Array writeMetadata(ArrayMetadata newArrayMetadata) throws ZarrException
     }
 
     /**
-     * Sets a new shape for the Zarr array. It only changes the metadata, no array data is modified or
-     * deleted. This method returns a new instance of the Zarr array class and the old instance
+     * Sets a new shape for the Zarr array. Only the metadata is updated by default.
+     * This method returns a new instance of the Zarr array class and the old instance
      * becomes invalid.
      *
      * @param newShape the new shape of the Zarr array
      * @throws ZarrException if the new metadata is invalid
      * @throws IOException   throws IOException if the new metadata cannot be serialized
      */
+    @Override
     public Array resize(long[] newShape) throws ZarrException, IOException {
+        return resize(newShape, true);
+    }
+
+    /**
+     * Sets a new shape for the Zarr array. This method returns a new instance of the Zarr array class
+     * and the old instance becomes invalid.
+     *
+     * @param newShape           the new shape of the Zarr array
+     * @param resizeMetadataOnly if true, only the metadata is updated; if false, chunks outside the new
+     *                           bounds are deleted and boundary chunks are trimmed
+     * @throws ZarrException if the new metadata is invalid
+     * @throws IOException   throws IOException if the new metadata cannot be serialized
+     */
+    @Override
+    public Array resize(long[] newShape, boolean resizeMetadataOnly) throws ZarrException, IOException {
+        return resize(newShape, resizeMetadataOnly, DEFAULT_PARALLELISM);
+    }
+
+    /**
+     * Sets a new shape for the Zarr array. This method returns a new instance of the Zarr array class
+     * and the old instance becomes invalid.
+     *
+     * @param newShape           the new shape of the Zarr array
+     * @param resizeMetadataOnly if true, only the metadata is updated; if false, chunks outside the new
+     *                           bounds are deleted and boundary chunks are trimmed
+     * @param parallel           utilizes parallelism if true when cleaning up chunks
+     * @throws ZarrException if the new metadata is invalid
+     * @throws IOException   throws IOException if the new metadata cannot be serialized
+     */
+    @Override
+    public Array resize(long[] newShape, boolean resizeMetadataOnly, boolean parallel) throws ZarrException, IOException {
         if (newShape.length != metadata.ndim()) {
             throw new IllegalArgumentException(
                     "'newShape' needs to have rank '" + metadata.ndim() + "'.");
         }
 
+        if (!resizeMetadataOnly) {
+            cleanupChunksForResize(newShape, parallel);
+        }
+
         ArrayMetadata newArrayMetadata = ArrayMetadataBuilder.fromArrayMetadata(metadata)
                 .withShape(newShape)
                 .build();
         return writeMetadata(newArrayMetadata);
     }
 
+
     /**
      * Sets the attributes of the Zarr array. It overwrites and removes any existing attributes. This
      * method returns a new instance of the Zarr array class and the old instance becomes invalid.
@@ -248,7 +285,8 @@ public Array setAttributes(Attributes newAttributes) throws ZarrException, IOExc
      * @throws IOException   throws IOException if the new metadata cannot be serialized
      */
     public Array updateAttributes(Function<Attributes, Attributes> attributeMapper) throws ZarrException, IOException {
-        return setAttributes(attributeMapper.apply(metadata.attributes));
+        Attributes currentAttributes = metadata.attributes != null ? new Attributes(metadata.attributes) : new Attributes();
+        return setAttributes(attributeMapper.apply(currentAttributes));
     }
 
     @Override

src/main/java/dev/zarr/zarrjava/v2/Group.java

@@ -27,7 +27,7 @@
 public class Group extends dev.zarr.zarrjava.core.Group implements Node {
     public GroupMetadata metadata;
 
-    protected Group(@Nonnull StoreHandle storeHandle, @Nonnull GroupMetadata groupMetadata) throws IOException {
+    protected Group(@Nonnull StoreHandle storeHandle, @Nonnull GroupMetadata groupMetadata) {
         super(storeHandle);
         this.metadata = groupMetadata;
     }
@@ -283,7 +283,8 @@ public Group setAttributes(Attributes newAttributes) throws ZarrException, IOExc
      */
     public Group updateAttributes(Function<Attributes, Attributes> attributeMapper)
             throws ZarrException, IOException {
-        return setAttributes(attributeMapper.apply(metadata.attributes));
+        Attributes currentAttributes = metadata.attributes != null ? new Attributes(metadata.attributes) : new Attributes();
+        return setAttributes(attributeMapper.apply(currentAttributes));
     }
 
 

src/main/java/dev/zarr/zarrjava/v3/Array.java

@@ -201,26 +201,63 @@ private Array writeMetadata(ArrayMetadata newArrayMetadata) throws ZarrException
     }
 
     /**
-     * Sets a new shape for the Zarr array. It only changes the metadata, no array data is modified or
-     * deleted. This method returns a new instance of the Zarr array class and the old instance
+     * Sets a new shape for the Zarr array. Only the metadata is updated by default.
+     * This method returns a new instance of the Zarr array class and the old instance
      * becomes invalid.
      *
      * @param newShape the new shape of the Zarr array
      * @throws ZarrException if the new metadata is invalid
      * @throws IOException   throws IOException if the new metadata cannot be serialized
      */
+    @Override
     public Array resize(long[] newShape) throws ZarrException, IOException {
+        return resize(newShape, true);
+    }
+
+    /**
+     * Sets a new shape for the Zarr array. This method returns a new instance of the Zarr array class
+     * and the old instance becomes invalid.
+     *
+     * @param newShape           the new shape of the Zarr array
+     * @param resizeMetadataOnly if true, only the metadata is updated; if false, chunks outside the new
+     *                           bounds are deleted and boundary chunks are trimmed
+     * @throws ZarrException if the new metadata is invalid
+     * @throws IOException   throws IOException if the new metadata cannot be serialized
+     */
+    @Override
+    public Array resize(long[] newShape, boolean resizeMetadataOnly) throws ZarrException, IOException {
+        return resize(newShape, resizeMetadataOnly, DEFAULT_PARALLELISM);
+    }
+
+    /**
+     * Sets a new shape for the Zarr array. This method returns a new instance of the Zarr array class
+     * and the old instance becomes invalid.
+     *
+     * @param newShape           the new shape of the Zarr array
+     * @param resizeMetadataOnly if true, only the metadata is updated; if false, chunks outside the new
+     *                           bounds are deleted and boundary chunks are trimmed
+     * @param parallel           utilizes parallelism if true when cleaning up chunks
+     * @throws ZarrException if the new metadata is invalid
+     * @throws IOException   throws IOException if the new metadata cannot be serialized
+     */
+    @Override
+    public Array resize(long[] newShape, boolean resizeMetadataOnly, boolean parallel) throws ZarrException, IOException {
         if (newShape.length != metadata.ndim()) {
             throw new IllegalArgumentException(
                     "'newShape' needs to have rank '" + metadata.ndim() + "'.");
         }
 
+        if (!resizeMetadataOnly) {
+            cleanupChunksForResize(newShape, parallel);
+        }
+
         ArrayMetadata newArrayMetadata = ArrayMetadataBuilder.fromArrayMetadata(metadata)
                 .withShape(newShape)
                 .build();
         return writeMetadata(newArrayMetadata);
     }
 
+
     /**
      * Sets the attributes of the Zarr array. It overwrites and removes any existing attributes. This
      * method returns a new instance of the Zarr array class and the old instance becomes invalid.
@@ -248,7 +285,8 @@ public Array setAttributes(Attributes newAttributes) throws ZarrException, IOExc
      * @throws IOException   throws IOException if the new metadata cannot be serialized
      */
     public Array updateAttributes(Function<Attributes, Attributes> attributeMapper) throws ZarrException, IOException {
-        return setAttributes(attributeMapper.apply(metadata.attributes));
+        Attributes currentAttributes = metadata.attributes != null ? new Attributes(metadata.attributes) : new Attributes();
+        return setAttributes(attributeMapper.apply(currentAttributes));
     }
 
     @Override

src/main/java/dev/zarr/zarrjava/v3/Group.java

@@ -289,7 +289,8 @@ private Group writeMetadata(GroupMetadata newGroupMetadata) throws IOException {
      * @throws IOException   if the metadata cannot be serialized
      */
     public Group updateAttributes(Function<Attributes, Attributes> attributeMapper) throws ZarrException, IOException {
-        return setAttributes(attributeMapper.apply(metadata.attributes));
+        Attributes currentAttributes = metadata.attributes != null ? new Attributes(metadata.attributes) : new Attributes();
+        return setAttributes(attributeMapper.apply(currentAttributes));
     }
 
     /**