diff --git a/src/main/java/dev/zarr/zarrjava/core/Array.java b/src/main/java/dev/zarr/zarrjava/core/Array.java
index 67461483..0b309c45 100644
--- a/src/main/java/dev/zarr/zarrjava/core/Array.java
+++ b/src/main/java/dev/zarr/zarrjava/core/Array.java
@@ -155,6 +155,45 @@ public void writeChunk(long[] chunkCoords, ucar.ma2.Array chunkArray) throws Zar
}
}
+ /**
+ * Writes already-encoded bytes for one chunk directly into the store, bypassing the codec
+ * pipeline entirely. No decoding or encoding is performed; the supplied bytes are stored
+ * verbatim.
+ *
+ * This is the encoded-bytes counterpart of {@link #writeChunk(long[], ucar.ma2.Array)}. It is
+ * useful when the data is already in exactly the form this array's codec pipeline would produce
+ * (e.g. ingesting JPEG data into an array that uses the {@code jpeg} codec), where a
+ * decode+encode round-trip would waste compute and, for lossy codecs, degrade quality.
+ *
+ * Unsafe: because the bytes are stored without decoding, this method cannot verify them.
+ * The caller is responsible for guaranteeing that {@code chunkBytes} is a whole chunk encoded in
+ * exactly the format this array expects (matching data type, chunk shape and codec
+ * configuration). Supplying incompatible bytes will silently corrupt the array.
+ *
+ * For a sharded array this operates at the shard level: {@code chunkCoords} addresses a whole
+ * shard and {@code chunkBytes} must be the fully-assembled shard (index and all inner chunks).
+ *
+ * @param chunkCoords The coordinates of the chunk as computed by the offset of the chunk divided
+ * by the chunk shape.
+ * @param chunkBytes The already-encoded bytes to store, or {@code null} to delete the chunk
+ * (a subsequent read then returns the fill value).
+ * @throws ZarrException throws ZarrException if the requested chunk is outside the array's domain
+ */
+ public void writeChunkDirect(long[] chunkCoords, @Nullable ByteBuffer chunkBytes) throws ZarrException {
+ if (!chunkIsInArray(chunkCoords)) {
+ throw new ZarrException("Attempting to write data outside of the array's domain.");
+ }
+ ArrayMetadata metadata = metadata();
+ String[] chunkKeys = metadata.chunkKeyEncoding().encodeChunkKey(chunkCoords);
+ StoreHandle chunkHandle = storeHandle.resolve(chunkKeys);
+
+ if (chunkBytes == null) {
+ chunkHandle.delete();
+ } else {
+ chunkHandle.set(chunkBytes);
+ }
+ }
+
/**
* Reads one chunk of the Zarr array as specified by the chunk coordinates into an
* ucar.ma2.Array.
@@ -181,6 +220,148 @@ public ucar.ma2.Array readChunk(long[] chunkCoords) throws ZarrException {
return codecPipeline.decode(chunkBytes);
}
+ /**
+ * Reads the already-encoded bytes of one chunk directly from the store, bypassing the codec
+ * pipeline entirely. No decoding is performed; the raw stored bytes are returned as-is.
+ *
+ * This is the encoded-bytes counterpart of {@link #readChunk(long[])}. It is useful for copying
+ * chunks out of an array in their encoded form (e.g. extracting JPEG data from an array that
+ * uses the {@code jpeg} codec) without a decode+encode round-trip.
+ *
+ * For a sharded array this operates at the shard level: {@code chunkCoords} addresses a whole
+ * shard and the returned bytes are the fully-assembled shard (index and all inner chunks). To
+ * get at a single inner chunk of a shard instead, use {@link #readInnerChunkDirect(long[])}.
+ *
+ * @param chunkCoords The coordinates of the chunk as computed by the offset of the chunk divided
+ * by the chunk shape.
+ * @return the raw encoded chunk bytes, or {@code null} if the chunk is not present in the store
+ * (i.e. it holds the fill value).
+ * @throws ZarrException throws ZarrException if the requested chunk is outside the array's domain
+ */
+ @Nullable
+ public ByteBuffer readChunkDirect(long[] chunkCoords) throws ZarrException {
+ if (!chunkIsInArray(chunkCoords)) {
+ throw new ZarrException("Attempting to read data outside of the array's domain.");
+ }
+ ArrayMetadata metadata = metadata();
+ final String[] chunkKeys = metadata.chunkKeyEncoding().encodeChunkKey(chunkCoords);
+ final StoreHandle chunkHandle = storeHandle.resolve(chunkKeys);
+
+ return chunkHandle.read();
+ }
+
+ /**
+ * The shape of the smallest unit this array encodes independently, i.e. the unit that the codec
+ * pipeline compresses and that {@link #readInnerChunkDirect(long[])} addresses.
+ *
+ * For a sharded array this is the inner chunk shape of the sharding codec (the innermost one, if
+ * shards are nested); for any other array it is simply {@link ArrayMetadata#chunkShape()}.
+ */
+ public int[] innerChunkShape() {
+ return codecPipeline.innerChunkShape();
+ }
+
+ /**
+ * Reads the already-encoded bytes of one inner chunk directly from the store, bypassing the codec
+ * pipeline entirely. No decoding is performed; the raw stored bytes are returned as-is.
+ *
+ * Unlike {@link #readChunkDirect(long[])}, which returns a whole stored chunk, this addresses the
+ * unit that the codec pipeline actually encodes. For a sharded array that is a single inner chunk
+ * of a shard: only the shard index and the bytes of that one inner chunk are read from the store,
+ * not the whole shard. So for an array using the {@code jpeg} codec this returns exactly one JPEG.
+ * For an unsharded array the two grids coincide and this is equivalent to
+ * {@link #readChunkDirect(long[])}.
+ *
+ * @param innerChunkCoords The coordinates of the inner chunk on the grid given by
+ * {@link #innerChunkShape()}, spanning the whole array.
+ * @return the raw encoded inner chunk bytes, or {@code null} if the inner chunk is not present in
+ * the store (i.e. it holds the fill value).
+ * @throws ZarrException throws ZarrException if the requested inner chunk is outside the array's
+ * domain, or if the codec pipeline does not allow addressing inner chunks
+ */
+ @Nullable
+ public ByteBuffer readInnerChunkDirect(long[] innerChunkCoords) throws ZarrException {
+ final long[][] splitCoords = splitInnerChunkCoords(innerChunkCoords);
+ final String[] chunkKeys = metadata().chunkKeyEncoding().encodeChunkKey(splitCoords[0]);
+ final StoreHandle chunkHandle = storeHandle.resolve(chunkKeys);
+
+ if (!codecPipeline.supportsPartialDecode()) {
+ return chunkHandle.read();
+ }
+ return codecPipeline.readInnerChunkEncoded(chunkHandle, splitCoords[1]);
+ }
+
+ /**
+ * Splits coordinates on the {@link #innerChunkShape()} grid into the coordinates of the stored
+ * chunk holding that inner chunk and the coordinates of the inner chunk within that stored chunk.
+ *
+ * @return an array of {@code {chunkCoords, coordsInChunk}}
+ * @throws IllegalArgumentException if {@code innerChunkCoords} has the wrong rank
+ * @throws ZarrException if the inner chunk is outside the array's domain
+ */
+ long[][] splitInnerChunkCoords(long[] innerChunkCoords) throws ZarrException {
+ final ArrayMetadata metadata = metadata();
+ final int ndim = metadata.ndim();
+ if (innerChunkCoords.length != ndim) {
+ throw new IllegalArgumentException(
+ "'innerChunkCoords' needs to have rank '" + ndim + "'.");
+ }
+ final int[] chunkShape = metadata.chunkShape();
+ final int[] innerChunkShape = innerChunkShape();
+
+ final long[] chunkCoords = new long[ndim];
+ final long[] coordsInChunk = new long[ndim];
+ for (int dimIdx = 0; dimIdx < ndim; dimIdx++) {
+ if (innerChunkCoords[dimIdx] < 0
+ || innerChunkCoords[dimIdx] * innerChunkShape[dimIdx] >= metadata.shape[dimIdx]) {
+ throw new ZarrException("Attempting to access data outside of the array's domain.");
+ }
+ final int innerChunksPerChunk = chunkShape[dimIdx] / innerChunkShape[dimIdx];
+ chunkCoords[dimIdx] = innerChunkCoords[dimIdx] / innerChunksPerChunk;
+ coordsInChunk[dimIdx] = innerChunkCoords[dimIdx] % innerChunksPerChunk;
+ }
+ return new long[][]{chunkCoords, coordsInChunk};
+ }
+
+ /**
+ * A writer for storing already-encoded inner chunks into this array without decoding or
+ * re-encoding anything, batching the changes so that each affected shard is rebuilt exactly once.
+ *
+ * This is the encoded-bytes counterpart of {@link #readInnerChunkDirect(long[])}. See
+ * {@link InnerChunkWriter} for the safety contract.
+ */
+ public InnerChunkWriter innerChunkWriter() {
+ return new InnerChunkWriter(this);
+ }
+
+ /**
+ * Stores the already-encoded bytes of one inner chunk, bypassing the codec pipeline entirely. No
+ * encoding is performed; the supplied bytes are stored verbatim.
+ *
+ * This is the encoded-bytes counterpart of {@link #readInnerChunkDirect(long[])}. For a sharded
+ * array the containing shard is read, rebuilt with this inner chunk's bytes spliced in, and stored
+ * again; every other inner chunk of that shard is copied through still encoded, so nothing is
+ * recompressed. For an unsharded array the inner chunk grid is the chunk grid and this is
+ * equivalent to {@link #writeChunkDirect(long[], ByteBuffer)}.
+ *
+ * To change several inner chunks, use {@link #innerChunkWriter()} instead: it rebuilds each
+ * affected shard once rather than once per inner chunk.
+ *
+ * Unsafe: because the bytes are stored without decoding, they cannot be verified. See
+ * {@link InnerChunkWriter} for the full contract.
+ *
+ * @param innerChunkCoords The coordinates of the inner chunk on the grid given by
+ * {@link #innerChunkShape()}, spanning the whole array.
+ * @param innerChunkBytes The already-encoded bytes to store, or {@code null} to remove the inner
+ * chunk (a subsequent read then returns the fill value).
+ * @throws ZarrException throws ZarrException if the requested inner chunk is outside the array's
+ * domain
+ */
+ public void writeInnerChunkDirect(long[] innerChunkCoords, @Nullable ByteBuffer innerChunkBytes)
+ throws ZarrException {
+ new InnerChunkWriter(this).put(innerChunkCoords, innerChunkBytes).flush();
+ }
+
/**
* Deletes chunks that are completely outside the new shape and trims boundary chunks.
*
diff --git a/src/main/java/dev/zarr/zarrjava/core/InnerChunkWriter.java b/src/main/java/dev/zarr/zarrjava/core/InnerChunkWriter.java
new file mode 100644
index 00000000..c701a063
--- /dev/null
+++ b/src/main/java/dev/zarr/zarrjava/core/InnerChunkWriter.java
@@ -0,0 +1,180 @@
+package dev.zarr.zarrjava.core;
+
+import dev.zarr.zarrjava.ZarrException;
+import dev.zarr.zarrjava.core.codec.ArrayBytesCodec.WithPartialDecode.InnerChunkUpdate;
+import dev.zarr.zarrjava.store.StoreHandle;
+
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.LinkedHashMap;
+import java.util.Map;
+
+/**
+ * Writes already-encoded inner chunks into a Zarr array without decoding or re-encoding anything,
+ * batching the changes so that each affected shard is rebuilt exactly once.
+ *
+ * Inner chunks are staged with {@link #put(long[], ByteBuffer)} and applied by {@link #flush()}. A
+ * shard holds an index plus the independently compressed bytes of its inner chunks, and stores offer
+ * no partial writes, so replacing one inner chunk means reading the shard, splicing the new bytes in,
+ * and storing the whole shard again. Every other inner chunk is copied through byte for byte, still
+ * encoded; only the shard index is ever decoded. Compare
+ * {@link Array#writeChunk(long[], ucar.ma2.Array)}, which for a sharded array decodes the shard and
+ * recompresses every inner chunk in it.
+ *
+ * Batching is what makes this cheap: staging several inner chunks that fall into the same shard and
+ * then flushing once rebuilds that shard once, not once per inner chunk.
+ *
+ * Unsafe: because nothing is decoded, nothing can be validated. The caller must guarantee that
+ * each buffer holds exactly one inner chunk of shape {@link Array#innerChunkShape()}, encoded exactly
+ * as this array's inner codec pipeline would produce it (matching data type, endianness, compressor
+ * and options). Incompatible bytes are stored happily and silently corrupt the array; the damage
+ * surfaces only when something reads it. {@link Array#readInnerChunkDirect(long[])} on an array with
+ * identical metadata is the one source of such bytes that is safe by construction.
+ *
+ * Whole shards are rewritten. Two writers that touch the same shard concurrently each read the
+ * shard, splice their own change in and store it, so the later store silently discards the earlier
+ * writer's inner chunk. Callers must serialize writes per shard. Storing a shard is also not
+ * crash-atomic in every store implementation, so an interrupted flush can leave a truncated shard,
+ * losing inner chunks that were not being changed.
+ *
+ * Instances are not thread-safe.
+ */
+public final class InnerChunkWriter {
+
+ @Nonnull
+ private final Array array;
+ /**
+ * The staged inner chunks, grouped by the stored chunk holding them.
+ */
+ private final Map pending = new LinkedHashMap<>();
+
+ public InnerChunkWriter(@Nonnull Array array) {
+ this.array = array;
+ }
+
+ /**
+ * Stages the already-encoded bytes of one inner chunk to be stored on the next {@link #flush()}.
+ * Nothing is read from or written to the store here.
+ *
+ * Only the {@code remaining()} bytes from the buffer's current position are used and the buffer's
+ * position is not advanced, but the bytes are not copied: do not modify the buffer's
+ * contents before {@link #flush()} returns.
+ *
+ * Staging the same coordinates twice replaces the earlier entry.
+ *
+ * @param innerChunkCoords The coordinates of the inner chunk on the grid given by
+ * {@link Array#innerChunkShape()}, spanning the whole array. This is the
+ * same coordinate space {@link Array#readInnerChunkDirect(long[])} uses.
+ * @param innerChunkBytes The already-encoded inner chunk bytes, or {@code null} to remove the
+ * inner chunk (a subsequent read then returns the fill value).
+ * @return this writer, so that calls can be chained
+ * @throws IllegalArgumentException if {@code innerChunkCoords} has the wrong rank
+ * @throws ZarrException if the inner chunk is outside the array's domain, or if
+ * {@code innerChunkBytes} has no remaining bytes
+ */
+ public InnerChunkWriter put(long[] innerChunkCoords, @Nullable ByteBuffer innerChunkBytes)
+ throws ZarrException {
+ final long[][] splitCoords = array.splitInnerChunkCoords(innerChunkCoords);
+ final long[] chunkCoords = splitCoords[0];
+ final long[] coordsInChunk = splitCoords[1];
+
+ if (innerChunkBytes != null && !innerChunkBytes.hasRemaining()) {
+ throw new ZarrException(
+ "The encoded bytes for inner chunk " + Arrays.toString(innerChunkCoords) + " are empty. "
+ + "Pass 'null' to remove the inner chunk instead.");
+ }
+ // duplicate() to snapshot position and limit without ever consuming the caller's buffer
+ final ByteBuffer stagedBytes = innerChunkBytes == null ? null : innerChunkBytes.duplicate();
+
+ final String[] chunkKeys = array.metadata().chunkKeyEncoding().encodeChunkKey(chunkCoords);
+ pending.computeIfAbsent(String.join("/", chunkKeys), key -> new ChunkBatch(chunkCoords))
+ .updates.put(Arrays.toString(coordsInChunk),
+ new InnerChunkUpdate(coordsInChunk, stagedBytes));
+ return this;
+ }
+
+ /**
+ * Stores everything staged since the last flush, rebuilding each affected shard exactly once, and
+ * clears the staged inner chunks. A stored chunk that would hold no inner chunks at all is deleted.
+ * Flushing with nothing staged is a no-op, and the writer is reusable afterwards.
+ *
+ * Stored chunks are processed one at a time and each is dropped from the staged set only once its
+ * write has succeeded, so a failed flush can be retried: rebuilding a shard twice from the same
+ * staged bytes produces byte-identical output.
+ *
+ * @throws ZarrException if an existing shard cannot be parsed, or if a rebuilt shard would be too
+ * large to address
+ */
+ public void flush() throws ZarrException {
+ final ArrayMetadata metadata = array.metadata();
+ final boolean addressesInnerChunks = array.codecPipeline.supportsPartialDecode();
+ final long innerChunksPerChunk = innerChunksPerChunk(metadata);
+
+ final Iterator batches = pending.values().iterator();
+ while (batches.hasNext()) {
+ final ChunkBatch batch = batches.next();
+ final String[] chunkKeys = metadata.chunkKeyEncoding().encodeChunkKey(batch.chunkCoords);
+ final StoreHandle chunkHandle = array.storeHandle.resolve(chunkKeys);
+
+ if (!addressesInnerChunks) {
+ // The inner chunk grid is the stored chunk grid, so the staged bytes are the whole
+ // stored chunk. There is exactly one staged inner chunk per stored chunk here.
+ final InnerChunkUpdate update = batch.updates.values().iterator().next();
+ if (update.innerChunkBytes == null) {
+ chunkHandle.delete();
+ } else {
+ chunkHandle.set(update.innerChunkBytes.duplicate());
+ }
+ batches.remove();
+ continue;
+ }
+
+ // Nothing of the existing shard survives when every inner chunk of it is being replaced, so
+ // in that case the shard does not need to be read at all.
+ final boolean replacesWholeChunk = batch.updates.size() == innerChunksPerChunk
+ && batch.updates.values().stream().allMatch(u -> u.innerChunkBytes != null);
+ final ByteBuffer oldChunkBytes = replacesWholeChunk ? null : chunkHandle.read();
+
+ final ByteBuffer newChunkBytes = array.codecPipeline.mergeInnerChunksEncoded(
+ oldChunkBytes, new ArrayList<>(batch.updates.values()));
+
+ if (newChunkBytes == null) {
+ chunkHandle.delete();
+ } else {
+ chunkHandle.set(newChunkBytes);
+ }
+ batches.remove();
+ }
+ }
+
+ /**
+ * How many inner chunks fit into one stored chunk.
+ */
+ private long innerChunksPerChunk(ArrayMetadata metadata) {
+ final int[] chunkShape = metadata.chunkShape();
+ final int[] innerChunkShape = array.innerChunkShape();
+ long innerChunksPerChunk = 1;
+ for (int dimIdx = 0; dimIdx < chunkShape.length; dimIdx++) {
+ innerChunksPerChunk *= chunkShape[dimIdx] / innerChunkShape[dimIdx];
+ }
+ return innerChunksPerChunk;
+ }
+
+ /**
+ * The staged inner chunks of one stored chunk, keyed by their coordinates within it so that
+ * staging the same inner chunk twice replaces the earlier entry.
+ */
+ private static final class ChunkBatch {
+
+ final long[] chunkCoords;
+ final Map updates = new LinkedHashMap<>();
+
+ ChunkBatch(long[] chunkCoords) {
+ this.chunkCoords = chunkCoords;
+ }
+ }
+}
diff --git a/src/main/java/dev/zarr/zarrjava/core/codec/ArrayBytesCodec.java b/src/main/java/dev/zarr/zarrjava/core/codec/ArrayBytesCodec.java
index f8505be0..118987f4 100644
--- a/src/main/java/dev/zarr/zarrjava/core/codec/ArrayBytesCodec.java
+++ b/src/main/java/dev/zarr/zarrjava/core/codec/ArrayBytesCodec.java
@@ -4,7 +4,9 @@
import dev.zarr.zarrjava.store.StoreHandle;
import ucar.ma2.Array;
+import javax.annotation.Nullable;
import java.nio.ByteBuffer;
+import java.util.List;
public abstract class ArrayBytesCodec extends AbstractCodec {
@@ -23,6 +25,72 @@ public abstract static class WithPartialDecode extends ArrayBytesCodec {
protected abstract Array decodePartial(
StoreHandle handle, long[] offset, int[] shape
) throws ZarrException;
+
+ /**
+ * The shape of the smallest unit that this codec encodes independently inside one stored
+ * chunk, and that {@link #readInnerChunkEncoded} can therefore address.
+ *
+ * Where this codec nests further codecs of the same kind, this is the innermost such shape
+ * that remains addressable by byte offset; the recursion stops as soon as a level's bytes
+ * would have to be decoded before its inner units could be located.
+ */
+ public abstract int[] innerChunkShape();
+
+ /**
+ * Reads the encoded bytes of a single inner chunk out of a stored chunk, without decoding
+ * them.
+ *
+ * @param handle the store handle of the stored chunk
+ * @param innerChunkCoords the coordinates of the inner chunk relative to the stored chunk, on
+ * the grid given by {@link #innerChunkShape()}
+ * @return the encoded inner chunk bytes, or {@code null} if the inner chunk is not present
+ */
+ @Nullable
+ protected abstract ByteBuffer readInnerChunkEncoded(
+ StoreHandle handle, long[] innerChunkCoords
+ ) throws ZarrException;
+
+ /**
+ * Rebuilds a stored chunk so that the given inner chunks hold the given already-encoded bytes,
+ * copying every inner chunk that is kept through without decoding it.
+ *
+ * This is a pure bytes-to-bytes operation; nothing is read from or written to a store.
+ *
+ * @param chunkBytes the current bytes of the stored chunk, or {@code null} if the stored chunk
+ * does not exist yet
+ * @param updates the inner chunks to replace or remove, with coordinates relative to the
+ * stored chunk on the grid given by {@link #innerChunkShape()}
+ * @return the new bytes of the stored chunk, positioned at 0 and sized exactly, or
+ * {@code null} if the stored chunk would hold no inner chunks at all and should
+ * therefore be removed
+ */
+ @Nullable
+ protected abstract ByteBuffer mergeInnerChunksEncoded(
+ @Nullable ByteBuffer chunkBytes, List updates
+ ) throws ZarrException;
+
+ /**
+ * The new encoded bytes of one inner chunk, or its removal.
+ */
+ public static final class InnerChunkUpdate {
+
+ /**
+ * The coordinates of the inner chunk relative to the stored chunk, on the grid given by
+ * {@link #innerChunkShape()}.
+ */
+ public final long[] innerChunkCoords;
+
+ /**
+ * The already-encoded inner chunk bytes, or {@code null} to remove the inner chunk.
+ */
+ @Nullable
+ public final ByteBuffer innerChunkBytes;
+
+ public InnerChunkUpdate(long[] innerChunkCoords, @Nullable ByteBuffer innerChunkBytes) {
+ this.innerChunkCoords = innerChunkCoords;
+ this.innerChunkBytes = innerChunkBytes;
+ }
+ }
}
}
diff --git a/src/main/java/dev/zarr/zarrjava/core/codec/CodecPipeline.java b/src/main/java/dev/zarr/zarrjava/core/codec/CodecPipeline.java
index d06b3165..f865e71c 100644
--- a/src/main/java/dev/zarr/zarrjava/core/codec/CodecPipeline.java
+++ b/src/main/java/dev/zarr/zarrjava/core/codec/CodecPipeline.java
@@ -6,8 +6,10 @@
import ucar.ma2.Array;
import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
import java.nio.ByteBuffer;
import java.util.Arrays;
+import java.util.List;
public class CodecPipeline {
@@ -83,6 +85,73 @@ public boolean supportsPartialDecode() {
return codecs.length == 1 && codecs[0] instanceof ArrayBytesCodec.WithPartialDecode;
}
+ /**
+ * The shape of the smallest unit this pipeline encodes independently, i.e. the unit that
+ * {@link #readInnerChunkEncoded} addresses and the grid its coordinates are on.
+ *
+ * For a sharded array this is the sharding codec's inner chunk shape. Shards may be nested, in
+ * which case this is the innermost inner chunk shape still addressable by byte offset: the
+ * recursion stops at any level whose inner codecs are not a single sharding codec, since the
+ * nested shard's bytes would have to be decoded before its index could be located. For any
+ * other pipeline this is the chunk shape itself.
+ */
+ public int[] innerChunkShape() {
+ if (!supportsPartialDecode()) {
+ return arrayMetadata.chunkShape;
+ }
+ return ((ArrayBytesCodec.WithPartialDecode) getArrayBytesCodec()).innerChunkShape();
+ }
+
+ /**
+ * Reads the encoded bytes of a single inner chunk out of a stored chunk, without decoding them.
+ *
+ * @param storeHandle the store handle of the stored chunk
+ * @param innerChunkCoords the coordinates of the inner chunk relative to the stored chunk, on the
+ * grid given by {@link #innerChunkShape()}
+ * @return the encoded inner chunk bytes, or {@code null} if the inner chunk is not present
+ */
+ @Nullable
+ public ByteBuffer readInnerChunkEncoded(
+ @Nonnull StoreHandle storeHandle, long[] innerChunkCoords
+ ) throws ZarrException {
+ if (!supportsPartialDecode()) {
+ throw new ZarrException(
+ "Reading individual inner chunks is not supported for these codecs. " + Arrays.toString(
+ codecs));
+ }
+ return ((ArrayBytesCodec.WithPartialDecode) getArrayBytesCodec()).readInnerChunkEncoded(
+ storeHandle, innerChunkCoords);
+ }
+
+ /**
+ * Rebuilds a stored chunk so that the given inner chunks hold the given already-encoded bytes,
+ * copying every inner chunk that is kept through without decoding it.
+ *
+ * Splicing encoded bytes into a stored chunk is only sound because a pipeline that supports this
+ * consists of the single {@link ArrayBytesCodec.WithPartialDecode} codec: no bytes-to-bytes codec
+ * wraps its output, so the bytes that codec produces are the stored object verbatim.
+ *
+ * @param chunkBytes the current bytes of the stored chunk, or {@code null} if the stored chunk does
+ * not exist yet
+ * @param updates the inner chunks to replace or remove, with coordinates relative to the stored
+ * chunk on the grid given by {@link #innerChunkShape()}
+ * @return the new bytes of the stored chunk, or {@code null} if the stored chunk would hold no
+ * inner chunks at all and should therefore be removed
+ */
+ @Nullable
+ public ByteBuffer mergeInnerChunksEncoded(
+ @Nullable ByteBuffer chunkBytes,
+ List updates
+ ) throws ZarrException {
+ if (!supportsPartialDecode()) {
+ throw new ZarrException(
+ "Writing individual inner chunks is not supported for these codecs. "
+ + Arrays.toString(codecs));
+ }
+ return ((ArrayBytesCodec.WithPartialDecode) getArrayBytesCodec()).mergeInnerChunksEncoded(
+ chunkBytes, updates);
+ }
+
@Nonnull
public Array decodePartial(
@Nonnull StoreHandle storeHandle,
diff --git a/src/main/java/dev/zarr/zarrjava/v3/codec/core/ShardingIndexedCodec.java b/src/main/java/dev/zarr/zarrjava/v3/codec/core/ShardingIndexedCodec.java
index faf076e9..4694f34a 100644
--- a/src/main/java/dev/zarr/zarrjava/v3/codec/core/ShardingIndexedCodec.java
+++ b/src/main/java/dev/zarr/zarrjava/v3/codec/core/ShardingIndexedCodec.java
@@ -18,10 +18,13 @@
import ucar.ma2.InvalidRangeException;
import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
+import java.util.HashMap;
import java.util.List;
+import java.util.Map;
public class ShardingIndexedCodec extends ArrayBytesCodec.WithPartialDecode implements Codec {
@@ -242,6 +245,218 @@ private Array decodeInternal(
return outputArray;
}
+ /**
+ * The nested sharding codec, if this shard's inner chunks are themselves shards. Only recognized
+ * when sharding is the sole inner codec: with additional inner codecs wrapping the nested shard,
+ * its bytes cannot be addressed without decoding them first.
+ */
+ @Nullable
+ private ShardingIndexedCodec nestedShardingCodec() {
+ if (configuration.codecs.length == 1
+ && configuration.codecs[0] instanceof ShardingIndexedCodec) {
+ return (ShardingIndexedCodec) configuration.codecs[0];
+ }
+ return null;
+ }
+
+ @Override
+ public int[] innerChunkShape() {
+ final ShardingIndexedCodec nested = nestedShardingCodec();
+ return nested == null ? configuration.chunkShape : nested.innerChunkShape();
+ }
+
+ @Override
+ @Nullable
+ protected ByteBuffer readInnerChunkEncoded(StoreHandle handle, long[] innerChunkCoords)
+ throws ZarrException {
+ return readInnerChunkEncoded(new StoreHandleDataProvider(handle), innerChunkCoords);
+ }
+
+ @Nullable
+ private ByteBuffer readInnerChunkEncoded(DataProvider dataProvider, long[] innerChunkCoords)
+ throws ZarrException {
+ final int ndim = arrayMetadata.ndim();
+ final ShardingIndexedCodec nested = nestedShardingCodec();
+ final int[] innerChunkShape = innerChunkShape();
+ final int[] chunksPerShard = getChunksPerShard(arrayMetadata);
+
+ // Split the innermost-grid coordinates into this level's chunk coordinates and the remainder
+ // that addresses the inner chunk within a nested shard.
+ final long[] chunkCoords = new long[ndim];
+ final long[] nestedChunkCoords = new long[ndim];
+ for (int dimIdx = 0; dimIdx < ndim; dimIdx++) {
+ final int chunksPerInnerChunk = configuration.chunkShape[dimIdx] / innerChunkShape[dimIdx];
+ chunkCoords[dimIdx] = innerChunkCoords[dimIdx] / chunksPerInnerChunk;
+ nestedChunkCoords[dimIdx] = innerChunkCoords[dimIdx] % chunksPerInnerChunk;
+ if (chunkCoords[dimIdx] < 0 || chunkCoords[dimIdx] >= chunksPerShard[dimIdx]) {
+ throw new ZarrException("Attempting to read an inner chunk outside of the shard.");
+ }
+ }
+
+ final int shardIndexByteLength = (int) getShardIndexSize(arrayMetadata);
+ final ByteBuffer shardIndexBytes = this.configuration.indexLocation.equals("start")
+ ? dataProvider.readPrefix(shardIndexByteLength)
+ : dataProvider.readSuffix(shardIndexByteLength);
+ if (shardIndexBytes == null) {
+ return null;
+ }
+
+ final Array shardIndexArray = indexCodecPipeline.decode(shardIndexBytes);
+ final long chunkByteOffset = getValueFromShardIndexArray(shardIndexArray, chunkCoords, 0);
+ final long chunkByteLength = getValueFromShardIndexArray(shardIndexArray, chunkCoords, 1);
+ if (chunkByteOffset == -1 || chunkByteLength == -1) {
+ return null;
+ }
+
+ final ByteBuffer chunkBytes = dataProvider.read(chunkByteOffset, chunkByteLength);
+ if (chunkBytes == null || nested == null) {
+ return chunkBytes;
+ }
+ return nested.readInnerChunkEncoded(new ByteBufferDataProvider(chunkBytes), nestedChunkCoords);
+ }
+
+ @Override
+ @Nullable
+ protected ByteBuffer mergeInnerChunksEncoded(
+ @Nullable ByteBuffer shardBytes, List updates) throws ZarrException {
+ final int ndim = arrayMetadata.ndim();
+ final ShardingIndexedCodec nested = nestedShardingCodec();
+ final int[] innerChunkShape = innerChunkShape();
+ final int[] chunksPerShard = getChunksPerShard(arrayMetadata);
+ final int shardIndexByteLength = (int) getShardIndexSize(arrayMetadata);
+ final boolean indexAtStart = this.configuration.indexLocation.equals("start");
+
+ // Split the innermost-grid coordinates of every update into this level's chunk coordinates and
+ // the remainder addressing the inner chunk within a nested shard, then group by the former so
+ // that each of this level's chunks is rebuilt at most once.
+ final Map> updatesPerChunk = new HashMap<>();
+ for (final InnerChunkUpdate update : updates) {
+ if (update.innerChunkCoords.length != ndim) {
+ throw new IllegalArgumentException(
+ "'innerChunkCoords' needs to have rank '" + ndim + "'.");
+ }
+ final long[] chunkCoords = new long[ndim];
+ final long[] nestedChunkCoords = new long[ndim];
+ for (int dimIdx = 0; dimIdx < ndim; dimIdx++) {
+ final int chunksPerInnerChunk = configuration.chunkShape[dimIdx] / innerChunkShape[dimIdx];
+ chunkCoords[dimIdx] = update.innerChunkCoords[dimIdx] / chunksPerInnerChunk;
+ nestedChunkCoords[dimIdx] = update.innerChunkCoords[dimIdx] % chunksPerInnerChunk;
+ if (chunkCoords[dimIdx] < 0 || chunkCoords[dimIdx] >= chunksPerShard[dimIdx]) {
+ throw new ZarrException("Attempting to write an inner chunk outside of the shard.");
+ }
+ }
+ updatesPerChunk.computeIfAbsent(Arrays.toString(chunkCoords), k -> new ArrayList<>())
+ .add(new InnerChunkUpdate(nestedChunkCoords, update.innerChunkBytes));
+ }
+
+ // Decode the existing shard index, if there is an existing shard. A shard that cannot be parsed
+ // is never overwritten.
+ ByteBufferDataProvider dataProvider = null;
+ Array oldShardIndexArray = null;
+ int shardByteLength = 0;
+ if (shardBytes != null && shardBytes.hasRemaining()) {
+ // slice() so that capacity == remaining and index offsets are relative to this buffer
+ final ByteBuffer oldShardBytes = shardBytes.slice();
+ shardByteLength = oldShardBytes.capacity();
+ if (shardByteLength < shardIndexByteLength) {
+ throw new ZarrException(
+ "The existing shard is " + shardByteLength + " bytes, which is too small to hold its "
+ + shardIndexByteLength + " byte shard index.");
+ }
+ dataProvider = new ByteBufferDataProvider(oldShardBytes);
+ // readPrefix/readSuffix return exactly sized slices, which the crc32c index codec requires
+ // because it verifies the checksum against the buffer's capacity.
+ final ByteBuffer shardIndexBytes = indexAtStart
+ ? dataProvider.readPrefix(shardIndexByteLength)
+ : dataProvider.readSuffix(shardIndexByteLength);
+ oldShardIndexArray = indexCodecPipeline.decode(shardIndexBytes);
+ }
+
+ final Array shardIndexArray = Array.factory(ucar.ma2.DataType.ULONG,
+ extendArrayBy1(chunksPerShard, 2));
+ // Array.factory zero-initializes, and offset 0 / length 0 reads back as a present but empty
+ // inner chunk rather than an absent one, so every entry has to be set to -1 explicitly. The
+ // fill value must be a long: MultiArrayUtils casts it to (long) without converting.
+ MultiArrayUtils.fill(shardIndexArray, -1L);
+
+ // Walk this level's chunk grid in row-major order, so that the layout is a deterministic
+ // function of the old shard and the updates. Untouched inner chunks are copied through still
+ // encoded; only the shard index above was ever decoded.
+ final ArrayMetadata.CoreArrayMetadata shardMetadata = codecPipeline.arrayMetadata;
+ final List chunkBytesList = new ArrayList<>();
+ long payloadByteLength = 0;
+ final long chunkByteOffsetShift = indexAtStart ? shardIndexByteLength : 0;
+
+ for (final long[] chunkCoords : IndexingUtils.computeChunkCoords(shardMetadata.shape,
+ shardMetadata.chunkShape)) {
+ ByteBuffer oldChunkBytes = null;
+ if (oldShardIndexArray != null) {
+ final long oldChunkByteOffset =
+ getValueFromShardIndexArray(oldShardIndexArray, chunkCoords, 0);
+ final long oldChunkByteLength =
+ getValueFromShardIndexArray(oldShardIndexArray, chunkCoords, 1);
+ if (oldChunkByteOffset != -1 || oldChunkByteLength != -1) {
+ if (oldChunkByteOffset < 0 || oldChunkByteLength < 0
+ || oldChunkByteOffset + oldChunkByteLength > shardByteLength) {
+ throw new ZarrException(
+ "The existing shard index is corrupt: inner chunk " + Arrays.toString(chunkCoords)
+ + " is at offset " + oldChunkByteOffset + " with length " + oldChunkByteLength
+ + ", which does not fit in a shard of " + shardByteLength + " bytes.");
+ }
+ oldChunkBytes = dataProvider.read(oldChunkByteOffset, oldChunkByteLength);
+ }
+ }
+
+ final List chunkUpdates = updatesPerChunk.get(Arrays.toString(chunkCoords));
+ final ByteBuffer newChunkBytes;
+ if (chunkUpdates == null) {
+ newChunkBytes = oldChunkBytes;
+ } else if (nested == null) {
+ // Without nesting the grids coincide, so there is exactly one update per chunk.
+ newChunkBytes = chunkUpdates.get(chunkUpdates.size() - 1).innerChunkBytes;
+ } else {
+ // A nested shard is just one inner chunk of this shard: rebuild it from its own bytes and
+ // splice the result in as an opaque blob.
+ newChunkBytes = nested.mergeInnerChunksEncoded(oldChunkBytes, chunkUpdates);
+ }
+
+ if (newChunkBytes == null) {
+ continue;
+ }
+ setValueFromShardIndexArray(shardIndexArray, chunkCoords, 0,
+ payloadByteLength + chunkByteOffsetShift);
+ setValueFromShardIndexArray(shardIndexArray, chunkCoords, 1, newChunkBytes.remaining());
+ chunkBytesList.add(newChunkBytes);
+ payloadByteLength += newChunkBytes.remaining();
+ }
+
+ if (chunkBytesList.isEmpty()) {
+ return null;
+ }
+
+ final long shardBytesLength = payloadByteLength + shardIndexByteLength;
+ if (shardBytesLength > Integer.MAX_VALUE) {
+ throw new ZarrException(
+ "The rebuilt shard would be " + shardBytesLength + " bytes, but a shard's contents are "
+ + "addressed with 32-bit offsets and cannot exceed " + Integer.MAX_VALUE + " bytes.");
+ }
+
+ final ByteBuffer newShardBytes = ByteBuffer.allocate((int) shardBytesLength);
+ if (indexAtStart) {
+ newShardBytes.put(indexCodecPipeline.encode(shardIndexArray));
+ }
+ for (final ByteBuffer chunkBytes : chunkBytesList) {
+ // duplicate() because put consumes the source position and the same buffer may have been
+ // supplied for more than one inner chunk
+ newShardBytes.put(chunkBytes.duplicate());
+ }
+ if (!indexAtStart) {
+ newShardBytes.put(indexCodecPipeline.encode(shardIndexArray));
+ }
+ newShardBytes.rewind();
+ return newShardBytes;
+ }
+
@Override
public Array decodePartial(StoreHandle chunkHandle, long[] offset, int[] shape) throws ZarrException {
if (Arrays.equals(shape, arrayMetadata.chunkShape)) {
diff --git a/src/test/java/dev/zarr/zarrjava/ZarrV3Test.java b/src/test/java/dev/zarr/zarrjava/ZarrV3Test.java
index 614a2b90..65abe2cf 100644
--- a/src/test/java/dev/zarr/zarrjava/ZarrV3Test.java
+++ b/src/test/java/dev/zarr/zarrjava/ZarrV3Test.java
@@ -4,6 +4,7 @@
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import dev.zarr.zarrjava.core.Attributes;
+import dev.zarr.zarrjava.core.InnerChunkWriter;
import dev.zarr.zarrjava.store.FilesystemStore;
import dev.zarr.zarrjava.store.HttpStore;
import dev.zarr.zarrjava.store.MemoryStore;
@@ -25,8 +26,10 @@
import org.junit.jupiter.params.provider.ValueSource;
import ucar.ma2.MAMath;
+import javax.annotation.Nullable;
import java.io.BufferedReader;
import java.io.IOException;
+import java.nio.ByteBuffer;
import java.nio.file.Files;
import java.nio.file.NoSuchFileException;
import java.nio.file.Path;
@@ -1052,4 +1055,586 @@ public void testEndianness(DataType dataType, BytesCodec.Endian endian) throws I
ucar.ma2.Array readData = reopenedArray.read();
assertIsTestdata(readData, dataType);
}
+
+ @Test
+ public void testDirectChunkReadWrite() throws IOException, ZarrException {
+ // Source array: write data normally so a real encoded chunk lands in the store.
+ StoreHandle sourceHandle = new FilesystemStore(TESTOUTPUT).resolve("testDirectChunkReadWriteV3", "source");
+ ArrayMetadata metadata = Array.metadataBuilder()
+ .withShape(4, 4)
+ .withDataType(DataType.UINT32)
+ .withChunkShape(2, 2)
+ .withCodecs(c -> c.withBytes("LITTLE").withGzip())
+ .build();
+ Array source = Array.create(sourceHandle, metadata);
+ int[] chunkData = new int[]{1, 2, 3, 4};
+ source.writeChunk(new long[]{0, 0}, ucar.ma2.Array.factory(ucar.ma2.DataType.UINT, new int[]{2, 2}, chunkData));
+
+ // readChunkDirect returns the raw encoded bytes, byte-for-byte identical to what is on disk.
+ ByteBuffer encoded = source.readChunkDirect(new long[]{0, 0});
+ Assertions.assertNotNull(encoded);
+ ByteBuffer rawFromStore = sourceHandle.resolve(metadata.chunkKeyEncoding().encodeChunkKey(new long[]{0, 0})).read();
+ Assertions.assertEquals(rawFromStore, encoded);
+
+ // Write those encoded bytes directly into a fresh array, bypassing the codec pipeline.
+ StoreHandle targetHandle = new FilesystemStore(TESTOUTPUT).resolve("testDirectChunkReadWriteV3", "target");
+ Array target = Array.create(targetHandle, metadata);
+ target.writeChunkDirect(new long[]{0, 0}, encoded);
+
+ // Decoding the directly-written chunk yields the original data: the bytes were a valid chunk.
+ ucar.ma2.Array roundTripped = target.readChunk(new long[]{0, 0});
+ Assertions.assertArrayEquals(chunkData, (int[]) roundTripped.get1DJavaArray(ucar.ma2.DataType.INT));
+
+ // writeChunkDirect(null) deletes the chunk: a subsequent direct read is null and a normal read is fill value.
+ target.writeChunkDirect(new long[]{0, 0}, null);
+ Assertions.assertNull(target.readChunkDirect(new long[]{0, 0}));
+
+ // Out-of-domain coordinates are rejected for both direct methods.
+ assertThrows(ZarrException.class, () -> target.readChunkDirect(new long[]{99, 99}));
+ assertThrows(ZarrException.class, () -> target.writeChunkDirect(new long[]{99, 99}, encoded));
+
+ // Without sharding the inner chunk grid is the chunk grid, so both direct reads agree.
+ Assertions.assertArrayEquals(new int[]{2, 2}, source.innerChunkShape());
+ Assertions.assertEquals(source.readChunkDirect(new long[]{0, 0}),
+ source.readInnerChunkDirect(new long[]{0, 0}));
+ }
+
+ /**
+ * Extracts a single encoded inner chunk out of a shard and checks it is a standalone encoded
+ * chunk, by writing it directly into an unsharded array whose codecs match the shard's inner
+ * codecs and decoding it there.
+ */
+ private void assertInnerChunkDecodesTo(
+ Array shardedArray, long[] innerChunkCoords, String storePath, int[] expected
+ ) throws IOException, ZarrException {
+ ByteBuffer innerChunkBytes = shardedArray.readInnerChunkDirect(innerChunkCoords);
+ Assertions.assertNotNull(innerChunkBytes);
+
+ Array plainArray = Array.create(
+ new FilesystemStore(TESTOUTPUT).resolve(storePath),
+ Array.metadataBuilder()
+ .withShape(2, 2)
+ .withDataType(DataType.UINT32)
+ .withChunkShape(2, 2)
+ .withCodecs(c -> c.withBytes("LITTLE"))
+ .build());
+ plainArray.writeChunkDirect(new long[]{0, 0}, innerChunkBytes);
+ Assertions.assertArrayEquals(expected,
+ (int[]) plainArray.readChunk(new long[]{0, 0}).get1DJavaArray(ucar.ma2.DataType.INT));
+ }
+
+ @Test
+ public void testDirectInnerChunkRead() throws IOException, ZarrException {
+ // 8x8 array, 4x4 shards, 2x2 inner chunks: 4 shards of 4 inner chunks each.
+ int[] testData = new int[8 * 8];
+ Arrays.setAll(testData, p -> p);
+
+ StoreHandle storeHandle = new FilesystemStore(TESTOUTPUT).resolve("testDirectInnerChunkRead", "sharded");
+ Array array = Array.create(storeHandle, Array.metadataBuilder()
+ .withShape(8, 8)
+ .withDataType(DataType.UINT32)
+ .withChunkShape(4, 4)
+ .withCodecs(c -> c.withSharding(new int[]{2, 2}, c1 -> c1.withBytes("LITTLE")))
+ .build());
+ array.write(ucar.ma2.Array.factory(ucar.ma2.DataType.UINT, new int[]{8, 8}, testData));
+
+ Assertions.assertArrayEquals(new int[]{2, 2}, array.innerChunkShape());
+
+ // Inner chunk (3,1) covers rows 6-7, cols 2-3: shard (1,0), inner chunk (1,1) within it.
+ assertInnerChunkDecodesTo(array, new long[]{3, 1},
+ "testDirectInnerChunkRead/plain", new int[]{50, 51, 58, 59});
+
+ // The extracted bytes are only that inner chunk, not the whole shard.
+ ByteBuffer innerChunkBytes = array.readInnerChunkDirect(new long[]{3, 1});
+ ByteBuffer shardBytes = array.readChunkDirect(new long[]{1, 0});
+ Assertions.assertNotNull(shardBytes);
+ Assertions.assertTrue(innerChunkBytes.remaining() < shardBytes.remaining());
+
+ // Out-of-domain inner chunk coordinates are rejected.
+ assertThrows(ZarrException.class, () -> array.readInnerChunkDirect(new long[]{4, 0}));
+ assertThrows(ZarrException.class, () -> array.readInnerChunkDirect(new long[]{0, -1}));
+ }
+
+ @Test
+ public void testDirectInnerChunkReadAbsent() throws IOException, ZarrException {
+ // Only the top-left 4x4 region is written, and it is all fill value except one inner chunk.
+ StoreHandle storeHandle = new FilesystemStore(TESTOUTPUT).resolve("testDirectInnerChunkReadAbsent", "sharded");
+ Array array = Array.create(storeHandle, Array.metadataBuilder()
+ .withShape(8, 8)
+ .withDataType(DataType.UINT32)
+ .withChunkShape(4, 4)
+ .withCodecs(c -> c.withSharding(new int[]{2, 2}, c1 -> c1.withBytes("LITTLE")))
+ .build());
+ int[] shardData = new int[4 * 4];
+ shardData[0] = 42;
+ array.write(new long[]{0, 0}, ucar.ma2.Array.factory(ucar.ma2.DataType.UINT, new int[]{4, 4}, shardData));
+
+ // Inner chunk (0,0) holds the 42 and is present.
+ Assertions.assertNotNull(array.readInnerChunkDirect(new long[]{0, 0}));
+ // Inner chunk (1,1) is all fill value, so the shard index marks it as missing.
+ Assertions.assertNull(array.readInnerChunkDirect(new long[]{1, 1}));
+ // Inner chunk (3,3) lives in shard (1,1), which was never written at all.
+ Assertions.assertNull(array.readInnerChunkDirect(new long[]{3, 3}));
+ }
+
+ @Test
+ public void testDirectInnerChunkReadNestedSharding() throws IOException, ZarrException {
+ // 8x8 array in one 8x8 shard of 4x4 shards of 2x2 inner chunks.
+ int[] testData = new int[8 * 8];
+ Arrays.setAll(testData, p -> p);
+
+ StoreHandle storeHandle = new FilesystemStore(TESTOUTPUT).resolve("testDirectInnerChunkReadNestedSharding", "sharded");
+ Array array = Array.create(storeHandle, Array.metadataBuilder()
+ .withShape(8, 8)
+ .withDataType(DataType.UINT32)
+ .withChunkShape(8, 8)
+ .withCodecs(c -> c.withSharding(new int[]{4, 4},
+ c1 -> c1.withSharding(new int[]{2, 2}, c2 -> c2.withBytes("LITTLE"))))
+ .build());
+ array.write(ucar.ma2.Array.factory(ucar.ma2.DataType.UINT, new int[]{8, 8}, testData));
+
+ // The addressable unit is the innermost chunk shape, not the intermediate shard shape.
+ Assertions.assertArrayEquals(new int[]{2, 2}, array.innerChunkShape());
+
+ // Inner chunk (1,2) covers rows 2-3, cols 4-5: outer shard (0,0), nested shard (0,1),
+ // inner chunk (1,0) within that.
+ assertInnerChunkDecodesTo(array, new long[]{1, 2},
+ "testDirectInnerChunkReadNestedSharding/plain", new int[]{20, 21, 28, 29});
+ }
+
+ /**
+ * The inverse of {@link #assertInnerChunkDecodesTo}: encodes a 2x2 UINT32 chunk through a plain
+ * array, yielding bytes that are a valid inner chunk for a shard whose inner codecs match.
+ */
+ private ByteBuffer encodedInnerChunk(String storePath, int[] values)
+ throws IOException, ZarrException {
+ Array plainArray = Array.create(
+ new FilesystemStore(TESTOUTPUT).resolve(storePath),
+ Array.metadataBuilder()
+ .withShape(2, 2)
+ .withDataType(DataType.UINT32)
+ .withChunkShape(2, 2)
+ .withCodecs(c -> c.withBytes("LITTLE"))
+ .build());
+ plainArray.writeChunk(new long[]{0, 0},
+ ucar.ma2.Array.factory(ucar.ma2.DataType.UINT, new int[]{2, 2}, values));
+ return plainArray.readChunkDirect(new long[]{0, 0});
+ }
+
+ /**
+ * An 8x8 UINT32 array of 4x4 shards of 2x2 inner chunks, filled with a 0..63 ramp.
+ */
+ private Array shardedRampArray(String storePath, String indexLocation)
+ throws IOException, ZarrException {
+ int[] testData = new int[8 * 8];
+ Arrays.setAll(testData, p -> p);
+
+ Array array = Array.create(new FilesystemStore(TESTOUTPUT).resolve(storePath),
+ Array.metadataBuilder()
+ .withShape(8, 8)
+ .withDataType(DataType.UINT32)
+ .withChunkShape(4, 4)
+ .withCodecs(c -> c.withSharding(new int[]{2, 2},
+ c1 -> c1.withBytes("LITTLE"), indexLocation))
+ .build());
+ array.write(ucar.ma2.Array.factory(ucar.ma2.DataType.UINT, new int[]{8, 8}, testData));
+ return array;
+ }
+
+ /**
+ * The expected contents of {@link #shardedRampArray}, with the 2x2 inner chunk at
+ * {@code innerChunkCoords} replaced by {@code values}, or by the fill value if {@code values} is
+ * null.
+ */
+ private int[] rampWithInnerChunk(long[] innerChunkCoords, @Nullable int[] values) {
+ int[] expected = new int[8 * 8];
+ Arrays.setAll(expected, p -> p);
+ for (int row = 0; row < 2; row++) {
+ for (int col = 0; col < 2; col++) {
+ expected[(int) (innerChunkCoords[0] * 2 + row) * 8 + (int) (innerChunkCoords[1] * 2 + col)] =
+ values == null ? 0 : values[row * 2 + col];
+ }
+ }
+ return expected;
+ }
+
+ private void assertArrayHolds(Array array, int[] expected) throws ZarrException {
+ Assertions.assertArrayEquals(expected,
+ (int[]) array.read().get1DJavaArray(ucar.ma2.DataType.INT));
+ }
+
+ @Test
+ public void testInnerChunkWriteRoundTrip() throws IOException, ZarrException {
+ Array array = shardedRampArray("testInnerChunkWriteRoundTrip/sharded", "end");
+ int[] newValues = new int[]{900, 901, 902, 903};
+ ByteBuffer newBytes =
+ encodedInnerChunk("testInnerChunkWriteRoundTrip/donor", newValues);
+
+ // Inner chunk (3,1) covers rows 6-7, cols 2-3: shard (1,0), inner chunk (1,1) within it.
+ array.innerChunkWriter().put(new long[]{3, 1}, newBytes).flush();
+
+ // Only that 2x2 block changed, and the stored bytes are the ones handed over verbatim.
+ assertArrayHolds(array, rampWithInnerChunk(new long[]{3, 1}, newValues));
+ Assertions.assertEquals(newBytes, array.readInnerChunkDirect(new long[]{3, 1}));
+ }
+
+ @Test
+ public void testInnerChunkWriteLeavesOtherBlobsByteIdentical() throws IOException, ZarrException {
+ Array array =
+ shardedRampArray("testInnerChunkWriteLeavesOtherBlobsByteIdentical/sharded", "end");
+
+ // Snapshot the three inner chunks of shard (0,0) that are not going to be touched.
+ long[][] untouched = new long[][]{{0, 1}, {1, 0}, {1, 1}};
+ ByteBuffer[] before = new ByteBuffer[untouched.length];
+ for (int i = 0; i < untouched.length; i++) {
+ before[i] = array.readInnerChunkDirect(untouched[i]);
+ Assertions.assertNotNull(before[i]);
+ }
+
+ int[] newValues = new int[]{700, 701, 702, 703};
+ array.innerChunkWriter()
+ .put(new long[]{0, 0},
+ encodedInnerChunk("testInnerChunkWriteLeavesOtherBlobsByteIdentical/donor", newValues))
+ .flush();
+
+ // The kept inner chunks were copied through still encoded, so their bytes are unchanged.
+ for (int i = 0; i < untouched.length; i++) {
+ Assertions.assertEquals(before[i], array.readInnerChunkDirect(untouched[i]),
+ "inner chunk " + Arrays.toString(untouched[i]) + " changed");
+ }
+ assertArrayHolds(array, rampWithInnerChunk(new long[]{0, 0}, newValues));
+ }
+
+ @Test
+ public void testInnerChunkWriteBatchesShards() throws IOException, ZarrException {
+ Array array = shardedRampArray("testInnerChunkWriteBatchesShards/sharded", "end");
+
+ // Two inner chunks in shard (0,0) and two in shard (1,1), all in one flush.
+ long[][] coords = new long[][]{{0, 0}, {1, 1}, {2, 2}, {3, 3}};
+ int[][] values = new int[][]{{10, 11, 12, 13}, {20, 21, 22, 23}, {30, 31, 32, 33},
+ {40, 41, 42, 43}};
+
+ InnerChunkWriter writer = array.innerChunkWriter();
+ for (int i = 0; i < coords.length; i++) {
+ writer.put(coords[i],
+ encodedInnerChunk("testInnerChunkWriteBatchesShards/donor" + i, values[i]));
+ }
+ writer.flush();
+
+ int[] expected = new int[8 * 8];
+ Arrays.setAll(expected, p -> p);
+ for (int i = 0; i < coords.length; i++) {
+ for (int row = 0; row < 2; row++) {
+ for (int col = 0; col < 2; col++) {
+ expected[(int) (coords[i][0] * 2 + row) * 8 + (int) (coords[i][1] * 2 + col)] =
+ values[i][row * 2 + col];
+ }
+ }
+ }
+ assertArrayHolds(array, expected);
+
+ // The writer is reusable and does not re-apply the first batch.
+ int[] secondValues = new int[]{50, 51, 52, 53};
+ writer.put(new long[]{0, 1},
+ encodedInnerChunk("testInnerChunkWriteBatchesShards/donorSecond", secondValues))
+ .flush();
+ for (int row = 0; row < 2; row++) {
+ for (int col = 0; col < 2; col++) {
+ expected[row * 8 + 2 + col] = secondValues[row * 2 + col];
+ }
+ }
+ assertArrayHolds(array, expected);
+ }
+
+ @Test
+ public void testInnerChunkWriteCreatesShard() throws IOException, ZarrException {
+ StoreHandle storeHandle =
+ new FilesystemStore(TESTOUTPUT).resolve("testInnerChunkWriteCreatesShard", "sharded");
+ Array array = Array.create(storeHandle, Array.metadataBuilder()
+ .withShape(8, 8)
+ .withDataType(DataType.UINT32)
+ .withChunkShape(4, 4)
+ .withCodecs(c -> c.withSharding(new int[]{2, 2}, c1 -> c1.withBytes("LITTLE")))
+ .build());
+
+ // Shard (1,1) has never been written, so there is no stored object to splice into.
+ StoreHandle shardHandle = storeHandle.resolve(
+ array.metadata().chunkKeyEncoding().encodeChunkKey(new long[]{1, 1}));
+ Assertions.assertFalse(shardHandle.exists());
+
+ int[] newValues = new int[]{61, 62, 63, 64};
+ array.innerChunkWriter()
+ .put(new long[]{3, 3}, encodedInnerChunk("testInnerChunkWriteCreatesShard/donor", newValues))
+ .flush();
+
+ Assertions.assertTrue(shardHandle.exists());
+ // The written inner chunk reads back, and the rest of the new shard is still fill value.
+ Assertions.assertArrayEquals(newValues, (int[]) array.read(new long[]{6, 6}, new long[]{2, 2})
+ .get1DJavaArray(ucar.ma2.DataType.INT));
+ Assertions.assertNull(array.readInnerChunkDirect(new long[]{2, 2}));
+ Assertions.assertArrayEquals(new int[]{0, 0, 0, 0},
+ (int[]) array.read(new long[]{4, 4}, new long[]{2, 2}).get1DJavaArray(ucar.ma2.DataType.INT));
+ }
+
+ @Test
+ public void testInnerChunkWriteDelete() throws IOException, ZarrException {
+ Array array = shardedRampArray("testInnerChunkWriteDelete/sharded", "end");
+
+ ByteBuffer siblingBefore = array.readInnerChunkDirect(new long[]{0, 1});
+ array.innerChunkWriter().put(new long[]{0, 0}, null).flush();
+
+ // The removed inner chunk is marked absent in the shard index and reads back as fill value.
+ Assertions.assertNull(array.readInnerChunkDirect(new long[]{0, 0}));
+ assertArrayHolds(array, rampWithInnerChunk(new long[]{0, 0}, null));
+ Assertions.assertEquals(siblingBefore, array.readInnerChunkDirect(new long[]{0, 1}));
+
+ // Removing every remaining inner chunk removes the shard object itself.
+ StoreHandle shardHandle = array.storeHandle.resolve(
+ array.metadata().chunkKeyEncoding().encodeChunkKey(new long[]{0, 0}));
+ Assertions.assertTrue(shardHandle.exists());
+ array.innerChunkWriter()
+ .put(new long[]{0, 1}, null)
+ .put(new long[]{1, 0}, null)
+ .put(new long[]{1, 1}, null)
+ .flush();
+ Assertions.assertFalse(shardHandle.exists());
+ Assertions.assertNull(array.readChunkDirect(new long[]{0, 0}));
+ }
+
+ @Test
+ public void testInnerChunkWriteIndexLocationStart() throws IOException, ZarrException {
+ Array array = shardedRampArray("testInnerChunkWriteIndexLocationStart/sharded", "start");
+
+ ByteBuffer siblingBefore = array.readInnerChunkDirect(new long[]{1, 1});
+ int[] newValues = new int[]{800, 801, 802, 803};
+ array.innerChunkWriter()
+ .put(new long[]{0, 0},
+ encodedInnerChunk("testInnerChunkWriteIndexLocationStart/donor", newValues))
+ .flush();
+
+ // Offsets are shifted past the leading index, so a wrong shift would corrupt every blob.
+ assertArrayHolds(array, rampWithInnerChunk(new long[]{0, 0}, newValues));
+ Assertions.assertEquals(siblingBefore, array.readInnerChunkDirect(new long[]{1, 1}));
+ }
+
+ @Test
+ public void testInnerChunkWriteCompressedInnerCodecs() throws IOException, ZarrException {
+ // Inner chunks are gzipped, so replacing one changes its encoded byte length.
+ int[] testData = new int[8 * 8];
+ Arrays.setAll(testData, p -> p * 7919);
+
+ StoreHandle storeHandle = new FilesystemStore(TESTOUTPUT)
+ .resolve("testInnerChunkWriteCompressedInnerCodecs", "sharded");
+ Array array = Array.create(storeHandle, Array.metadataBuilder()
+ .withShape(8, 8)
+ .withDataType(DataType.UINT32)
+ .withChunkShape(4, 4)
+ .withCodecs(c -> c.withSharding(new int[]{2, 2}, c1 -> c1.withBytes("LITTLE").withGzip()))
+ .build());
+ array.write(ucar.ma2.Array.factory(ucar.ma2.DataType.UINT, new int[]{8, 8}, testData));
+
+ // A uniform inner chunk compresses better than the ramp, so the new blob is shorter.
+ Array donor = Array.create(
+ new FilesystemStore(TESTOUTPUT).resolve("testInnerChunkWriteCompressedInnerCodecs", "donor"),
+ Array.metadataBuilder()
+ .withShape(2, 2)
+ .withDataType(DataType.UINT32)
+ .withChunkShape(2, 2)
+ .withCodecs(c -> c.withBytes("LITTLE").withGzip())
+ .build());
+ int[] newValues = new int[]{5, 5, 5, 5};
+ donor.writeChunk(new long[]{0, 0},
+ ucar.ma2.Array.factory(ucar.ma2.DataType.UINT, new int[]{2, 2}, newValues));
+ ByteBuffer newBytes = donor.readChunkDirect(new long[]{0, 0});
+
+ ByteBuffer siblingBefore = array.readInnerChunkDirect(new long[]{0, 1});
+ int oldShardLength = array.readChunkDirect(new long[]{0, 0}).remaining();
+ int oldBlobLength = array.readInnerChunkDirect(new long[]{0, 0}).remaining();
+ Assertions.assertTrue(newBytes.remaining() < oldBlobLength,
+ "expected the replacement blob to be shorter than the one it replaces");
+
+ array.innerChunkWriter().put(new long[]{0, 0}, newBytes).flush();
+
+ // The shard shrank by exactly the difference, the kept blobs are untouched, and the data decodes.
+ Assertions.assertEquals(oldShardLength - (oldBlobLength - newBytes.remaining()),
+ array.readChunkDirect(new long[]{0, 0}).remaining());
+ Assertions.assertEquals(siblingBefore, array.readInnerChunkDirect(new long[]{0, 1}));
+
+ int[] expected = new int[8 * 8];
+ Arrays.setAll(expected, p -> p * 7919);
+ for (int row = 0; row < 2; row++) {
+ for (int col = 0; col < 2; col++) {
+ expected[row * 8 + col] = newValues[row * 2 + col];
+ }
+ }
+ assertArrayHolds(array, expected);
+ }
+
+ @Test
+ public void testInnerChunkWriteNestedSharding() throws IOException, ZarrException {
+ // 8x8 array in one 8x8 shard of 4x4 shards of 2x2 inner chunks.
+ int[] testData = new int[8 * 8];
+ Arrays.setAll(testData, p -> p);
+
+ StoreHandle storeHandle =
+ new FilesystemStore(TESTOUTPUT).resolve("testInnerChunkWriteNestedSharding", "sharded");
+ Array array = Array.create(storeHandle, Array.metadataBuilder()
+ .withShape(8, 8)
+ .withDataType(DataType.UINT32)
+ .withChunkShape(8, 8)
+ .withCodecs(c -> c.withSharding(new int[]{4, 4},
+ c1 -> c1.withSharding(new int[]{2, 2}, c2 -> c2.withBytes("LITTLE"))))
+ .build());
+ array.write(ucar.ma2.Array.factory(ucar.ma2.DataType.UINT, new int[]{8, 8}, testData));
+
+ // (1,3) sits in the same nested shard as the target (1,2); (3,3) sits in a different one.
+ ByteBuffer sameNestedShardBefore = array.readInnerChunkDirect(new long[]{1, 3});
+ ByteBuffer otherNestedShardBefore = array.readInnerChunkDirect(new long[]{3, 3});
+
+ int[] newValues = new int[]{600, 601, 602, 603};
+ array.innerChunkWriter()
+ .put(new long[]{1, 2}, encodedInnerChunk("testInnerChunkWriteNestedSharding/donor", newValues))
+ .flush();
+
+ // The nested shard was rebuilt and spliced back into the outer shard, inside out.
+ assertArrayHolds(array, rampWithInnerChunk(new long[]{1, 2}, newValues));
+ Assertions.assertEquals(sameNestedShardBefore, array.readInnerChunkDirect(new long[]{1, 3}));
+ Assertions.assertEquals(otherNestedShardBefore, array.readInnerChunkDirect(new long[]{3, 3}));
+ }
+
+ @Test
+ public void testInnerChunkWriteIdempotent() throws IOException, ZarrException {
+ Array array = shardedRampArray("testInnerChunkWriteIdempotent/sharded", "end");
+ ByteBuffer newBytes =
+ encodedInnerChunk("testInnerChunkWriteIdempotent/donor", new int[]{1, 2, 3, 4});
+
+ array.innerChunkWriter().put(new long[]{1, 1}, newBytes).flush();
+ ByteBuffer afterFirst = array.readChunkDirect(new long[]{0, 0});
+
+ // Rebuilding a shard from the same staged bytes is byte-identical, so a retry is safe.
+ array.innerChunkWriter().put(new long[]{1, 1}, newBytes).flush();
+ Assertions.assertEquals(afterFirst, array.readChunkDirect(new long[]{0, 0}));
+ }
+
+ @Test
+ public void testInnerChunkWriteBoundaryShard() throws IOException, ZarrException {
+ // 6x6 array of 4x4 shards: the right and bottom shards overhang the array.
+ StoreHandle storeHandle =
+ new FilesystemStore(TESTOUTPUT).resolve("testInnerChunkWriteBoundaryShard", "sharded");
+ Array array = Array.create(storeHandle, Array.metadataBuilder()
+ .withShape(6, 6)
+ .withDataType(DataType.UINT32)
+ .withChunkShape(4, 4)
+ .withCodecs(c -> c.withSharding(new int[]{2, 2}, c1 -> c1.withBytes("LITTLE")))
+ .build());
+ int[] testData = new int[6 * 6];
+ Arrays.setAll(testData, p -> p);
+ array.write(ucar.ma2.Array.factory(ucar.ma2.DataType.UINT, new int[]{6, 6}, testData));
+
+ // Inner chunk (2,2) covers rows 4-5, cols 4-5, entirely inside the array.
+ int[] newValues = new int[]{91, 92, 93, 94};
+ array.innerChunkWriter()
+ .put(new long[]{2, 2}, encodedInnerChunk("testInnerChunkWriteBoundaryShard/donor", newValues))
+ .flush();
+ Assertions.assertArrayEquals(newValues, (int[]) array.read(new long[]{4, 4}, new long[]{2, 2})
+ .get1DJavaArray(ucar.ma2.DataType.INT));
+
+ // Inner chunk (3,3) starts at row 6, outside the array, and is rejected like on the read side.
+ assertThrows(ZarrException.class, () -> array.readInnerChunkDirect(new long[]{3, 3}));
+ assertThrows(ZarrException.class,
+ () -> array.innerChunkWriter().put(new long[]{3, 3}, ByteBuffer.allocate(16)));
+ }
+
+ @Test
+ public void testInnerChunkWriteRejects() throws IOException, ZarrException {
+ Array array = shardedRampArray("testInnerChunkWriteRejects/sharded", "end");
+ ByteBuffer bytes = encodedInnerChunk("testInnerChunkWriteRejects/donor", new int[]{1, 2, 3, 4});
+
+ // Wrong rank, out-of-domain and negative coordinates are rejected before any store access.
+ assertThrows(IllegalArgumentException.class,
+ () -> array.innerChunkWriter().put(new long[]{0, 0, 0}, bytes));
+ assertThrows(ZarrException.class, () -> array.innerChunkWriter().put(new long[]{4, 0}, bytes));
+ assertThrows(ZarrException.class, () -> array.innerChunkWriter().put(new long[]{0, -1}, bytes));
+ // An empty buffer would be stored as a present but zero-length inner chunk.
+ assertThrows(ZarrException.class,
+ () -> array.innerChunkWriter().put(new long[]{0, 0}, ByteBuffer.allocate(0)));
+
+ // A shard that cannot be parsed is never overwritten.
+ array.writeChunkDirect(new long[]{0, 0}, ByteBuffer.wrap(new byte[]{1, 2, 3}));
+ assertThrows(ZarrException.class,
+ () -> array.innerChunkWriter().put(new long[]{0, 0}, bytes).flush());
+ Assertions.assertEquals(3, array.readChunkDirect(new long[]{0, 0}).remaining());
+ }
+
+ @Test
+ public void testInnerChunkWriteUnsharded() throws IOException, ZarrException {
+ StoreHandle storeHandle =
+ new FilesystemStore(TESTOUTPUT).resolve("testInnerChunkWriteUnsharded", "plain");
+ Array array = Array.create(storeHandle, Array.metadataBuilder()
+ .withShape(4, 4)
+ .withDataType(DataType.UINT32)
+ .withChunkShape(2, 2)
+ .withCodecs(c -> c.withBytes("LITTLE"))
+ .build());
+
+ // Without sharding the inner chunk grid is the chunk grid, so this stores the chunk verbatim.
+ int[] newValues = new int[]{7, 8, 9, 10};
+ ByteBuffer bytes = encodedInnerChunk("testInnerChunkWriteUnsharded/donor", newValues);
+ array.writeInnerChunkDirect(new long[]{1, 1}, bytes);
+ Assertions.assertEquals(bytes, array.readChunkDirect(new long[]{1, 1}));
+ Assertions.assertArrayEquals(newValues, (int[]) array.readChunk(new long[]{1, 1})
+ .get1DJavaArray(ucar.ma2.DataType.INT));
+
+ array.writeInnerChunkDirect(new long[]{1, 1}, null);
+ Assertions.assertNull(array.readChunkDirect(new long[]{1, 1}));
+ }
+
+ @Test
+ public void testInnerChunkWriteDegenerateSharding() throws IOException, ZarrException {
+ // Sharding where the inner chunk shape equals the chunk shape: one inner chunk per shard, but
+ // the stored object still carries a shard index, so the inner chunk is not the whole object.
+ StoreHandle storeHandle = new FilesystemStore(TESTOUTPUT)
+ .resolve("testInnerChunkWriteDegenerateSharding", "sharded");
+ Array array = Array.create(storeHandle, Array.metadataBuilder()
+ .withShape(4, 4)
+ .withDataType(DataType.UINT32)
+ .withChunkShape(2, 2)
+ .withCodecs(c -> c.withSharding(new int[]{2, 2}, c1 -> c1.withBytes("LITTLE")))
+ .build());
+ int[] testData = new int[4 * 4];
+ Arrays.setAll(testData, p -> p);
+ array.write(ucar.ma2.Array.factory(ucar.ma2.DataType.UINT, new int[]{4, 4}, testData));
+
+ // The inner chunk bytes are shorter than the shard object, which also holds the index.
+ ByteBuffer innerChunkBytes = array.readInnerChunkDirect(new long[]{0, 0});
+ Assertions.assertNotNull(innerChunkBytes);
+ Assertions.assertTrue(
+ innerChunkBytes.remaining() < array.readChunkDirect(new long[]{0, 0}).remaining());
+
+ int[] newValues = new int[]{11, 12, 13, 14};
+ array.writeInnerChunkDirect(new long[]{1, 1},
+ encodedInnerChunk("testInnerChunkWriteDegenerateSharding/donor", newValues));
+ Assertions.assertArrayEquals(newValues, (int[]) array.read(new long[]{2, 2}, new long[]{2, 2})
+ .get1DJavaArray(ucar.ma2.DataType.INT));
+ }
+
+ @Test
+ public void testInnerChunkWriteMatchesNormalWrite() throws IOException, ZarrException {
+ // The same logical content, written once through the codec pipeline and once by splicing
+ // encoded inner chunks, must read back identically.
+ int[] newValues = new int[]{321, 322, 323, 324};
+
+ Array spliced = shardedRampArray("testInnerChunkWriteMatchesNormalWrite/spliced", "end");
+ spliced.innerChunkWriter()
+ .put(new long[]{2, 1},
+ encodedInnerChunk("testInnerChunkWriteMatchesNormalWrite/donor", newValues))
+ .flush();
+
+ Array normal = shardedRampArray("testInnerChunkWriteMatchesNormalWrite/normal", "end");
+ normal.write(new long[]{4, 2},
+ ucar.ma2.Array.factory(ucar.ma2.DataType.UINT, new int[]{2, 2}, newValues));
+
+ Assertions.assertArrayEquals((int[]) normal.read().get1DJavaArray(ucar.ma2.DataType.INT),
+ (int[]) spliced.read().get1DJavaArray(ucar.ma2.DataType.INT));
+ }
}