Fixing local vcf reading with no compression

Author: mwiewior

datafusion/vcf/examples/datafusion_integration.rs

@@ -10,20 +10,21 @@ async fn main() -> datafusion::error::Result<()> {
     // let path = "gs://genomics-public-data/platinum-genomes/vcf/NA12878_S1.genome.vcf".to_string();
     // let path ="/tmp/gnomad.exomes.v4.1.sites.chr21.vcf.bgz".to_string();
     let path ="/tmp/gnomad.v4.1.sv.sites.vcf.gz".to_string();
+    // let path ="/tmp//NA12878_S1.genome.vcf".to_string();
     // let infos  = Some(Vec::from(["AC".to_string(), "AF".to_string(), "AN".to_string(), "FS".to_string(), "AN_raw".to_string(), "variant_type".to_string(), "AS_culprit".to_string(), "only_het".to_string()]));
-    let infos  = Some(Vec::from(["SVTYPE".to_string()]));
-    // let infos  = Some(Vec::from(["AF".to_string()]));
-    // let infos  = None;
+    // let infos  = Some(Vec::from(["SVTYPE".to_string()]));
+    // let infos  = Some(Vec::from(["AF".to_string()])); 50722005
+    let infos  = None;
     // Create a new context with the default configuration
     let ctx = SessionContext::new();
     ctx.sql("set datafusion.execution.skip_physical_aggregate_schema_check=true").await?;
     // let table_provider = VcfTableProvider::new("/tmp/gnomad.exomes.v4.1.sites.chr21.vcf.bgz".parse().unwrap(), vec!["SVTYPE".parse().unwrap()], vec![], Some(8))?;
     let table_provider = VcfTableProvider::new(path, infos, None, Some(4), None,None)?;
     ctx.register_table("custom_table", Arc::new(table_provider)).expect("TODO: panic message");
     // let df = ctx.sql("SELECT svtype, count(*) as cnt FROM custom_table group by svtype").await?;
-    // let df = ctx.sql("SELECT count(*) as cnt FROM custom_table").await?;
+    let df = ctx.sql("SELECT count(*) as cnt FROM custom_table").await?;
     // df.clone().write_csv("/tmp/gnomad.exomes.v4.1.sites.chr21-old.csv", DataFrameWriteOptions::default(), Some(CsvOptions::default())).await?;
-    let df = ctx.sql("SELECT chrom FROM custom_table LIMIT 5").await?;
+    // let df = ctx.sql("SELECT chrom FROM custom_table LIMIT 5").await?;
     // println!("{:?}", df.explain(false, false)?);
     df.show().await.expect("TODO: panic message");
     // println!("{:?}", );

datafusion/vcf/src/physical_exec.rs

@@ -20,7 +20,7 @@ use noodles::vcf::header::Infos;
 use noodles::vcf::variant::Record;
 use noodles::vcf::variant::record::{AlternateBases, Filters, Ids, ReferenceBases};
 use noodles::vcf::variant::record::info::field::{Value, value::Array as ValueArray};
-use crate::storage::{get_local_vcf_bgzf_reader, get_storage_type, StorageType, VcfRemoteReader};
+use crate::storage::{get_local_vcf_bgzf_reader, get_storage_type, StorageType, VcfLocalReader, VcfRemoteReader};
 use crate::table_provider::{info_to_arrow_type, OptionalField};
 
 fn build_record_batch(
@@ -184,64 +184,81 @@ async fn get_local_vcf(file_path: String, schema_ref: SchemaRef,
     let schema = Arc::clone(&schema_ref);
     let file_path = file_path.clone();
     let thread_num = thread_num.unwrap_or(1);
-    let mut reader = get_local_vcf_bgzf_reader(file_path, thread_num)?;
-    let header = reader.read_header()?;
+    let mut reader = VcfLocalReader::new(file_path.clone(), thread_num).await;
+    let header = reader.read_header().await?;
     let infos = header.infos();
-
+    let mut record_num = 0;
     let mut info_builders: (Vec<String>, Vec<DataType>, Vec<OptionalField>) = (Vec::new(), Vec::new(), Vec::new());
     set_info_builders(batch_size, info_fields, &infos, &mut info_builders);
 
-    let iter  = std::iter::from_fn(move || {
+    let stream = try_stream! {
 
-        let mut records = reader.records();
+        let mut records = reader.read_records();
         let iter_start_time = Instant::now();
-        while count < batch_size {
-            let record = records.next();
-            if record.is_none() {
-                break;
-            }
-            let record = record.unwrap().unwrap();
-            // For each record, fill the fixed columns.
+        while let Some(result) = records.next().await {
+            let record = result?;  // propagate errors if any
             chroms.push(record.reference_sequence_name().to_string());
-            poss.push(record.variant_start().unwrap().unwrap().get() as u32);
+            poss.push(record.variant_start().unwrap()?.get() as u32);
             pose.push(get_variant_end(&record, &header));
             ids.push(record.ids().iter().map(|v| v.to_string()).collect::<Vec<String>>().join(";"));
             refs.push(record.reference_bases().to_string());
             alts.push(record.alternate_bases().iter().map(|v| v.unwrap_or(".").to_string()).collect::<Vec<String>>().join("|"));
             quals.push(record.quality_score().unwrap_or(Ok(0.0)).unwrap() as f64);
             filters.push(record.filters().iter(&header).map(|v| v.unwrap_or(".").to_string()).collect::<Vec<String>>().join(";"));
             load_infos(Box::new(record), &header, &mut info_builders);
-            count += 1;
+            record_num += 1;
+            // Once the batch size is reached, build and yield a record batch.
+            if record_num % batch_size == 0 {
+                debug!("Record number: {}", record_num);
+                let batch = build_record_batch(
+                    Arc::clone(&schema.clone()),
+                    &chroms,
+                    &poss,
+                    &pose,
+                    &ids,
+                    &refs,
+                    &alts,
+                    &quals,
+                    &filters,
+                    Some(&builders_to_arrays(&mut info_builders.2)), projection.clone(),
+                    // if infos.is_empty() { None } else { Some(&infos) },
+
+                )?;
+                batch_num += 1;
+                debug!("Batch number: {}", batch_num);
+                yield batch;
+                // Clear vectors for the next batch.
+                chroms.clear();
+                poss.clear();
+                pose.clear();
+                ids.clear();
+                refs.clear();
+                alts.clear();
+                quals.clear();
+                filters.clear();
+
+            }
         }
-        if count == 0 {
-            return None;
+        // If there are remaining records that don't fill a complete batch,
+        // yield them as well.
+        if !chroms.is_empty() {
+            let batch = build_record_batch(
+                Arc::clone(&schema.clone()),
+                &chroms,
+                &poss,
+                &pose,
+                &ids,
+                &refs,
+                &alts,
+                &quals,
+                &filters,
+                Some(&builders_to_arrays(&mut info_builders.2)), projection.clone(),
+                // if infos.is_empty() { None } else { Some(&infos) },
+            )?;
+            yield batch;
         }
-        let duration = iter_start_time.elapsed();
-        debug!("Time elapsed in iterating records: {:?}", duration);
-        debug!("Batch number: {}", batch_num);
-        let start_time = Instant::now();
-        let batch = build_record_batch(Arc::clone(&schema),
-                                       &chroms, &poss, &pose,
-                                       &ids, &refs, &alts,
-                                       &quals, &filters,
-                                       Some(&builders_to_arrays(&mut info_builders.2)),
-                                       projection.clone()
-        ).unwrap();
-        let duration = start_time.elapsed();
-        debug!("Time elapsed in building batch: {:?}", duration);
-        count = 0;
-        chroms.clear();
-        poss.clear();
-        pose.clear();
-        ids.clear();
-        refs.clear();
-        alts.clear();
-        quals.clear();
-        filters.clear();
-        batch_num += 1;
-        Some(Ok(batch))
-    });
-    Ok(stream::iter(iter))
+    };
+    Ok(stream)
 }
 
 

datafusion/vcf/src/storage.rs

@@ -1,8 +1,9 @@
 use std::fs::File;
 use std::io::Error;
 use std::num::NonZero;
+use async_stream::stream;
 use bytes::Bytes;
-use futures::StreamExt;
+use futures::{stream, StreamExt};
 use futures::stream::BoxStream;
 use log::debug;
 use noodles::{bgzf, vcf};
@@ -244,6 +245,8 @@ pub enum VcfRemoteReader {
     PLAIN( vcf::r#async::io::Reader<StreamReader<FuturesBytesStream, Bytes>>)
 }
 
+
+
 impl VcfRemoteReader {
     pub async fn new(file_path: String, chunk_size: usize, concurrent_fetches: usize) -> Self {
         let compression_type = get_compression_type(file_path.clone());
@@ -281,3 +284,44 @@ impl VcfRemoteReader {
     }
 }
 
+pub enum VcfLocalReader {
+    BGZF(Reader<MultithreadedReader<File>>),
+    PLAIN(vcf::r#async::io::Reader<BufReader<tokio::fs::File>>)
+}
+
+impl VcfLocalReader {
+    pub async fn new(file_path: String, thread_num: usize) -> Self {
+        let compression_type = get_compression_type(file_path.clone());
+        match compression_type {
+            CompressionType::BGZF | CompressionType::GZIP=> {
+                let reader = get_local_vcf_bgzf_reader(file_path, thread_num).unwrap();
+                VcfLocalReader::BGZF(reader)
+            }
+            CompressionType::NONE => {
+                let reader = get_local_vcf_reader(file_path).await.unwrap();
+                VcfLocalReader::PLAIN(reader)
+            }
+        }
+    }
+    pub async fn read_header(&mut self) -> Result<vcf::Header, Error> {
+        match self {
+            VcfLocalReader::BGZF(reader) => {
+                reader.read_header()
+            }
+            VcfLocalReader::PLAIN(reader) => {
+                reader.read_header().await
+            }
+        }
+    }
+
+    pub fn read_records(&mut self) -> BoxStream<'_, Result<Record, Error>> {
+        match self {
+            VcfLocalReader::BGZF(reader) => {
+                stream::iter(reader.records()).boxed()
+            }
+            VcfLocalReader::PLAIN(reader) => {
+                reader.records().boxed()
+            }
+        }
+    }
+}