From 7f158fbef50328f0cfd70be414c5aad507ee0bf5 Mon Sep 17 00:00:00 2001 From: Konstantin Date: Tue, 4 Aug 2026 17:10:06 +0200 Subject: [PATCH 1/4] add direct array read / Write --- .../java/dev/zarr/zarrjava/core/Array.java | 68 +++++++++++++++++++ .../java/dev/zarr/zarrjava/ZarrV3Test.java | 39 +++++++++++ 2 files changed, 107 insertions(+) diff --git a/src/main/java/dev/zarr/zarrjava/core/Array.java b/src/main/java/dev/zarr/zarrjava/core/Array.java index 67461483..92d96569 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,35 @@ 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). + * + * @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(); + } + /** * Deletes chunks that are completely outside the new shape and trims boundary chunks. * diff --git a/src/test/java/dev/zarr/zarrjava/ZarrV3Test.java b/src/test/java/dev/zarr/zarrjava/ZarrV3Test.java index 614a2b90..bb1da4a8 100644 --- a/src/test/java/dev/zarr/zarrjava/ZarrV3Test.java +++ b/src/test/java/dev/zarr/zarrjava/ZarrV3Test.java @@ -27,6 +27,7 @@ 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 +1053,42 @@ 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)); + } } From 7173f31defceb57d75061e2fea1cd137bdd054ba Mon Sep 17 00:00:00 2001 From: Konstantin Date: Wed, 12 Aug 2026 17:30:30 +0200 Subject: [PATCH 2/4] readInnerChunksdirect now works --- .../java/dev/zarr/zarrjava/core/Array.java | 66 ++++++++++- .../zarrjava/core/codec/ArrayBytesCodec.java | 21 ++++ .../zarrjava/core/codec/CodecPipeline.java | 33 ++++++ .../v3/codec/core/ShardingIndexedCodec.java | 71 ++++++++++++ .../java/dev/zarr/zarrjava/ZarrV3Test.java | 108 ++++++++++++++++++ 5 files changed, 298 insertions(+), 1 deletion(-) diff --git a/src/main/java/dev/zarr/zarrjava/core/Array.java b/src/main/java/dev/zarr/zarrjava/core/Array.java index 92d96569..43e03a6e 100644 --- a/src/main/java/dev/zarr/zarrjava/core/Array.java +++ b/src/main/java/dev/zarr/zarrjava/core/Array.java @@ -229,7 +229,8 @@ public ucar.ma2.Array readChunk(long[] chunkCoords) throws ZarrException { * 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). + * 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. @@ -249,6 +250,69 @@ public ByteBuffer readChunkDirect(long[] chunkCoords) throws ZarrException { 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 { + 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(); + + // Split the inner chunk coordinates into the coordinates of the stored chunk holding them and + // the coordinates of the inner chunk within that stored chunk. + 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 read data outside of the array's domain."); + } + final int innerChunksPerChunk = chunkShape[dimIdx] / innerChunkShape[dimIdx]; + chunkCoords[dimIdx] = innerChunkCoords[dimIdx] / innerChunksPerChunk; + coordsInChunk[dimIdx] = innerChunkCoords[dimIdx] % innerChunksPerChunk; + } + + final String[] chunkKeys = metadata.chunkKeyEncoding().encodeChunkKey(chunkCoords); + final StoreHandle chunkHandle = storeHandle.resolve(chunkKeys); + + if (Arrays.equals(innerChunkShape, chunkShape)) { + return chunkHandle.read(); + } + return codecPipeline.readInnerChunkEncoded(chunkHandle, coordsInChunk); + } + /** * Deletes chunks that are completely outside the new shape and trims boundary chunks. * 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..d3ed828d 100644 --- a/src/main/java/dev/zarr/zarrjava/core/codec/ArrayBytesCodec.java +++ b/src/main/java/dev/zarr/zarrjava/core/codec/ArrayBytesCodec.java @@ -4,6 +4,7 @@ import dev.zarr.zarrjava.store.StoreHandle; import ucar.ma2.Array; +import javax.annotation.Nullable; import java.nio.ByteBuffer; public abstract class ArrayBytesCodec extends AbstractCodec { @@ -23,6 +24,26 @@ 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. + */ + 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; } } 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..1be859d8 100644 --- a/src/main/java/dev/zarr/zarrjava/core/codec/CodecPipeline.java +++ b/src/main/java/dev/zarr/zarrjava/core/codec/CodecPipeline.java @@ -6,6 +6,7 @@ import ucar.ma2.Array; import javax.annotation.Nonnull; +import javax.annotation.Nullable; import java.nio.ByteBuffer; import java.util.Arrays; @@ -83,6 +84,38 @@ public boolean supportsPartialDecode() { return codecs.length == 1 && codecs[0] instanceof ArrayBytesCodec.WithPartialDecode; } + /** + * The shape of the smallest unit this pipeline encodes independently. For a sharded array this + * is the inner chunk shape of the sharding codec; otherwise it 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); + } + @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..0a4e8721 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,6 +18,7 @@ import ucar.ma2.InvalidRangeException; import javax.annotation.Nonnull; +import javax.annotation.Nullable; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.Arrays; @@ -242,6 +243,76 @@ 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 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 bb1da4a8..c74fecc1 100644 --- a/src/test/java/dev/zarr/zarrjava/ZarrV3Test.java +++ b/src/test/java/dev/zarr/zarrjava/ZarrV3Test.java @@ -1090,5 +1090,113 @@ public void testDirectChunkReadWrite() throws IOException, ZarrException { // 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}); } } From 10d109140cc26610e6dd850a94916fb498a5a0b8 Mon Sep 17 00:00:00 2001 From: Konstantin Date: Thu, 13 Aug 2026 16:46:50 +0200 Subject: [PATCH 3/4] Clarify innerChunkShape docs for nested shards Co-Authored-By: Claude Opus 5 (1M context) --- .../dev/zarr/zarrjava/core/codec/ArrayBytesCodec.java | 4 ++++ .../dev/zarr/zarrjava/core/codec/CodecPipeline.java | 10 ++++++++-- 2 files changed, 12 insertions(+), 2 deletions(-) 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 d3ed828d..a13e31a1 100644 --- a/src/main/java/dev/zarr/zarrjava/core/codec/ArrayBytesCodec.java +++ b/src/main/java/dev/zarr/zarrjava/core/codec/ArrayBytesCodec.java @@ -28,6 +28,10 @@ protected abstract Array decodePartial( /** * 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(); 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 1be859d8..f45f3e7d 100644 --- a/src/main/java/dev/zarr/zarrjava/core/codec/CodecPipeline.java +++ b/src/main/java/dev/zarr/zarrjava/core/codec/CodecPipeline.java @@ -85,8 +85,14 @@ public boolean supportsPartialDecode() { } /** - * The shape of the smallest unit this pipeline encodes independently. For a sharded array this - * is the inner chunk shape of the sharding codec; otherwise it is the chunk shape itself. + * 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()) { From a85e947530172cff8370e348dd7fe84dc1fa5d40 Mon Sep 17 00:00:00 2001 From: Konstantin Date: Wed, 19 Aug 2026 17:02:24 +0200 Subject: [PATCH 4/4] adjusted writing mechanism --- .../java/dev/zarr/zarrjava/core/Array.java | 69 ++- .../zarr/zarrjava/core/InnerChunkWriter.java | 180 +++++++ .../zarrjava/core/codec/ArrayBytesCodec.java | 43 ++ .../zarrjava/core/codec/CodecPipeline.java | 30 ++ .../v3/codec/core/ShardingIndexedCodec.java | 144 ++++++ .../java/dev/zarr/zarrjava/ZarrV3Test.java | 438 ++++++++++++++++++ 6 files changed, 894 insertions(+), 10 deletions(-) create mode 100644 src/main/java/dev/zarr/zarrjava/core/InnerChunkWriter.java diff --git a/src/main/java/dev/zarr/zarrjava/core/Array.java b/src/main/java/dev/zarr/zarrjava/core/Array.java index 43e03a6e..0b309c45 100644 --- a/src/main/java/dev/zarr/zarrjava/core/Array.java +++ b/src/main/java/dev/zarr/zarrjava/core/Array.java @@ -281,7 +281,26 @@ public int[] innerChunkShape() { */ @Nullable public ByteBuffer readInnerChunkDirect(long[] innerChunkCoords) throws ZarrException { - ArrayMetadata metadata = metadata(); + 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( @@ -290,27 +309,57 @@ public ByteBuffer readInnerChunkDirect(long[] innerChunkCoords) throws ZarrExcep final int[] chunkShape = metadata.chunkShape(); final int[] innerChunkShape = innerChunkShape(); - // Split the inner chunk coordinates into the coordinates of the stored chunk holding them and - // the coordinates of the inner chunk within that stored chunk. 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 read data outside of the array's domain."); + 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}; + } - final String[] chunkKeys = metadata.chunkKeyEncoding().encodeChunkKey(chunkCoords); - final StoreHandle chunkHandle = storeHandle.resolve(chunkKeys); + /** + * 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); + } - if (Arrays.equals(innerChunkShape, chunkShape)) { - return chunkHandle.read(); - } - return codecPipeline.readInnerChunkEncoded(chunkHandle, coordsInChunk); + /** + * 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(); } /** 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 a13e31a1..118987f4 100644 --- a/src/main/java/dev/zarr/zarrjava/core/codec/ArrayBytesCodec.java +++ b/src/main/java/dev/zarr/zarrjava/core/codec/ArrayBytesCodec.java @@ -6,6 +6,7 @@ import javax.annotation.Nullable; import java.nio.ByteBuffer; +import java.util.List; public abstract class ArrayBytesCodec extends AbstractCodec { @@ -48,6 +49,48 @@ protected abstract Array decodePartial( 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 f45f3e7d..f865e71c 100644 --- a/src/main/java/dev/zarr/zarrjava/core/codec/CodecPipeline.java +++ b/src/main/java/dev/zarr/zarrjava/core/codec/CodecPipeline.java @@ -9,6 +9,7 @@ import javax.annotation.Nullable; import java.nio.ByteBuffer; import java.util.Arrays; +import java.util.List; public class CodecPipeline { @@ -122,6 +123,35 @@ public ByteBuffer 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 0a4e8721..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 @@ -22,7 +22,9 @@ 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 { @@ -313,6 +315,148 @@ private ByteBuffer readInnerChunkEncoded(DataProvider dataProvider, long[] inner 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 c74fecc1..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,6 +26,7 @@ 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; @@ -1199,4 +1201,440 @@ public void testDirectInnerChunkReadNestedSharding() throws IOException, ZarrExc 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)); + } }