vcf/reader: Add query iterator

zaeleus · Apr 28, 2021 · 302033d · brainstorm · May 11, 2021 · zaeleus
1 parent 3ca4b73
commit 302033d
Show file tree

Hide file tree

Showing 3 changed files with 227 additions and 1 deletion.
diff --git a/noodles-vcf/Cargo.toml b/noodles-vcf/Cargo.toml
@@ -9,4 +9,6 @@ edition = "2018"
 indexmap = "1.4.0"
 nom = "6.1.2"
 noodles-bgzf = { path = "../noodles-bgzf" }
+noodles-core = { path = "../noodles-core" }
+noodles-tabix = { path = "../noodles-tabix" }
 percent-encoding = "2.1.0"
diff --git a/noodles-vcf/src/reader.rs b/noodles-vcf/src/reader.rs
@@ -1,12 +1,15 @@
 //! VCF reader and iterators.
 
+mod query;
 mod records;
 
-pub use self::records::Records;
+pub use self::{query::Query, records::Records};
 
 use std::io::{self, BufRead, Read, Seek};
 
 use noodles_bgzf as bgzf;
+use noodles_core::Region;
+use noodles_tabix as tabix;
 
 const LINE_FEED: char = '\n';
 const CARRIAGE_RETURN: char = '\r';
@@ -241,6 +244,64 @@ where
     pub fn seek(&mut self, pos: bgzf::VirtualPosition) -> io::Result<bgzf::VirtualPosition> {
         self.inner.seek(pos)
     }
+
+    /// Returns an iterator over records that intersects the given region.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// # use std::{fs::File, io};
+    /// use noodles_core::Region;
+    /// use noodles_bgzf as bgzf;;
+    /// use noodles_tabix as tabix;
+    /// use noodles_vcf as vcf;
+    ///
+    /// let mut reader = File::open("sample.vcf.gz")
+    ///     .map(bgzf::Reader::new)
+    ///     .map(vcf::Reader::new)?;
+    ///
+    /// let index = tabix::read("sample.vcf.gz.tbi")?;
+    /// let region = Region::mapped("sq0", 8, 13);
+    /// let query = reader.query(&index, &region)?;
+    ///
+    /// for result in query {
+    ///     let record = result?;
+    ///     println!("{:?}", record);
+    /// }
+    /// Ok::<(), io::Error>(())
+    /// ```
+    pub fn query(&mut self, index: &tabix::Index, region: &Region) -> io::Result<Query<'_, R>> {
+        let (i, reference_sequence_name, start, end) = resolve_region(index, region)?;
+
+        let index_reference_sequence = index.reference_sequences().get(i).ok_or_else(|| {
+            io::Error::new(
+                io::ErrorKind::InvalidInput,
+                format!(
+                    "could not find reference in index: {} >= {}",
+                    i,
+                    index.reference_sequences().len()
+                ),
+            )
+        })?;
+
+        let query_bins = index_reference_sequence.query(start, end);
+
+        let chunks: Vec<_> = query_bins
+            .iter()
+            .flat_map(|bin| bin.chunks())
+            .cloned()
+            .collect();
+
+        // TODO: optimize chunks
+
+        Ok(Query::new(
+            self,
+            chunks,
+            reference_sequence_name,
+            start,
+            end,
+        ))
+    }
 }
 
 // Reads all bytes until a line feed ('\n') or EOF is reached.
@@ -267,6 +328,32 @@ where
     }
 }
 
