wip: minor change to writeBlocks, implement readBlocks

* where readBlocks batches when possible
saalfeldlab · Jan 13, 2025 · 6fac5f2 · 6fac5f2
1 parent 52aaa2a
commit 6fac5f2
Show file tree

Hide file tree

Showing 3 changed files with 81 additions and 6 deletions.
diff --git a/src/main/java/org/janelia/saalfeldlab/n5/GsonKeyValueN5Reader.java b/src/main/java/org/janelia/saalfeldlab/n5/GsonKeyValueN5Reader.java
@@ -27,9 +27,13 @@
 
 import java.io.IOException;
 import java.io.UncheckedIOException;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
 
 import org.janelia.saalfeldlab.n5.N5Exception.N5IOException;
+import org.janelia.saalfeldlab.n5.shard.InMemoryShard;
 import org.janelia.saalfeldlab.n5.shard.Shard;
 import org.janelia.saalfeldlab.n5.shard.ShardParameters;
 import org.janelia.saalfeldlab.n5.shard.VirtualShard;
@@ -124,6 +128,49 @@ default DataBlock<?> readBlock(
 		}
 	}
 
+	@Override
+	default List<DataBlock<?>> readBlocks(
+			final String pathName,
+			final DatasetAttributes datasetAttributes,
+			final List<long[]> blockPositions) throws N5Exception {
+
+		// TODO which interface should have this implementation?
+		if (datasetAttributes instanceof ShardParameters) {
+
+			/* Group by shard index */
+			final HashMap<Integer, Shard<?>> shardBlockMap = new HashMap<>();
+			final HashMap<Integer, List<long[]>> shardPositionMap = new HashMap<>();
+			final ShardParameters shardAttributes = (ShardParameters)datasetAttributes;
+
+			for ( long[] blockPosition : blockPositions ) {
+				final long[] shardPosition = shardAttributes.getShardPositionForBlock(blockPosition);
+				final int shardHash = Arrays.hashCode(shardPosition);
+				if (!shardBlockMap.containsKey(shardHash)) {
+					final Shard<?> shard = getShard(pathName, (DatasetAttributes & ShardParameters)shardAttributes, shardPosition);
+					shardBlockMap.put(shardHash, shard);
+
+					final ArrayList<long[]> positionList = new ArrayList<>();
+					positionList.add(blockPosition);
+					shardPositionMap.put(shardHash, positionList);
+				}
+				else
+					shardPositionMap.get(shardBlockMap.get(shardHash)).add(blockPosition);
+			}
+
+			final ArrayList<DataBlock<?>> blocks = new ArrayList<>();
+			for (Shard<?> shard : shardBlockMap.values()) {
+				/* Add existing blocks before overwriting shard */
+				final int shardHash = Arrays.hashCode(shard.getGridPosition());
+				for( final long[] blkPosition : shardPositionMap.get(shardHash)) {
+					blocks.add(shard.getBlock(blkPosition));
+				}
+			}
+			return blocks;
+		} else
+			return GsonN5Reader.super.readBlocks(pathName, datasetAttributes, blockPositions);
+
+	}
+
 	@Override
 	default String[] list(final String pathName) throws N5Exception {
 

diff --git a/src/main/java/org/janelia/saalfeldlab/n5/GsonKeyValueN5Writer.java b/src/main/java/org/janelia/saalfeldlab/n5/GsonKeyValueN5Writer.java
@@ -216,8 +216,10 @@ default boolean removeAttributes(final String pathName, final List<String> attri
 		return removed;
 	}
 
-	@SuppressWarnings({ "rawtypes", "unchecked" })
-	@Override default <T> void writeBlocks(final String datasetPath, final DatasetAttributes datasetAttributes, final DataBlock<T>... dataBlocks) throws N5Exception {
+	@Override default <T> void writeBlocks(
+			final String datasetPath,
+			final DatasetAttributes datasetAttributes,
+			final DataBlock<T>... dataBlocks) throws N5Exception {
 
 		if (datasetAttributes instanceof ShardParameters) {
 			/* Group by shard index */
@@ -236,6 +238,7 @@ default boolean removeAttributes(final String pathName, final List<String> attri
 
 			for (InMemoryShard<T> shard : shardBlockMap.values()) {
 				/* Add existing blocks before overwriting shard */
+				@SuppressWarnings("unchecked")
 				final Shard<T> currentShard = (Shard<T>)getShard(datasetPath, (DatasetAttributes & ShardParameters)shardAttributes, shard.getGridPosition());
 				for (DataBlock<T> currentBlock : currentShard.getBlocks()) {
 					if (shard.getBlock(currentBlock.getGridPosition()) == null)
@@ -246,10 +249,7 @@ default boolean removeAttributes(final String pathName, final List<String> attri
 			}
 
 		} else {
-			/* Just write each block */
-			for (DataBlock<T> dataBlock : dataBlocks) {
-				writeBlock(datasetPath, datasetAttributes, dataBlock);
-			}
+			GsonN5Writer.super.writeBlocks(datasetPath, datasetAttributes, dataBlocks);
 		}
 	}
 

diff --git a/src/main/java/org/janelia/saalfeldlab/n5/N5Reader.java b/src/main/java/org/janelia/saalfeldlab/n5/N5Reader.java
@@ -293,6 +293,34 @@ DataBlock<?> readBlock(
 			final DatasetAttributes datasetAttributes,
 			final long... gridPosition) throws N5Exception;
 
+	/**
+	 * Reads multiple {@link DataBlock}s.
+	 * <p>
+	 * Implementations may optimize / batch read operations when possible, e.g.
+	 * in the case that the datasets are sharded.
+	 *
+	 * @param pathName
+	 *            dataset path
+	 * @param datasetAttributes
+	 *            the dataset attributes
+	 * @param gridPositions
+	 *            a list of grid positions
+	 * @return a list of data blocks
+	 * @throws N5Exception
+	 *             the exception
+	 */
+	default List<DataBlock<?>> readBlocks(
+			final String pathName,
+			final DatasetAttributes datasetAttributes,
+			final List<long[]> gridPositions) throws N5Exception {
+
+		final ArrayList<DataBlock<?>> blocks = new ArrayList<>();
+		for( final long[] p : gridPositions )
+			blocks.add(readBlock(pathName, datasetAttributes, p));
+
+		return blocks;
+	}
+
 	/**
 	 * Load a {@link DataBlock} as a {@link Serializable}. The offset is given
 	 * in