Adding support for remote reading of uncompressed VCFs

mwiewior · mwiewior · commit 95e6854435d5 · 2025-02-28T12:09:13.000+01:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/datafusion/vcf/Cargo.toml b/datafusion/vcf/Cargo.toml
@@ -14,7 +14,7 @@ noodles = { version = "0.91.0",  features = ["bam", "sam", "vcf", "bgzf", "async
 noodles-bgzf = { version = "0.36.0" ,features = ["libdeflate"] }
 tokio = { version = "1.43.0", features = ["rt-multi-thread", "rt", "macros"] }
 object_store = {version = "0.11.2", features = ["gcp"] }
-tokio-util = {  version="0.7.13", features =["io-util"] }
+tokio-util = {  version="0.7.13", features = ["io-util", "compat"] }
 bytes = "1.10.0"
 opendal = { version = "0.51.2", features = ["services-gcs", "services-s3","layers-blocking"] }
 log = "0.4.22"
diff --git a/datafusion/vcf/examples/bgzf.rs b/datafusion/vcf/examples/bgzf.rs
@@ -7,12 +7,6 @@ use noodles::{bgzf, vcf};
 use noodles::vcf::io::Reader;
 use noodles_bgzf::MultithreadedReader;
 
-pub fn get_local_vcf_reader(file_path: String, thread_num: usize) -> Result<Reader<MultithreadedReader<File>>, Error> {
-    debug!("Reading VCF file from local storage with {} threads", thread_num);
-    File::open(file_path)
-        .map(|f| noodles_bgzf::MultithreadedReader::with_worker_count(NonZero::new(thread_num).unwrap(), f))
-        .map(vcf::io::Reader::new)
-}
 
 #[tokio::main(flavor = "multi_thread")]
 async fn main() -> datafusion::error::Result<()> {
diff --git a/datafusion/vcf/examples/datafusion_integration.rs b/datafusion/vcf/examples/datafusion_integration.rs
@@ -7,7 +7,8 @@ async fn main() -> datafusion::error::Result<()> {
     env_logger::init();
     // let path = "gs://gcp-public-data--gnomad/release/4.1/vcf/exomes/gnomad.exomes.v4.1.sites.chr21.vcf.bgz".to_string();
     // let path = "gs://gcp-public-data--gnomad/release/4.1/genome_sv/gnomad.v4.1.sv.sites.vcf.gz".to_string();
-    let path ="/tmp/gnomad.exomes.v4.1.sites.chr21.vcf.bgz".to_string();
+    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 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()]));
diff --git a/datafusion/vcf/src/physical_exec.rs b/datafusion/vcf/src/physical_exec.rs
@@ -32,7 +32,7 @@ use noodles::vcf::variant::Record;
 use noodles::vcf::variant::record::{AlternateBases, Filters, Ids};
 use noodles::vcf::variant::record::info::field::{Value, value::Array as ValueArray};
 use noodles_bgzf::MultithreadedReader;
-use crate::storage::{get_local_vcf_reader, get_remote_stream_bgzf, get_remote_vcf_reader, get_storage_type, StorageType};
+use crate::storage::{get_compression_type, get_local_vcf_bgzf_reader, get_remote_stream_bgzf, get_remote_vcf_bgzf_reader, get_remote_vcf_header, get_remote_vcf_reader, get_storage_type, CompressionType, StorageType, VcfRemoteReader};
 use crate::table_provider::{info_to_arrow_type, OptionalField};
 
 fn build_record_batch(
@@ -148,7 +148,7 @@ async fn get_local_vcf(file_path: String, schema_ref: SchemaRef, batch_size: usi
     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_reader(file_path, thread_num)?;
+    let mut reader = get_local_vcf_bgzf_reader(file_path, thread_num)?;
     let header = reader.read_header()?;
     let infos = header.infos();
 
@@ -206,12 +206,14 @@ async fn get_local_vcf(file_path: String, schema_ref: SchemaRef, batch_size: usi
 
 
 async fn get_remote_vcf_stream(file_path: String, schema: SchemaRef, batch_size: usize, info_fields: Option<Vec<String>>) -> datafusion::error::Result<AsyncStream<datafusion::error::Result<RecordBatch>, impl Future<Output=()> + Sized>> {
-    let mut reader = get_remote_vcf_reader(file_path.clone()).await;
+    let mut reader = VcfRemoteReader::new(file_path.clone()).await;
     let header = reader.read_header().await?;
     let infos = header.infos();
     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 stream = try_stream! {
         // Create vectors for accumulating record data.
         let mut chroms: Vec<String> = Vec::with_capacity(batch_size);
@@ -231,7 +233,8 @@ async fn get_remote_vcf_stream(file_path: String, schema: SchemaRef, batch_size:
 
 
         // Process records one by one.
-        let mut records = reader.records();
+
+        let mut records = reader.read_records().await;
         while let Some(result) = records.next().await {
             let record = result?;  // propagate errors if any
             chroms.push(record.reference_sequence_name().to_string());
@@ -244,9 +247,6 @@ async fn get_remote_vcf_stream(file_path: String, schema: SchemaRef, batch_size:
             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);
             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);
diff --git a/datafusion/vcf/src/storage.rs b/datafusion/vcf/src/storage.rs
@@ -1,14 +1,20 @@
 use std::fs::File;
+use std::io;
 use std::io::Error;
 use std::num::NonZero;
 use std::sync::Arc;
+use bytes::Bytes;
+use futures::{Stream, StreamExt};
+use futures::stream::BoxStream;
 use log::debug;
 use noodles::{bgzf, vcf};
 use noodles::vcf::io::Reader;
+use noodles::vcf::Record;
 use noodles_bgzf::{AsyncReader, MultithreadedReader};
 use opendal::{FuturesBytesStream, Operator};
 use opendal::layers::{LoggingLayer, RetryLayer, TimeoutLayer};
 use opendal::services::{Gcs, S3};
+use tokio::io::{AsyncRead, BufReader};
 use tokio_util::io::StreamReader;
 
 
@@ -28,7 +34,7 @@ impl CompressionType {
     }
 
     fn from_string(compression_type: String) -> Self {
-        match compression_type.as_str() {
+        match compression_type.to_lowercase().as_str() {
             "gz" => CompressionType::GZIP,
             "bgz" => CompressionType::BGZF,
             "none" => CompressionType::NONE,
@@ -79,6 +85,9 @@ fn get_file_path(file_path: String) -> String {
 
 pub fn get_compression_type(file_path: String) -> CompressionType {
     //extract the file extension from path
+    if file_path.to_lowercase().ends_with(".vcf") {
+        return CompressionType::NONE;
+    }
     let file_extension = file_path.split('.').last().unwrap();
     //return the compression type
     CompressionType::from_string(file_extension.to_string())
@@ -153,20 +162,108 @@ pub async fn get_remote_stream(file_path: String) ->  Result<FuturesBytesStream,
     }
 }
 
-pub async fn get_remote_vcf_reader(file_path: String) -> vcf::r#async::io::Reader<AsyncReader<StreamReader<FuturesBytesStream, bytes::Bytes>>> {
+pub async fn get_remote_vcf_bgzf_reader(file_path: String) -> vcf::r#async::io::Reader<AsyncReader<StreamReader<FuturesBytesStream, Bytes>>> {
     let inner = get_remote_stream_bgzf(file_path.clone()).await.unwrap();
     let mut reader = vcf::r#async::io::Reader::new(inner);
     reader
 }
 
+pub async fn get_remote_vcf_reader(file_path: String) -> vcf::r#async::io::Reader<StreamReader<FuturesBytesStream, Bytes>> {
+    let inner = StreamReader::new(get_remote_stream(file_path.clone()).await.unwrap());
+    let mut reader = vcf::r#async::io::Reader::new(inner);
+    reader
+}
+
+
 
-pub fn get_local_vcf_reader(file_path: String, thread_num: usize) -> Result<Reader<MultithreadedReader<File>>, Error> {
+
+pub fn get_local_vcf_bgzf_reader(file_path: String, thread_num: usize) -> Result<Reader<MultithreadedReader<File>>, Error> {
     debug!("Reading VCF file from local storage with {} threads", thread_num);
     File::open(file_path)
         .map(|f| noodles_bgzf::MultithreadedReader::with_worker_count(NonZero::new(thread_num).unwrap(), f))
         .map(vcf::io::Reader::new)
 }
 
 
+pub async fn get_local_vcf_reader(file_path: String) -> Result<vcf::r#async::io::Reader<BufReader<tokio::fs::File>>, Error> {
+    debug!("Reading VCF file from local storage with async reader");
+    let reader = tokio::fs::File::open("sample.vcf")
+        .await
+        .map(BufReader::new)
+        .map(vcf::r#async::io::Reader::new)?;
+    Ok(reader)
+}
+
+
+pub async fn get_local_vcf_header(file_path: String, thread_num: usize) -> Result<vcf::Header, Error> {
+    let compression_type = get_compression_type(file_path.clone());
+    let header = match compression_type {
+        CompressionType::BGZF | CompressionType::GZIP=> {
+            let mut reader = get_local_vcf_bgzf_reader(file_path, thread_num)?;
+            reader.read_header()?
+        }
+        CompressionType::NONE => {
+            let mut reader = get_local_vcf_reader(file_path).await?;
+            reader.read_header().await?
+        }
+    };
+    Ok(header)
+}
 
+pub async fn get_remote_vcf_header(file_path: String) -> Result<vcf::Header, Error> {
+    let compression_type = get_compression_type(file_path.clone());
+    let header = match compression_type {
+        CompressionType::BGZF | CompressionType::GZIP=> {
+            let mut reader = get_remote_vcf_bgzf_reader(file_path).await;
+            reader.read_header().await?
+        }
+        CompressionType::NONE => {
+            let mut reader = get_remote_vcf_reader(file_path).await;
+            reader.read_header().await?
+        }
+    };
+    Ok(header)
+}
+
+pub enum VcfRemoteReader {
+    BGZF( vcf::r#async::io::Reader<AsyncReader<StreamReader<FuturesBytesStream, Bytes>>>),
+    PLAIN( vcf::r#async::io::Reader<StreamReader<FuturesBytesStream, Bytes>>)
+}
+
+impl VcfRemoteReader {
+    pub async fn new(file_path: String) -> Self {
+        let compression_type = get_compression_type(file_path.clone());
+        match compression_type {
+            CompressionType::BGZF | CompressionType::GZIP=> {
+                let reader = get_remote_vcf_bgzf_reader(file_path).await;
+                VcfRemoteReader::BGZF(reader)
+            }
+            CompressionType::NONE => {
+                let reader = get_remote_vcf_reader(file_path).await;
+                VcfRemoteReader::PLAIN(reader)
+            }
+        }
+    }
+    pub async fn read_header(&mut self) -> Result<vcf::Header, Error> {
+        match self {
+            VcfRemoteReader::BGZF(reader) => {
+                reader.read_header().await
+            }
+            VcfRemoteReader::PLAIN(reader) => {
+                reader.read_header().await
+            }
+        }
+    }
+
+    pub async fn read_records(&mut self) -> BoxStream<'_, Result<Record, Error>> {
+        match self {
+            VcfRemoteReader::BGZF(reader) => {
+                reader.records().boxed()
+            }
+            VcfRemoteReader::PLAIN(reader) => {
+                reader.records().boxed()
+            }
+        }
+    }
+}
 
diff --git a/datafusion/vcf/src/table_provider.rs b/datafusion/vcf/src/table_provider.rs
@@ -17,7 +17,7 @@ use noodles::vcf::Header;
 use noodles::vcf::header::Infos;
 use noodles::vcf::header::record::value::map::info::{Number, Type};
 use crate::physical_exec::VcfExec;
-use crate::storage::{get_local_vcf_reader, get_remote_vcf_reader, get_storage_type, StorageType};
+use crate::storage::{get_local_vcf_bgzf_reader, get_local_vcf_header, get_remote_vcf_bgzf_reader, get_remote_vcf_header, get_storage_type, StorageType};
 
 async fn determine_schema_from_header(
     file_path: &str,
@@ -26,10 +26,7 @@ async fn determine_schema_from_header(
 ) -> datafusion::common::Result<SchemaRef> {
 
     let storage_type = get_storage_type(file_path.to_string());
-    let header = match storage_type {
-        StorageType::LOCAL =>  get_local_vcf_reader(file_path.to_string(), 1)?.read_header()?,
-        _ => get_remote_vcf_reader(file_path.to_string()).await.read_header().await?
-    };
+    let header = get_remote_vcf_header(file_path.to_string()).await?;
     let header_infos = header.infos();
 
     let mut fields = vec![
@@ -64,6 +61,7 @@ pub struct VcfTableProvider {
     format_fields: Option<Vec<String>>,
     schema: SchemaRef,
     thread_num: Option<usize>,
+
 }
 
 impl VcfTableProvider {
@@ -73,7 +71,7 @@ impl VcfTableProvider {
         format_fields: Option<Vec<String>>,
         thread_num: Option<usize>,
     ) -> datafusion::common::Result<Self> {
-        let schema = block_on(determine_schema_from_header(&file_path, &info_fields, &format_fields)).unwrap();
+        let schema = block_on(determine_schema_from_header(&file_path, &info_fields, &format_fields))?;
         Ok(Self {
             file_path,
             info_fields,
@@ -101,6 +99,7 @@ impl TableProvider for VcfTableProvider {
     }
 
     async fn scan(&self, state: &dyn Session, projection: Option<&Vec<usize>>, _filters: &[Expr], limit: Option<usize>) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
+
         Ok(Arc::new(VcfExec {
             cache: PlanProperties::new( EquivalenceProperties::new(self.schema.clone()),
                                         Partitioning::UnknownPartitioning(1),