+fn resolve_region(index: &tabix::Index, region: &Region) -> io::Result<(usize, String, i32, i32)> {
+    match region {
+        Region::Mapped { name, start, end } => {
+            let i = index
+                .reference_sequence_names()
+                .iter()
+                .position(|n| name == n)
+                .ok_or_else(|| {
+                    io::Error::new(
+                        io::ErrorKind::InvalidInput,
+                        format!(
+                            "region reference sequence does not exist in reference sequences: {:?}",
+                            region
+                        ),
+                    )
+                })?;
+
+            Ok((i, name.into(), *start, *end))
+        }
+        _ => Err(io::Error::new(
+            io::ErrorKind::InvalidData,
+            "region is not mapped",
+        )),
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::io::BufReader;

diff --git a/noodles-vcf/src/reader/query.rs b/noodles-vcf/src/reader/query.rs
@@ -0,0 +1,137 @@
+use std::io::{self, Read, Seek};
+
+use noodles_bgzf as bgzf;
+use noodles_tabix::index::reference_sequence::bin::Chunk;
+
+use crate::{record::Chromosome, Record};
+
+use super::Reader;
+
+enum State {
+    Seek,
+    Read(bgzf::VirtualPosition),
+    End,
+}
+
+/// An iterator over records of a BAM reader that intersects a given region.
+///
+/// This is created by calling [`Reader::query`].
+pub struct Query<'a, R>
+where
+    R: Read + Seek,
+{
+    reader: &'a mut Reader<bgzf::Reader<R>>,
+    chunks: Vec<Chunk>,
+    reference_sequence_name: String,
+    start: i32,
+    end: i32,
+    i: usize,
+    state: State,
+    line_buf: String,
+}
+
+impl<'a, R> Query<'a, R>
+where
+    R: Read + Seek,
+{
+    pub(crate) fn new(
+        reader: &'a mut Reader<bgzf::Reader<R>>,
+        chunks: Vec<Chunk>,
+        reference_sequence_name: String,
+        start: i32,
+        end: i32,
+    ) -> Self {
+        Self {
+            reader,
+            chunks,
+            reference_sequence_name,
+            start,
+            end,
+            i: 0,
+            state: State::Seek,
+            line_buf: String::new(),
+        }
+    }
+
+    fn next_chunk(&mut self) -> io::Result<Option<bgzf::VirtualPosition>> {
+        if self.i >= self.chunks.len() {
+            return Ok(None);
+        }
+
+        let chunk = self.chunks[self.i];
+        self.reader.seek(chunk.start())?;
+
+        self.i += 1;
+
+        Ok(Some(chunk.end()))
+    }
+
+    fn read_and_parse_record(&mut self) -> Option<io::Result<Record>> {
+        self.line_buf.clear();
+
+        match self.reader.read_record(&mut self.line_buf) {
+            Ok(0) => None,
+            Ok(_) => Some(
+                self.line_buf
+                    .parse()
+                    .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e)),
+            ),
+            Err(e) => Some(Err(e)),
+        }
+    }
+}
+
+impl<'a, R> Iterator for Query<'a, R>
+where
+    R: Read + Seek,
+{
+    type Item = io::Result<Record>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            match self.state {
+                State::Seek => {
+                    self.state = match self.next_chunk() {
+                        Ok(Some(chunk_end)) => State::Read(chunk_end),
+                        Ok(None) => State::End,
+                        Err(e) => return Some(Err(e)),
+                    }
+                }
+                State::Read(chunk_end) => match self.read_and_parse_record() {
+                    Some(result) => {
+                        if self.reader.virtual_position() >= chunk_end {
+                            self.state = State::Seek;
+                        }
+
+                        match result {
+                            Ok(record) => {
+                                let reference_sequence_name = match record.chromosome() {
+                                    Chromosome::Name(n) => n.into(),
+                                    Chromosome::Symbol(n) => n.to_string(),
+                                };
+
+                                let start = i32::from(record.position());
+                                let end = start + 1; // TODO: records with END
+
+                                if reference_sequence_name == self.reference_sequence_name
+                                    && in_interval(start, end, self.start, self.end)
+                                {
+                                    return Some(Ok(record));
+                                }
+                            }
+                            Err(e) => return Some(Err(e)),
+                        }
+                    }
+                    None => {
+                        self.state = State::Seek;
+                    }
+                },
+                State::End => return None,
+            }
+        }
+    }
+}
+
+fn in_interval(a_start: i32, a_end: i32, b_start: i32, b_end: i32) -> bool {
+    a_start <= b_end && b_start <= a_end
+}