feat: writeBlocks respects existing blocks in a given shard if not ov…

…erwriting those blocks explicitly
saalfeldlab · Jan 8, 2025 · 52f762e · 52f762e
1 parent 5badbb7
commit 52f762e
Show file tree

Hide file tree

Showing 5 changed files with 109 additions and 22 deletions.
diff --git a/src/main/java/org/janelia/saalfeldlab/n5/GsonKeyValueN5Writer.java b/src/main/java/org/janelia/saalfeldlab/n5/GsonKeyValueN5Writer.java
@@ -232,6 +232,14 @@ default boolean removeAttributes(final String pathName, final List<String> attri
 			}
 
 			for (InMemoryShard<T> shard : shardBlockMap.values()) {
+
+				/* Add existing blocks before overwriting shard */
+				final Shard<T> currentShard = (Shard<T>)getShard(datasetPath, shardAttributes, shard.getGridPosition());
+				for (DataBlock<T> currentBlock : currentShard.getBlocks()) {
+					if (shard.getBlock(currentBlock.getGridPosition()) == null)
+						shard.addBlock(currentBlock);
+				}
+
 				writeShard(datasetPath, shardAttributes, shard);
 			}
 		} else {
@@ -240,7 +248,6 @@ default boolean removeAttributes(final String pathName, final List<String> attri
 				writeBlock(datasetPath, datasetAttributes, dataBlock);
 			}
 		}
-
 	}
 
 	@Override

diff --git a/src/main/java/org/janelia/saalfeldlab/n5/shard/InMemoryShard.java b/src/main/java/org/janelia/saalfeldlab/n5/shard/InMemoryShard.java
@@ -67,6 +67,7 @@ public int numBlocks() {
 		return blocks.size();
 	}
 
