Skip to content
153 changes: 147 additions & 6 deletions src/main/java/dev/zarr/zarrjava/core/Array.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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);
Expand Down Expand Up @@ -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
Expand All @@ -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);
}

/**
Expand Down Expand Up @@ -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);
}

/**
Expand Down Expand Up @@ -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);
}
Expand Down
44 changes: 41 additions & 3 deletions src/main/java/dev/zarr/zarrjava/v2/Array.java
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down
5 changes: 3 additions & 2 deletions src/main/java/dev/zarr/zarrjava/v2/Group.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down Expand Up @@ -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));
}


Expand Down
44 changes: 41 additions & 3 deletions src/main/java/dev/zarr/zarrjava/v3/Array.java
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down
3 changes: 2 additions & 1 deletion src/main/java/dev/zarr/zarrjava/v3/Group.java
Original file line number Diff line number Diff line change
Expand Up @@ -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));
}

/**
Expand Down
Loading
Loading