Skip to content

Commit

Permalink
test: ShardIndexTest
Browse files Browse the repository at this point in the history
  • Loading branch information
@bogovicj
bogovicj committed Jan 8, 2025
1 parent 1161859 commit 6e3cbe5
Show file tree
Hide file tree
Showing 2 changed files with 154 additions and 13 deletions.
74 changes: 61 additions & 13 deletions src/main/java/org/janelia/saalfeldlab/n5/shard/ShardIndex.java
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
package org.janelia.saalfeldlab.n5.shard;

import org.apache.commons.io.input.BoundedInputStream;
import org.janelia.saalfeldlab.n5.DataBlock;
import org.janelia.saalfeldlab.n5.DataType;
import org.janelia.saalfeldlab.n5.DatasetAttributes;
Expand All @@ -15,6 +16,7 @@
import org.janelia.saalfeldlab.n5.codec.DeterministicSizeCodec;
import org.janelia.saalfeldlab.n5.shard.ShardingCodec.IndexLocation;

import java.io.ByteArrayInputStream;
import java.io.DataInputStream;
import java.io.IOException;
import java.io.InputStream;
Expand Down Expand Up @@ -113,6 +115,28 @@ public long numBytes() {
return totalNumBytes;
}

public static ShardIndex read(byte[] data, final ShardIndex index) throws IOException {

final IndexByteBounds byteBounds = byteBounds(index, data.length);
final ByteArrayInputStream is = new ByteArrayInputStream(data);
is.skip(byteBounds.start);
BoundedInputStream bIs = BoundedInputStream.builder()
.setInputStream(is)
.setMaxCount(byteBounds.size).get();

return read(bIs, index);
}

public static ShardIndex read(InputStream in, final ShardIndex index) throws IOException {

@SuppressWarnings("unchecked")
final DataBlock<long[]> indexBlock = (DataBlock<long[]>) DefaultBlockReader.readBlock(in,
index.getIndexAttributes(), index.gridPosition);
final long[] indexData = indexBlock.getData();
System.arraycopy(indexData, 0, index.data, 0, index.data.length);
return index;
}

