Feat/read write direct

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.
+     * <p>
+     * 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.
+     * <p>
+     * <b>Unsafe:</b> 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.
+     * <p>
+     * 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,152 @@ 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.
+     * <p>
+     * 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.
+     * <p>
+     * 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.
+     * <p>
+     * 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.
+     * <p>
+     * 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.
+     * <p>
+     * The inner chunk is addressed level by level: first the coordinates of the shard on the array's
+     * chunk grid, then the coordinates of the inner chunk within that shard, and, if shards are
+     * nested, one more entry per nesting level. For example, in an 8x8 array with 4x4 shards of 2x2
+     * inner chunks, the inner chunk covering rows 6-7 and columns 2-3 is
+     * {@code readInnerChunkDirect(new long[]{1, 0}, new long[]{1, 1})}: shard (1,0), inner chunk (1,1)
+     * within it. For an unsharded array only the chunk coordinates are given, and this is equivalent to
+     * {@link #readChunkDirect(long[])}.
+     *
+     * @param innerChunkCoords The chunk coordinates, followed by one entry per sharding level.
+     * @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 IllegalArgumentException if the number of entries does not match the sharding depth, or
+     *                                  an entry has the wrong rank
+     * @throws ZarrException            if the requested inner chunk is outside the shard or the
+     *                                  array's domain
+     */
+    @Nullable
+    public ByteBuffer readInnerChunkDirect(long[]... innerChunkCoords) throws ZarrException {
+        checkInnerChunkCoords(innerChunkCoords);
+        final String[] chunkKeys = metadata().chunkKeyEncoding().encodeChunkKey(innerChunkCoords[0]);
+        final StoreHandle chunkHandle = storeHandle.resolve(chunkKeys);
+
+        if (innerChunkCoords.length == 1) {
+            return chunkHandle.read();
+        }
+        return codecPipeline.readInnerChunkDirect(chunkHandle,
+                Arrays.copyOfRange(innerChunkCoords, 1, innerChunkCoords.length));
+    }
+
+    /**
+     * Checks that {@code innerChunkCoords} holds chunk coordinates followed by one entry per sharding
+     * level, that every entry is within its shard, and that the addressed inner chunk starts inside
+     * the array's domain.
+     *
+     * @throws IllegalArgumentException if the number of entries does not match the sharding depth, or
+     *                                  an entry has the wrong rank
+     * @throws ZarrException            if the inner chunk is outside the shard or the array's domain
+     */
+    void checkInnerChunkCoords(long[][] innerChunkCoords) throws ZarrException {
+        final ArrayMetadata metadata = metadata();
+        final int ndim = metadata.ndim();
+        final int[][] innerChunkShapes = codecPipeline.innerChunkShapes();
+        if (innerChunkCoords.length != innerChunkShapes.length + 1) {
+            throw new IllegalArgumentException(
+                    "'innerChunkCoords' needs " + (innerChunkShapes.length + 1) + " entries, the chunk "
+                            + "coordinates followed by one per sharding level, but has "
+                            + innerChunkCoords.length + ".");
+        }
+
+        // Walk the levels from the stored chunk inwards, accumulating where the inner chunk starts.
+        final long[] start = new long[ndim];
+        int[] parentShape = null;
+        for (int level = 0; level < innerChunkCoords.length; level++) {
+            final long[] coords = innerChunkCoords[level];
+            if (coords == null || coords.length != ndim) {
+                throw new IllegalArgumentException(
+                        "'innerChunkCoords' entries need to have rank '" + ndim + "'.");
+            }
+            final int[] shape = level == 0 ? metadata.chunkShape() : innerChunkShapes[level - 1];
+            for (int dimIdx = 0; dimIdx < ndim; dimIdx++) {
+                if (coords[dimIdx] < 0
+                        || (parentShape != null && coords[dimIdx] >= parentShape[dimIdx] / shape[dimIdx])) {
+                    throw new ZarrException("Attempting to access an inner chunk outside of its shard.");
+                }
+                start[dimIdx] += coords[dimIdx] * shape[dimIdx];
+            }
+            parentShape = shape;
+        }
+        for (int dimIdx = 0; dimIdx < ndim; dimIdx++) {
+            if (start[dimIdx] >= metadata.shape[dimIdx]) {
+                throw new ZarrException("Attempting to access data outside of the array's domain.");
+            }
+        }
+    }
+
+    /**
+     * Opens a writer for storing already-encoded inner chunks into one shard of this array without
+     * decoding or re-encoding anything. All staged inner chunks are applied together, so the shard is
+     * rebuilt once no matter how many inner chunks change. Use it with try-with-resources so that the
+     * staged inner chunks are flushed when the block ends:
+     * <pre>{@code
+     * try (InnerChunkWriter writer = array.writeInnerChunkDirect(new long[]{1, 0})) {
+     *     writer.put(new long[]{1, 1}, innerChunkBytes);
+     * }
+     * }</pre>
+     * This is the encoded-bytes counterpart of {@link #readInnerChunkDirect(long[]...)}. See
+     * {@link InnerChunkWriter} for the safety contract.
+     *
+     * @param shardCoords The coordinates of the shard on the array's chunk grid.
+     * @throws IllegalArgumentException if {@code shardCoords} has the wrong rank
+     * @throws ZarrException            if the shard is outside the array's domain
+     */
+    public InnerChunkWriter writeInnerChunkDirect(long[] shardCoords) throws ZarrException {
+        return new InnerChunkWriter(this, shardCoords);
+    }
+
     /**
      * Deletes chunks that are completely outside the new shape and trims boundary chunks.
      *