Add custom sharding option

glencoesoftware · Jul 30, 2024 · 5620d04 · 5620d04
1 parent 3116a33
commit 5620d04
Show file tree

Hide file tree

Showing 2 changed files with 75 additions and 5 deletions.
diff --git a/src/main/java/com/glencoesoftware/zarr/Convert.java b/src/main/java/com/glencoesoftware/zarr/Convert.java
@@ -67,6 +67,7 @@ public class Convert implements Callable<Integer> {
   private boolean writeV2;
 
   private ShardConfiguration shardConfig;
+  private int[] requestedShard; // the requested size for custom sharding
   private String[] codecs;
 
   /**
@@ -136,8 +137,12 @@ public void setSharding(String shard) {
         shardConfig = Enum.valueOf(ShardConfiguration.class, shard);
       }
       catch (IllegalArgumentException e) {
-        // TODO
         shardConfig = ShardConfiguration.CUSTOM;
+        String[] shardSize = shard.split(",");
+        requestedShard = new int[shardSize.length];
+        for (int i=0; i<shardSize.length; i++) {
+          requestedShard[i] = Integer.parseInt(shardSize[i]);
+        }
       }
     }
   }
@@ -318,7 +323,7 @@ public void convertToV3() throws Exception {
                   }
                   break;
                 case CUSTOM:
-                  // TODO
+                  chunkSizes = requestedShard;
                   break;
               }
 

diff --git a/src/test/java/com/glencoesoftware/zarr/test/ConversionTest.java b/src/test/java/com/glencoesoftware/zarr/test/ConversionTest.java
@@ -317,7 +317,7 @@ public void testCodecs() throws Exception {
   }
 
   /**
-   * Test different sharding options
+   * Test different default sharding options
    */
   @Test
   public void testSharding() throws Exception {
@@ -432,12 +432,77 @@ public void test3DSharding() throws Exception {
     }
   }
 
+  /**
+   * Test different custom sharding options
+   */
+  @Test
+  public void testCustomSharding() throws Exception {
+    input = fake("sizeX", "4096", "sizeY", "4096", "sizeT", "2", "sizeC", "3");
+    assertBioFormats2Raw();
+
+    String[] shardOptions = new String[] {
+      "1,1,1,2048,2048",
+      "2,1,1,1024,1024",
+      "1,3,1,4096,4096"
+    };
+    int[][] shardSizes = new int[][] {
+      {1, 1, 1, 2048, 2048},
+      {2, 1, 1, 1024, 1024},
+      {1, 3, 1, 4096, 4096}
+    };
+
+    for (int opt=0; opt<shardOptions.length; opt++) {
+      // first convert v2 produced by bioformats2raw to v3
+      Path v3Output = tmp.newFolder().toPath().resolve("v3-test");
+      Convert v3Converter = new Convert();
+      v3Converter.setInput(output.toString());
+      v3Converter.setOutput(v3Output.toString());
+
+      v3Converter.setSharding(shardOptions[opt]);
+      v3Converter.convertToV3();
+
+      // check list of codecs in the v3 arrays
+
+      Store store = new FilesystemStore(v3Output);
+      Array resolution = Array.open(store.resolve("0", "0"));
+
+      int[] shardSize = shardSizes[opt];
+      Assert.assertArrayEquals(resolution.metadata.chunkShape(), shardSize);
+
+      // now convert v3 back to v2
+      Path roundtripOutput = tmp.newFolder().toPath().resolve("v2-roundtrip-test");
+      Convert v2Converter = new Convert();
+      v2Converter.setInput(v3Output.toString());
+      v2Converter.setOutput(roundtripOutput.toString());
+      v2Converter.setWriteV2(true);
+      v2Converter.convertToV2();
+
+      Path originalOMEXML = output.resolve("OME").resolve("METADATA.ome.xml");
+      Path roundtripOMEXML = roundtripOutput.resolve("OME").resolve("METADATA.ome.xml");
+
+      // make sure the OME-XML is present and not changed
+      Assert.assertEquals(Files.readAllLines(originalOMEXML), Files.readAllLines(roundtripOMEXML));
+
+      // since the image is small, make sure all pixels are identical in both resolutions
+      for (int r=0; r<4; r++) {
+        ZarrArray original = ZarrGroup.open(output.resolve("0")).openArray(String.valueOf(r));
+        ZarrArray roundtrip = ZarrGroup.open(roundtripOutput.resolve("0")).openArray(String.valueOf(r));
+
+        compareZarrArrays(original, roundtrip);
+      }
+    }
+  }
+
   private void compareZarrArrays(ZarrArray original, ZarrArray roundtrip) throws Exception {
     Assert.assertArrayEquals(original.getShape(), roundtrip.getShape());
 
     int[] shape = original.getShape();
-    byte[] originalImage = new byte[shape[2] * shape[3] * shape[4]];
-    byte[] roundtripImage = new byte[shape[2] * shape[3] * shape[4]];
+    int arraySize = 1;
+    for (int s : shape) {
+      arraySize *= s;
+    }
+    byte[] originalImage = new byte[arraySize];
+    byte[] roundtripImage = new byte[arraySize];
     original.read(originalImage, shape);
     roundtrip.read(roundtripImage, shape);