+	@Override
 	public List<DataBlock<T>> getBlocks() {
 
 		return new ArrayList<>(blocks.values());

diff --git a/src/main/java/org/janelia/saalfeldlab/n5/shard/Shard.java b/src/main/java/org/janelia/saalfeldlab/n5/shard/Shard.java
@@ -1,7 +1,9 @@
 package org.janelia.saalfeldlab.n5.shard;
 
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Iterator;
+import java.util.List;
 
 import org.janelia.saalfeldlab.n5.DataBlock;
 import org.janelia.saalfeldlab.n5.ShardedDatasetAttributes;
@@ -106,9 +108,23 @@ default Iterator<DataBlock<T>> iterator() {
 		return new DataBlockIterator<T>(this);
 	}
 
-	default DataBlock<T>[] getAllBlocks(long... position) {
-		//TODO Caleb: Do we want this?
-		return null;
+	default List<DataBlock<T>> getBlocks() {
+
+		final ShardIndex shardIndex = getIndex();
+		final ShardedDatasetAttributes attrs = getDatasetAttributes();
+		final List<DataBlock<T>> blocks = new ArrayList<>();
+		for (long blockIdx = 0; blockIdx < attrs.getNumBlocks(); blockIdx++) {
+			int shardOffset = (int)blockIdx * 2;
+			final long[] index = shardIndex.getData();
+			if (index[shardOffset] == Shard.EMPTY_INDEX_NBYTES || index[shardOffset+1] == EMPTY_INDEX_NBYTES)
+				continue;
+
+			final long[] blockPosInShard = ShardIndex.shardPositionFromIndexOffset(shardOffset, attrs.getBlocksPerShard());
+			final long[] blockPosInImg = attrs.getBlockPositionFromShardPosition(getGridPosition(), blockPosInShard);
+			blocks.add(getBlock(blockPosInImg));
+		}
+
+		return blocks;
 	}
 
 	public ShardIndex getIndex();

diff --git a/src/main/java/org/janelia/saalfeldlab/n5/shard/ShardIndex.java b/src/main/java/org/janelia/saalfeldlab/n5/shard/ShardIndex.java
@@ -236,18 +236,31 @@ private static int[] prepend(final int value, final int[] array) {
 		return indexBlockSize;
 	}
 
-	public static void main(String[] args) {
-
-		final ShardIndex ib = new ShardIndex(new int[]{2, 2});
+	/**
+	 * Calculate the block position in the shard grid for a given index offset.
+	 *
+	 * @param offset          the offset into the index
+	 * @param blocksPerShard the dimensions of the shard in blocks
+	 * @return the relative position in the shard grid
+	 */
+	public static long[] shardPositionFromIndexOffset(int offset, int[] blocksPerShard) {
+
+		int maxOffset = 1;
+		for (int i = 0; i < blocksPerShard.length; i++) {
+			maxOffset *= blocksPerShard[i];
+		}
+		if (offset >= maxOffset*2) {
+			throw new IllegalArgumentException("Shard Index Offset " + offset + " is out of bounds for shard dimensions " + Arrays.toString(blocksPerShard));
+		}
 
-		ib.set(8, 9, new long[]{1, 1});
+		final long[] position = new long[blocksPerShard.length];
+		int remainder = offset / 2;
 
-		// System.out.println(ib.getIndex(0, 0));
-		// System.out.println(ib.getIndex(1, 0));
-		// System.out.println(ib.getIndex(0, 1));
-		// System.out.println(ib.getIndex(1, 1));
+		for (int dim = blocksPerShard.length - 1; dim >= 0; dim--) { // Iterate backwards
+			position[dim] = remainder % blocksPerShard[dim]; // Calculate position for this dimension
+			remainder /= blocksPerShard[dim]; // Update the remainder
+		}
 
-		System.out.println("done");
+		return position;
 	}
-
 }
diff --git a/src/test/java/org/janelia/saalfeldlab/n5/shard/ShardTest.java b/src/test/java/org/janelia/saalfeldlab/n5/shard/ShardTest.java
@@ -69,7 +69,7 @@ private ShardedDatasetAttributes getTestAttributes(long[] dimensions, int[] shar
 				shardSize,
 				blockSize,
 				DataType.UINT8,
-				new Codec[]{new N5BlockCodec(dataByteOrder) , new GzipCompression(4)},
+				new Codec[]{new N5BlockCodec(dataByteOrder)}, // , new GzipCompression(4)},
 				new DeterministicSizeCodec[]{new BytesCodec(indexByteOrder), new Crc32cChecksumCodec()},
 				indexLocation
 		);
@@ -99,6 +99,7 @@ public void writeReadBlocksTest() {
 			data[i] = (byte)((100) + (10) + i);
 		}
 
+
 		writer.writeBlocks(
 				"shard",
 				datasetAttributes,
@@ -117,15 +118,64 @@ public void writeReadBlocksTest() {
 		);
 
 		final KeyValueAccess kva = ((N5KeyValueWriter)writer).getKeyValueAccess();
-		String p = writer.getURI().getPath();
-		final String shard00 = kva.compose(writer.getURI(), "shard", "0", "0");
-		kva.exists(shard00);
 
-		final String shard10 = kva.compose(writer.getURI(), "shard", "0", "0");
-		kva.exists(shard10);
+		final String[][] keys = new String[][]{
+				{"shard", "0", "0"},
+				{"shard", "1", "0"},
+				{"shard", "2", "2"}
+		};
+		for (String[] key : keys) {
+			final String shard = kva.compose(writer.getURI(), key);
+			Assert.assertTrue("Shard at" + Arrays.toString(key) + "Does not exist", kva.exists(shard));
+		}
+
+		final long[][] blockIndices = new long[][]{ {0,0}, {0,1}, {1,0}, {1,1}, {4,0}, {5,0}, {11,11}};
+		for (long[] blockIndex : blockIndices) {
+			final DataBlock<?> block = writer.readBlock("shard", datasetAttributes, blockIndex);
+			Assert.assertArrayEquals("Read from shard doesn't match", data, (byte[])block.getData());
+		}
+
+		final byte[] data2 = new byte[numElements];
+		for (int i = 0; i < data2.length; i++) {
+			data2[i] = (byte)(10 + i);
+		}
+		writer.writeBlocks(
+				"shard",
+				datasetAttributes,
+				/* shard (0, 0) */
+				new ByteArrayDataBlock(blockSize, new long[]{0,0}, data2),
+				new ByteArrayDataBlock(blockSize, new long[]{1,1}, data2),
 
-		final String shard33 = kva.compose(writer.getURI(), "shard", "0", "0");
-		kva.exists(shard33);
+				/* shard (0, 1) */
+				new ByteArrayDataBlock(blockSize, new long[]{0,4}, data2),
+				new ByteArrayDataBlock(blockSize, new long[]{0,5}, data2),
+
+				/* shard (2, 2) */
+				new ByteArrayDataBlock(blockSize, new long[]{10,10}, data2)
+		);
+
+		final String[][] keys2 = new String[][]{
+				{"shard", "0", "0"},
+				{"shard", "1", "0"},
+				{"shard", "0", "1"},
+				{"shard", "2", "2"}
+		};
+		for (String[] key : keys2) {
+			final String shard = kva.compose(writer.getURI(), key);
+			Assert.assertTrue("Shard at" + Arrays.toString(key) + "Does not exist", kva.exists(shard));
+		}
+
+		final long[][] oldBlockIndices = new long[][]{{0,1}, {1,0}, {4,0}, {5,0}, {11,11}};
+		for (long[] blockIndex : oldBlockIndices) {
+			final DataBlock<?> block = writer.readBlock("shard", datasetAttributes, blockIndex);
+			Assert.assertArrayEquals("Read from shard doesn't match", data, (byte[])block.getData());
+		}
+
+		final long[][] newBlockIndices = new long[][]{{0,0}, {1,1}, {0,4}, {0,5}, {10,10}};
+		for (long[] blockIndex : newBlockIndices) {
+			final DataBlock<?> block = writer.readBlock("shard", datasetAttributes, blockIndex);
+			Assert.assertArrayEquals("Read from shard doesn't match", data2, (byte[])block.getData());
+		}
 	}
 
 	@Test