Merge pull request #6 from scalableminds/check-sharding-bounds

Check sharding bounds

Author: normanrz

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

@@ -10,9 +10,12 @@
 import dev.zarr.zarrjava.v3.chunkgrid.RegularChunkGrid;
 import dev.zarr.zarrjava.v3.chunkkeyencoding.ChunkKeyEncoding;
 import dev.zarr.zarrjava.v3.codec.Codec;
+import dev.zarr.zarrjava.v3.codec.core.ShardingIndexedCodec;
+
 import java.nio.ByteBuffer;
 import java.util.Arrays;
 import java.util.Map;
+import java.util.Optional;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
@@ -91,6 +94,35 @@ public ArrayMetadata(
           "Expected node type '" + this.nodeType + "', got '" + nodeType + "'.");
     }
 
+    if (chunkGrid instanceof RegularChunkGrid) {
+      int[] chunkShape = ((RegularChunkGrid) chunkGrid).configuration.chunkShape;
+      if (shape.length != chunkShape.length) {
+        throw new ZarrException("Shape (ndim=" + shape.length + ") and chunk shape (ndim=" +
+            chunkShape.length + ") need to have the same number of dimensions.");
+      }
+      for (int i = 0; i < shape.length; i++) {
+        if (shape[i] < chunkShape[i]) {
+          throw new ZarrException("Shape " + Arrays.toString(shape) + " can not contain chunk shape "
+              + Arrays.toString(chunkShape));
+        }
+      }
+
+      Optional<Codec> shardingCodec = getShardingIndexedCodec(codecs);
+      int[] outerChunkShape = chunkShape;
+      while (shardingCodec.isPresent()) {
+        ShardingIndexedCodec.Configuration shardingConfig = ((ShardingIndexedCodec) shardingCodec.get()).configuration;
+        int[] innerChunkShape = shardingConfig.chunkShape;
+        if (outerChunkShape.length != innerChunkShape.length)
+          throw new ZarrException("Sharding dimensions mismatch of outer chunk shape " + Arrays.toString(outerChunkShape) + " and inner chunk shape" + Arrays.toString(innerChunkShape));
+        for (int i = 0; i < outerChunkShape.length; i++) {
+          if (outerChunkShape[i] < innerChunkShape[i])
+            throw new ZarrException("Sharding outer chunk shape " + Arrays.toString(outerChunkShape) + " can not contain inner chunk shape " + Arrays.toString(innerChunkShape));
+        }
+        outerChunkShape = innerChunkShape;
+        shardingCodec = getShardingIndexedCodec(shardingConfig.codecs);
+      }
+    }
+
     this.shape = shape;
     this.dataType = dataType;
     this.chunkGrid = chunkGrid;
@@ -227,6 +259,10 @@ public int ndim() {
     return shape.length;
   }
 
+  public static Optional<Codec> getShardingIndexedCodec(Codec[] codecs) {
+    return Arrays.stream(codecs).filter(codec -> codec instanceof ShardingIndexedCodec).findFirst();
+  }
+
   public int[] chunkShape() {
     return ((RegularChunkGrid) this.chunkGrid).configuration.chunkShape;
   }

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

@@ -11,6 +11,9 @@
 import dev.zarr.zarrjava.v3.codec.CodecBuilder;
 import dev.zarr.zarrjava.v3.codec.core.BytesCodec;
 import dev.zarr.zarrjava.v3.codec.core.BytesCodec.Endian;
+import dev.zarr.zarrjava.v3.codec.core.ShardingIndexedCodec;
+
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.function.Function;
@@ -133,12 +136,8 @@ public ArrayMetadata build() throws ZarrException {
     if (chunkGrid == null) {
       throw new ZarrException("Chunk grid needs to be provided. Please call `.withChunkShape`.");
     }
-    if (chunkGrid instanceof RegularChunkGrid
-        && shape.length != ((RegularChunkGrid) chunkGrid).configuration.chunkShape.length) {
-      throw new ZarrException("Shape (ndim=" + shape.length + ") and chunk shape (ndim=" +
-          ((RegularChunkGrid) chunkGrid).configuration.chunkShape.length +
-          ") need to have the same number of dimensions.");
-    }
+
+
     return new ArrayMetadata(shape, dataType, chunkGrid, chunkKeyEncoding, fillValue, codecs,
         dimensionNames,
         attributes

src/test/java/dev/zarr/zarrjava/ZarrTest.java

@@ -19,6 +19,7 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.CsvSource;
+import org.junit.jupiter.params.provider.MethodSource;
 import org.junit.jupiter.params.provider.ValueSource;
 import ucar.ma2.MAMath;
 
@@ -230,6 +231,60 @@ public void testWriteReadWithZarrita(String codec, String codecParam) throws Exc
         assert exitCode == 0;
     }
 
+    static Stream<int[]> invalidchunkSizes() {
+        return Stream.of(
+            new int[] {1} ,
+            new int[] {1, 1, 1},
+            new int[] {5, 1},
+            new int[] {1, 5}
+        );
+    }
+
+    @ParameterizedTest
+    @MethodSource("invalidchunkSizes")
+    public void testCheckInvalidChunkBounds(int[] chunkSize) throws Exception {
+        long[] shape = new long[] {4, 4};
+
+        StoreHandle storeHandle = new FilesystemStore(TESTOUTPUT).resolve("invalid_chunksize");
+        ArrayMetadataBuilder builder = Array.metadataBuilder()
+            .withShape(shape)
+            .withDataType(DataType.UINT32)
+            .withChunkShape(chunkSize);
+
+        assertThrows(ZarrException.class, builder::build);
+    }
+
+    @ParameterizedTest
+    @ValueSource(strings = {"large", "small", "nested", "wrong dims", "correct"})
+    public void testCheckShardingBounds(String scenario) throws Exception {
+        long[] shape = new long[] {4, 4};
+        int[] shardSize = new int[] {2, 2};
+        int[] chunkSize = new int[] {2, 2};
+
+        if (scenario.equals("large"))
+            shardSize = new int[] {8, 8};
+        if (scenario.equals("small"))
+            shardSize = new int[] {1, 1};
+        if (scenario.equals("wrong dims"))
+            shardSize = new int[] {1};
+        StoreHandle storeHandle = new FilesystemStore(TESTOUTPUT).resolve("illegal_shardsize");
+        ArrayMetadataBuilder builder = Array.metadataBuilder()
+            .withShape(shape)
+            .withDataType(DataType.UINT32).withChunkShape(shardSize);
+
+        if (scenario.equals("nested")) {
+            int[] nestedChunkSize = new int[]{4, 4};
+            builder = builder.withCodecs(c -> c.withSharding(chunkSize, c1 -> c1.withSharding(nestedChunkSize, c2 -> c2.withBytes("LITTLE"))));
+        } else {
+            builder = builder.withCodecs(c -> c.withSharding(chunkSize, c1 -> c1.withBytes("LITTLE")));
+        }
+        if (scenario.equals("correct")){
+            builder.build();
+        }else{
+            assertThrows(ZarrException.class, builder::build);
+        }
+    }
+
     @ParameterizedTest
     @CsvSource({"0,true", "0,false", "5, true", "5, false"})
     public void testZstdCodecReadWrite(int clevel, boolean checksum) throws ZarrException, IOException {