-
Notifications
You must be signed in to change notification settings - Fork 242
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add support for reading and writing SBI files for BAMs.
- Loading branch information
Showing
7 changed files
with
589 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,65 @@ | ||
package htsjdk.samtools; | ||
|
||
import htsjdk.samtools.cram.io.InputStreamUtils; | ||
import htsjdk.samtools.seekablestream.SeekablePathStream; | ||
import htsjdk.samtools.seekablestream.SeekableStream; | ||
import htsjdk.samtools.util.BlockCompressedInputStream; | ||
import htsjdk.samtools.util.IOUtil; | ||
|
||
import java.io.EOFException; | ||
import java.io.IOException; | ||
import java.io.OutputStream; | ||
import java.nio.ByteBuffer; | ||
import java.nio.ByteOrder; | ||
import java.nio.file.Files; | ||
import java.nio.file.Path; | ||
|
||
/** | ||
* Writes SBI files for BAM files, as understood by {@link SBIIndex}. | ||
*/ | ||
public final class BAMSBIIndexer { | ||
|
||
/** | ||
* Perform indexing on the given BAM file, at the granularity level specified. | ||
* | ||
* @param bamFile the path to the BAM file | ||
* @param granularity write the offset of every n-th alignment to the index | ||
* @throws IOException as per java IO contract | ||
*/ | ||
public static void createIndex(final Path bamFile, final long granularity) throws IOException { | ||
Path splittingBaiFile = IOUtil.addExtension(bamFile, SBIIndex.FILE_EXTENSION); | ||
try (SeekableStream in = new SeekablePathStream(bamFile); OutputStream out = Files.newOutputStream(splittingBaiFile)) { | ||
createIndex(in, out, granularity); | ||
} | ||
} | ||
|
||
/** | ||
* Perform indexing on the given BAM file, at the granularity level specified. | ||
* | ||
* @param in a seekable stream for reading the BAM file from | ||
* @param out the stream to write the index to | ||
* @param granularity write the offset of every n-th alignment to the index | ||
* @throws IOException as per java IO contract | ||
*/ | ||
public static void createIndex(final SeekableStream in, final OutputStream out, final long granularity) throws IOException { | ||
long recordStart = SAMUtils.findVirtualOffsetOfFirstRecordInBam(in); | ||
try (BlockCompressedInputStream blockIn = new BlockCompressedInputStream(in)) { | ||
blockIn.seek(recordStart); | ||
final ByteBuffer byteBuffer = ByteBuffer.allocate(4).order(ByteOrder.LITTLE_ENDIAN); // BAM is little-endian | ||
SBIIndexWriter indexWriter = new SBIIndexWriter(out, granularity); | ||
while (true) { | ||
try { | ||
recordStart = blockIn.getFilePointer(); | ||
InputStreamUtils.readFully(blockIn, byteBuffer.array(), 0, 4); | ||
final int blockSize = byteBuffer.getInt(0); // length of remainder of alignment record | ||
indexWriter.processRecord(recordStart); | ||
InputStreamUtils.skipFully(blockIn, blockSize); | ||
} catch (EOFException e) { | ||
break; | ||
} | ||
} | ||
indexWriter.writeVirtualOffset(recordStart); | ||
indexWriter.finish(in.length()); | ||
} | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,251 @@ | ||
package htsjdk.samtools; | ||
|
||
import htsjdk.samtools.util.BinaryCodec; | ||
import htsjdk.samtools.util.BlockCompressedFilePointerUtil; | ||
|
||
import java.io.BufferedInputStream; | ||
import java.io.IOException; | ||
import java.io.InputStream; | ||
import java.nio.file.Files; | ||
import java.nio.file.Path; | ||
import java.util.ArrayList; | ||
import java.util.Arrays; | ||
import java.util.List; | ||
import java.util.NavigableSet; | ||
import java.util.TreeSet; | ||
|
||
/** | ||
* SBI is an index into BGZF-compressed data files, which has an entry for the file position of the start of every | ||
* <i>n</i>th record. Reads files that were created by {@link BAMSBIIndexer}. | ||
*/ | ||
public final class SBIIndex { | ||
|
||
public static class Header { | ||
private final long fileLength; | ||
private final byte[] md5; | ||
private final byte[] uuid; | ||
private final long totalNumberOfRecords; | ||
private final long granularity; | ||
|
||
public Header(long fileLength, byte[] md5, byte[] uuid, long totalNumberOfRecords, long granularity) { | ||
this.fileLength = fileLength; | ||
this.md5 = md5; | ||
this.uuid = uuid; | ||
this.totalNumberOfRecords = totalNumberOfRecords; | ||
this.granularity = granularity; | ||
} | ||
|
||
public long getFileLength() { | ||
return fileLength; | ||
} | ||
|
||
public byte[] getMd5() { | ||
return md5; | ||
} | ||
|
||
public byte[] getUuid() { | ||
return uuid; | ||
} | ||
|
||
public long getTotalNumberOfRecords() { | ||
return totalNumberOfRecords; | ||
} | ||
|
||
public long getGranularity() { | ||
return granularity; | ||
} | ||
} | ||
|
||
public static final String FILE_EXTENSION = ".sbi"; | ||
|
||
/** | ||
* SBI magic number. | ||
*/ | ||
static final byte[] SBI_MAGIC = "SBI\1".getBytes(); | ||
|
||
private final Header header; | ||
private final NavigableSet<Long> virtualOffsets; | ||
|
||
/** | ||
* Create an in-memory SBI with the given virtual offsets. | ||
* @param virtualOffsets the offsets in the index | ||
*/ | ||
public SBIIndex(final Header header, final NavigableSet<Long> virtualOffsets) { | ||
this.header = header; | ||
this.virtualOffsets = new TreeSet<>(virtualOffsets); | ||
if (this.virtualOffsets.isEmpty()) { | ||
throw new RuntimeException("Invalid SBI format: should contain at least one offset"); | ||
} | ||
} | ||
|
||
/** | ||
* Load an SBI into memory from a path. | ||
* @param path the path to the SBI file | ||
* @throws IOException as per java IO contract | ||
*/ | ||
public static SBIIndex load(final Path path) throws IOException { | ||
try (InputStream in = new BufferedInputStream(Files.newInputStream(path))) { | ||
return readIndex(in); | ||
} | ||
} | ||
|
||
/** | ||
* Load an SBI into memory from a stream. | ||
* @param in the stream to read the SBI from | ||
*/ | ||
public static SBIIndex load(final InputStream in) { | ||
return readIndex(in); | ||
} | ||
|
||
private static SBIIndex readIndex(final InputStream in) { | ||
BinaryCodec binaryCodec = new BinaryCodec(in); | ||
Header header = readHeader(binaryCodec); | ||
long numOffsets = binaryCodec.readLong(); | ||
NavigableSet<Long> virtualOffsets = new TreeSet<>(); | ||
long prev = -1; | ||
for (long i = 0; i < numOffsets; i++) { | ||
long cur = binaryCodec.readLong(); | ||
if (prev > cur) { | ||
throw new RuntimeException(String.format( | ||
"Invalid SBI; offsets not in order: %#x > %#x", | ||
prev, cur)); | ||
} | ||
virtualOffsets.add(prev = cur); | ||
} | ||
return new SBIIndex(header, virtualOffsets); | ||
} | ||
|
||
private static Header readHeader(BinaryCodec binaryCodec) { | ||
final byte[] buffer = new byte[SBI_MAGIC.length]; | ||
binaryCodec.readBytes(buffer); | ||
if (!Arrays.equals(buffer, SBI_MAGIC)) { | ||
throw new RuntimeException("Invalid file header in SBI: " + new String(buffer) + " (" + Arrays.toString(buffer) + ")"); | ||
} | ||
long fileLength = binaryCodec.readLong(); | ||
byte[] md5 = new byte[16]; | ||
binaryCodec.readBytes(md5); | ||
byte[] uuid = new byte[16]; | ||
binaryCodec.readBytes(uuid); | ||
long totalNumberOfRecords = binaryCodec.readLong(); | ||
long granularity = binaryCodec.readLong(); | ||
return new Header(fileLength, md5, uuid, totalNumberOfRecords, granularity); | ||
} | ||
|
||
/** | ||
* Returns the granularity of the index, that is the number of alignments between subsequent entries in the index, | ||
* or zero if not specified. | ||
* @return the granularity of the index | ||
*/ | ||
public long getGranularity() { | ||
return header.getGranularity(); | ||
} | ||
|
||
/** | ||
* Returns the entries in the index. | ||
* | ||
* @return a set of file pointers for all the alignment offsets in the index, in ascending order. The last | ||
* virtual file pointer is the position at which the next record would start if it were added to the file. | ||
*/ | ||
public NavigableSet<Long> getVirtualOffsets() { | ||
return new TreeSet<>(virtualOffsets); | ||
} | ||
|
||
/** | ||
* Returns number of entries in the index. | ||
* | ||
* @return the number of virtual offsets in the index | ||
*/ | ||
public long size() { | ||
return virtualOffsets.size(); | ||
} | ||
|
||
/** | ||
* Returns the length of the data file in bytes. | ||
* | ||
* @return the length of the data file in bytes | ||
*/ | ||
public long dataFileLength() { | ||
return header.getFileLength(); | ||
} | ||
|
||
/** | ||
* Split the data file for this index into non-overlapping chunks of roughly the given size that cover the whole | ||
* file and that can be read independently of one another. | ||
* | ||
* @param splitSize the rough size of each split in bytes | ||
* @return a list of contiguous, non-overlapping, sorted chunks that cover the whole data file | ||
* @see #getChunk(long, long) | ||
*/ | ||
public List<Chunk> split(long splitSize) { | ||
if (splitSize <= 0) { | ||
throw new IllegalArgumentException(String.format("Split size must be positive: %s", splitSize)); | ||
} | ||
long fileSize = dataFileLength(); | ||
List<Chunk> chunks = new ArrayList<>(); | ||
for (long splitStart = 0; splitStart < fileSize; splitStart += splitSize) { | ||
Chunk chunk = getChunk(splitStart, splitStart + splitSize); | ||
if (chunk != null) { | ||
chunks.add(chunk); | ||
} | ||
} | ||
return chunks; | ||
} | ||
|
||
/** | ||
* Return a chunk that corresponds to the given range in the data file. Note that the chunk does not necessarily | ||
* completely cover the given range, however this method will map a set of contiguous, non-overlapping file ranges | ||
* that cover the whole data file to a set of contiguous, non-overlapping chunks that cover the whole data file. | ||
* | ||
* @param splitStart the start of the file range (inclusive) | ||
* @param splitEnd the start of the file range (exclusive) | ||
* @return a chunk whose virtual start is at the first alignment start position that is greater than or equal to the | ||
* given split start position, and whose virtual end is at the first alignment start position that is greater than | ||
* or equal to the given split end position, or null if the chunk would be empty. | ||
* @see #split(long) | ||
*/ | ||
public Chunk getChunk(long splitStart, long splitEnd) { | ||
if (splitStart >= splitEnd) { | ||
throw new IllegalArgumentException(String.format("Split start (%s) must be less than end (%s)", splitStart, splitEnd)); | ||
} | ||
long maxEnd = BlockCompressedFilePointerUtil.getBlockAddress(virtualOffsets.last()); | ||
splitStart = Math.min(splitStart, maxEnd); | ||
splitEnd = Math.min(splitEnd, maxEnd); | ||
long virtualSplitStart = BlockCompressedFilePointerUtil.makeFilePointer(splitStart); | ||
long virtualSplitEnd = BlockCompressedFilePointerUtil.makeFilePointer(splitEnd); | ||
Long virtualSplitStartAlignment = virtualOffsets.ceiling(virtualSplitStart); | ||
Long virtualSplitEndAlignment = virtualOffsets.ceiling(virtualSplitEnd); | ||
// neither virtualSplitStartAlignment nor virtualSplitEndAlignment should ever be null, but check anyway | ||
if (virtualSplitStartAlignment == null) { | ||
throw new IllegalArgumentException(String.format("No virtual offset found for virtual file pointer %s, last virtual offset %s", | ||
BlockCompressedFilePointerUtil.asString(virtualSplitStart), BlockCompressedFilePointerUtil.asString(virtualOffsets.last()))); | ||
} | ||
if (virtualSplitEndAlignment == null) { | ||
throw new IllegalArgumentException(String.format("No virtual offset found for virtual file pointer %s, last virtual offset %s", | ||
BlockCompressedFilePointerUtil.asString(virtualSplitEnd), BlockCompressedFilePointerUtil.asString(virtualOffsets.last()))); | ||
} | ||
if (virtualSplitStartAlignment.longValue() == virtualSplitEndAlignment.longValue()) { | ||
return null; | ||
} | ||
return new Chunk(virtualSplitStartAlignment, virtualSplitEndAlignment); | ||
} | ||
|
||
@Override | ||
public boolean equals(Object o) { | ||
if (this == o) return true; | ||
if (o == null || getClass() != o.getClass()) return false; | ||
|
||
SBIIndex that = (SBIIndex) o; | ||
|
||
return virtualOffsets.equals(that.virtualOffsets); | ||
} | ||
|
||
@Override | ||
public int hashCode() { | ||
return virtualOffsets.hashCode(); | ||
} | ||
|
||
@Override | ||
public String toString() { | ||
return virtualOffsets.toString(); | ||
} | ||
} |
Oops, something went wrong.