diff --git a/src/main/java/org/janelia/saalfeldlab/n5/shard/InMemoryShard.java b/src/main/java/org/janelia/saalfeldlab/n5/shard/InMemoryShard.java index c527df26..70340b8e 100644 --- a/src/main/java/org/janelia/saalfeldlab/n5/shard/InMemoryShard.java +++ b/src/main/java/org/janelia/saalfeldlab/n5/shard/InMemoryShard.java @@ -1,20 +1,27 @@ package org.janelia.saalfeldlab.n5.shard; +import java.io.ByteArrayInputStream; import java.io.IOException; +import java.io.InputStream; import java.io.OutputStream; import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; import java.util.List; +import org.apache.commons.io.input.BoundedInputStream; import org.apache.commons.io.output.ByteArrayOutputStream; import org.apache.commons.io.output.CountingOutputStream; import org.apache.commons.io.output.ProxyOutputStream; import org.checkerframework.checker.units.qual.A; import org.janelia.saalfeldlab.n5.DataBlock; import org.janelia.saalfeldlab.n5.DatasetAttributes; +import org.janelia.saalfeldlab.n5.DefaultBlockReader; import org.janelia.saalfeldlab.n5.DefaultBlockWriter; +import org.janelia.saalfeldlab.n5.KeyValueAccess; +import org.janelia.saalfeldlab.n5.LockedChannel; import org.janelia.saalfeldlab.n5.shard.ShardingCodec.IndexLocation; +import org.janelia.saalfeldlab.n5.util.GridIterator; public class InMemoryShard extends AbstractShard { @@ -27,7 +34,6 @@ public class InMemoryShard extends AbstractShard { * Use morton- or c-ording instead of writing blocks out in the order they're added? * (later) */ - public InMemoryShard(final A datasetAttributes, final long[] shardPosition) { this( datasetAttributes, shardPosition, null); @@ -99,6 +105,61 @@ public void write(final OutputStream out) throws IOException { writeShardStart(out, this); } + public static InMemoryShard readShard( + final KeyValueAccess kva, final String key, final long[] gridPosition, final A attributes) + throws IOException { + + try (final LockedChannel lockedChannel = kva.lockForReading(key)) { + try (final InputStream is = lockedChannel.newInputStream()) { + return readShard(is, gridPosition, attributes); + } + } + } + + public static InMemoryShard readShard( + final InputStream inputStream, final long[] gridPosition, final A attributes) throws IOException { + + try (ByteArrayOutputStream result = new ByteArrayOutputStream()) { + byte[] buffer = new byte[1024]; + for (int length; (length = inputStream.read(buffer)) != -1;) { + result.write(buffer, 0, length); + } + return readShard(result.toByteArray(), gridPosition, attributes); + } + } + + public static InMemoryShard readShard(final byte[] data, + long[] shardPosition, final A attributes) throws IOException { + + final ShardIndex index = attributes.createIndex(); + ShardIndex.read(data, index); + + final InMemoryShard shard = new InMemoryShard(attributes, shardPosition, index); + final GridIterator it = new GridIterator(attributes.getBlocksPerShard()); + while (it.hasNext()) { + + final long[] p = it.next(); + final int[] pInt = GridIterator.long2int(p); + + if (index.exists(pInt)) { + + final ByteArrayInputStream is = new ByteArrayInputStream(data); + is.skip(index.getOffset(pInt)); + BoundedInputStream bIs = BoundedInputStream.builder().setInputStream(is) + .setMaxCount(index.getNumBytes(pInt)).get(); + + final long[] blockGridPosition = attributes.getBlockPositionFromShardPosition(shardPosition, p); + @SuppressWarnings("unchecked") + final DataBlock blk = (DataBlock) DefaultBlockReader.readBlock(bIs, attributes, + blockGridPosition); + shard.addBlock(blk); + bIs.close(); + } + } + + return shard; + } + public static void writeShard(final OutputStream out, final Shard shard) throws IOException { fromShard(shard).write(out); diff --git a/src/main/java/org/janelia/saalfeldlab/n5/shard/VirtualShard.java b/src/main/java/org/janelia/saalfeldlab/n5/shard/VirtualShard.java index 3ebe47d3..3a36ac14 100644 --- a/src/main/java/org/janelia/saalfeldlab/n5/shard/VirtualShard.java +++ b/src/main/java/org/janelia/saalfeldlab/n5/shard/VirtualShard.java @@ -4,6 +4,7 @@ import java.io.InputStream; import java.io.OutputStream; import java.io.UncheckedIOException; +import java.util.function.Supplier; import org.checkerframework.checker.units.qual.A; import org.janelia.saalfeldlab.n5.DataBlock; @@ -28,6 +29,20 @@ public VirtualShard(final A data this.path = path; } + public VirtualShard(final A datasetAttributes, long[] gridPosition) { + + this(datasetAttributes, gridPosition, null, null); + } + + @SuppressWarnings("unchecked") + public DataBlock getBlock(Supplier inputSupplier, long... blockGridPosition) throws IOException { + + // TODO this method is just a wrapper around readBlock and probably not worth keeping + try (InputStream is = inputSupplier.get()) { + return (DataBlock) DefaultBlockReader.readBlock(is, datasetAttributes, blockGridPosition); + } + } + @SuppressWarnings("unchecked") @Override public DataBlock getBlock(long... blockGridPosition) { diff --git a/src/main/java/org/janelia/saalfeldlab/n5/util/GridIterator.java b/src/main/java/org/janelia/saalfeldlab/n5/util/GridIterator.java index 8b9ab8fc..1ea16efb 100644 --- a/src/main/java/org/janelia/saalfeldlab/n5/util/GridIterator.java +++ b/src/main/java/org/janelia/saalfeldlab/n5/util/GridIterator.java @@ -63,6 +63,10 @@ public long[] next() { return position; } + public int[] nextAsInt() { + return long2int(next()); + } + public int getIndex() { return index; } @@ -93,11 +97,4 @@ final static public long[] int2long(final int[] i) { return l; } - public static void main(String[] args) { - - final GridIterator it = new GridIterator(new int[]{2, 2, 2}); - while (it.hasNext()) { - System.out.println(Arrays.toString(it.next())); - } - } }