public static ShardIndex read(
final KeyValueAccess keyValueAccess,
final String key,
Expand All @@ -121,16 +145,9 @@ public static ShardIndex read(

final IndexByteBounds byteBounds = byteBounds(index, keyValueAccess.size(key));
try (final LockedChannel lockedChannel = keyValueAccess.lockForReading(key, byteBounds.start, byteBounds.end)) {
final long[] indexData;
try (final InputStream in = lockedChannel.newInputStream()) {
final DataBlock<long[]> indexBlock = (DataBlock<long[]>)DefaultBlockReader.readBlock(
in,
index.getIndexAttributes(),
index.gridPosition);
indexData = indexBlock.getData();
return read(in,index);
}
System.arraycopy(indexData, 0, index.data, 0, index.data.length);
return index;
} catch (final N5Exception.N5NoSuchKeyException e) {
return null;
} catch (final IOException | UncheckedIOException e) {
Expand All @@ -144,7 +161,7 @@ public static void write(
final String key
) throws IOException {

final long start = index.location == IndexLocation.START ? 0 : keyValueAccess.size(key);
final long start = index.location == IndexLocation.START ? 0 : sizeOrZero( keyValueAccess, key) ;
try (final LockedChannel lockedChannel = keyValueAccess.lockForWriting(key, start, index.numBytes())) {
try (final OutputStream os = lockedChannel.newOutputStream()) {
write(index, os);
Expand All @@ -154,6 +171,14 @@ public static void write(
}
}

private static long sizeOrZero(final KeyValueAccess keyValueAccess, final String key) {
try {
return keyValueAccess.size(key);
} catch (Exception e) {
return 0;
}
}

public static void write(final ShardIndex index, OutputStream out) throws IOException {

DefaultBlockWriter.writeBlock(out, index.getIndexAttributes(), index);
Expand Down Expand Up @@ -191,15 +216,17 @@ public static IndexByteBounds byteBounds(final long indexSize, final IndexLocati
}
}

private static class IndexByteBounds {
public static class IndexByteBounds {

private final long start;
private final long end;
public final long start;
public final long end;
public final long size;

private IndexByteBounds(long start, long end) {
public IndexByteBounds(long start, long end) {

this.start = start;
this.end = end;
this.size = end - start + 1;
}
}

Expand Down Expand Up @@ -241,4 +268,25 @@ private static int[] prepend(final int value, final int[] array) {
System.arraycopy(array, 0, indexBlockSize, 1, array.length);
return indexBlockSize;
}

@Override
public boolean equals(Object other) {

if (other instanceof ShardIndex) {

final ShardIndex index = (ShardIndex) other;
if (this.location != index.location)
return false;

if (!Arrays.equals(this.size, index.size))
return false;

if (!Arrays.equals(this.data, index.data))
return false;

}
return true;
}

}

93 changes: 93 additions & 0 deletions src/test/java/org/janelia/saalfeldlab/n5/shard/ShardIndexTest.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,93 @@
package org.janelia.saalfeldlab.n5.shard;

import static org.junit.Assert.assertEquals;

import java.io.IOException;
import java.nio.file.Paths;

import org.janelia.saalfeldlab.n5.DataType;
import org.janelia.saalfeldlab.n5.GzipCompression;
import org.janelia.saalfeldlab.n5.KeyValueAccess;
import org.janelia.saalfeldlab.n5.N5FSTest;
import org.janelia.saalfeldlab.n5.N5KeyValueWriter;
import org.janelia.saalfeldlab.n5.ShardedDatasetAttributes;
import org.janelia.saalfeldlab.n5.codec.BytesCodec;
import org.janelia.saalfeldlab.n5.codec.Codec;
import org.janelia.saalfeldlab.n5.codec.DeterministicSizeCodec;
import org.janelia.saalfeldlab.n5.codec.N5BlockCodec;
import org.janelia.saalfeldlab.n5.codec.checksum.Crc32cChecksumCodec;
import org.janelia.saalfeldlab.n5.shard.ShardingCodec.IndexLocation;
import org.junit.After;
import org.junit.Ignore;
import org.junit.Test;

public class ShardIndexTest {

private static final N5FSTest tempN5Factory = new N5FSTest();

@After
public void removeTempWriters() {
tempN5Factory.removeTempWriters();
}

@Test
public void testReadVirtual() throws IOException {

final N5KeyValueWriter writer = (N5KeyValueWriter) tempN5Factory.createTempN5Writer();
final KeyValueAccess kva = writer.getKeyValueAccess();

final int[] shardBlockGridSize = new int[] { 6, 5 };
final IndexLocation indexLocation = IndexLocation.END;
final DeterministicSizeCodec[] indexCodecs = new DeterministicSizeCodec[] { new BytesCodec(),
new Crc32cChecksumCodec() };

final String path = Paths.get(Paths.get(writer.getURI()).toAbsolutePath().toString(), "0").toString();

final ShardIndex index = new ShardIndex(shardBlockGridSize, indexLocation, indexCodecs);
index.set(0, 6, new int[] { 0, 0 });
index.set(19, 32, new int[] { 1, 0 });
index.set(93, 111, new int[] { 3, 0 });
index.set(143, 1, new int[] { 1, 2 });
ShardIndex.write(index, kva, path);

final ShardIndex other = new ShardIndex(shardBlockGridSize, indexLocation, indexCodecs);
ShardIndex.read(kva, path, other);

assertEquals(index, other);
}

@Test
@Ignore
public void testReadInMemory() throws IOException {

final N5KeyValueWriter writer = (N5KeyValueWriter) tempN5Factory.createTempN5Writer();
final KeyValueAccess kva = writer.getKeyValueAccess();

final int[] shardBlockGridSize = new int[] { 6, 5 };
final IndexLocation indexLocation = IndexLocation.END;
final DeterministicSizeCodec[] indexCodecs = new DeterministicSizeCodec[] { new BytesCodec(),
new Crc32cChecksumCodec() };
final String path = Paths.get(Paths.get(writer.getURI()).toAbsolutePath().toString(), "0").toString();

final ShardIndex index = new ShardIndex(shardBlockGridSize, indexLocation, indexCodecs);
index.set(0, 6, new int[] { 0, 0 });
index.set(19, 32, new int[] { 1, 0 });
index.set(93, 111, new int[] { 3, 0 });
index.set(143, 1, new int[] { 1, 2 });
ShardIndex.write(index, kva, path);

ShardedDatasetAttributes attrs = new ShardedDatasetAttributes(
new long[]{6,5},
shardBlockGridSize,
new int[]{1,1},
DataType.UINT8,
new Codec[]{new N5BlockCodec(), new GzipCompression(4)},
new DeterministicSizeCodec[]{new BytesCodec(), new Crc32cChecksumCodec()},
indexLocation
);

final InMemoryShard shard = InMemoryShard.readShard(kva, path, new long[] {0,0}, attrs);

assertEquals(index, shard.index);
}
}

0 comments on commit 6e3cbe5

Please sign in to comment.