feat: writeShardEndStream

saalfeldlab · Jan 2, 2025 · 1ced570 · 1ced570
1 parent b62dc00
commit 1ced570
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 11 deletions.
diff --git a/src/main/java/org/janelia/saalfeldlab/n5/GsonKeyValueN5Writer.java b/src/main/java/org/janelia/saalfeldlab/n5/GsonKeyValueN5Writer.java
@@ -255,7 +255,6 @@ default <T> void writeShard(
 		try (final LockedChannel lock = getKeyValueAccess().lockForWriting(shardPath)) {
 			try (final OutputStream out = lock.newOutputStream()) {
 				InMemoryShard.fromShard(shard).write(out);
-				out.close();
 			}
 		} catch (final IOException | UncheckedIOException e) {
 			throw new N5IOException(

diff --git a/src/main/java/org/janelia/saalfeldlab/n5/shard/InMemoryShard.java b/src/main/java/org/janelia/saalfeldlab/n5/shard/InMemoryShard.java
@@ -8,6 +8,7 @@
 
 import org.apache.commons.io.output.ByteArrayOutputStream;
 import org.apache.commons.io.output.CountingOutputStream;
+import org.apache.commons.io.output.ProxyOutputStream;
 import org.janelia.saalfeldlab.n5.DataBlock;
 import org.janelia.saalfeldlab.n5.DefaultBlockWriter;
 import org.janelia.saalfeldlab.n5.ShardedDatasetAttributes;
@@ -80,7 +81,7 @@ public ShardIndex getIndex() {
 	public void write(final OutputStream out) throws IOException {
 
 		if (indexLocation() == IndexLocation.END)
-			writeShardEnd(out, this);
+			writeShardEndStream(out, this);
 		else
 			writeShardStart(out, this);
 	}
@@ -110,21 +111,31 @@ protected static <T> void writeShardEndStream(
 
 		final ShardIndexBuilder indexBuilder = new ShardIndexBuilder(shard);
 		indexBuilder.indexLocation(IndexLocation.END);
+		indexBuilder.setCodecs(datasetAttributes.getShardingCodec().getIndexCodecs());
 
-		final CountingOutputStream cout = new CountingOutputStream(out);
-
-		long offset = 0;
+		final ProxyOutputStream nop = new ProxyOutputStream(out) {
+
+			@Override public void close() {
+				//nop
+			}
+		};
+
+		final CountingOutputStream cout = new CountingOutputStream(nop);
+
+		long bytesWritten = 0;
 		for (int i = 0; i < shard.numBlocks(); i++) {
 
 			final DataBlock<T> block = shard.getBlock(i);
 			DefaultBlockWriter.writeBlock(cout, datasetAttributes, block);
-
-			indexBuilder.addBlock( block.getGridPosition(), offset);
-			offset = cout.getByteCount();
+
+
+			final long size = cout.getByteCount() - bytesWritten;
+			bytesWritten = cout.getByteCount();
+
+			indexBuilder.addBlock( block.getGridPosition(), size);
 		}
 
-		final ShardIndex index = indexBuilder.build();
-		DefaultBlockWriter.writeBlock(out, datasetAttributes, index);
+		ShardIndex.write(indexBuilder.build(), out);
 	}
 
 	protected static <T> void writeShardEnd(

diff --git a/src/main/java/org/janelia/saalfeldlab/n5/shard/ShardIndexBuilder.java b/src/main/java/org/janelia/saalfeldlab/n5/shard/ShardIndexBuilder.java
@@ -55,7 +55,7 @@ public ShardIndexBuilder setCodecs(DeterministicSizeCodec... codecs) {
 	}
 
 	public ShardIndexBuilder addBlock(long[] blockPosition, long numBytes) {
-
+		//TODO Caleb: Maybe move to ShardIndex?
 		final long[] blockPositionInShard = shard.getDatasetAttributes().getBlockPositionInShard(
 				shard.getGridPosition(),
 				blockPosition);

diff --git a/src/test/java/org/janelia/saalfeldlab/n5/shard/ShardDemos.java b/src/test/java/org/janelia/saalfeldlab/n5/shard/ShardDemos.java
@@ -71,6 +71,7 @@ public static void main(String[] args) throws MalformedURLException {
 
 	@Parameterized.Parameters(name = "IndexLocation({0}), Block ByteOrder({1}), Index ByteOrder({2})")
 	public static Collection<Object[]> data() {
+
 		final ArrayList<Object[]> params = new ArrayList<>();
 		for (IndexLocation indexLoc : IndexLocation.values()) {
 			for (ByteOrder blockByteOrder : new ByteOrder[]{ByteOrder.BIG_ENDIAN, ByteOrder.LITTLE_ENDIAN}) {