Fixing local vcf without compression and gcs reads optimization

mwiewior · mwiewior · commit 6797e9a72d29 · 2025-03-04T20:25:21.000+01:00
diff --git a/datafusion/vcf/examples/datafusion_integration.rs b/datafusion/vcf/examples/datafusion_integration.rs
@@ -1,6 +1,4 @@
 use std::sync::Arc;
-use datafusion::config::CsvOptions;
-use datafusion::dataframe::DataFrameWriteOptions;
 use datafusion::prelude::SessionContext;
 use datafusion_vcf::table_provider::VcfTableProvider;
 
@@ -20,12 +18,12 @@ async fn main() -> datafusion::error::Result<()> {
     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))?;
+    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 * 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!("{:?}", );
diff --git a/datafusion/vcf/examples/noodles_test.rs b/datafusion/vcf/examples/noodles_test.rs
@@ -67,10 +67,10 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 }
 
 
-// let builder = Gcs::default()
-//     .bucket(BUCKET)
-//     .disable_vm_metadata()
-//     .allow_anonymous();
+let builder = Gcs::default()
+    .bucket(BUCKET)
+    .disable_vm_metadata()
+    .allow_anonymous();
 //
 // let operator = Operator::new(builder)?.finish();
 //
diff --git a/datafusion/vcf/src/physical_exec.rs b/datafusion/vcf/src/physical_exec.rs
@@ -2,37 +2,25 @@ use std::any::Any;
 use std::fmt::{Debug, Formatter};
 use std::sync::Arc;
 
-use datafusion::arrow::array::{Array, Float64Array, Int32Array, NullArray, StringArray, UInt32Array};
-use datafusion::arrow::datatypes::{DataType, Field, SchemaRef};
+use datafusion::arrow::array::{Array, Float64Array, NullArray, StringArray, UInt32Array};
+use datafusion::arrow::datatypes::{DataType, SchemaRef};
 use datafusion::arrow::record_batch::RecordBatch;
 use datafusion::common::DataFusionError;
 use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
 use datafusion_execution::{SendableRecordBatchStream, TaskContext};
 use futures::{stream, StreamExt, TryStreamExt};
-use std::{io, str};
-use std::fs::File;
-use std::hash::Hasher;
-use std::io::Error;
-use std::num::NonZero;
-use std::ops::Deref;
+use std::str;
 use std::time::Instant;
 use async_stream::__private::AsyncStream;
 use async_stream::try_stream;
-use datafusion::arrow::ipc::FieldBuilder;
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
-use env_logger::builder;
-use log::{debug, info};
-use noodles::vcf;
+use log::debug;
 use noodles::vcf::Header;
 use noodles::vcf::header::Infos;
-use noodles::vcf::header::record::value::map::Info;
-use noodles::vcf::header::record::value::map::info::{Number, Type};
-use noodles::vcf::io::Reader;
 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 noodles_bgzf::MultithreadedReader;
-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::storage::{get_local_vcf_bgzf_reader, get_storage_type, StorageType, VcfRemoteReader};
 use crate::table_provider::{info_to_arrow_type, OptionalField};
 
 fn build_record_batch(
@@ -192,7 +180,6 @@ async fn get_local_vcf(file_path: String, schema_ref: SchemaRef,
     let mut filters: Vec<String> = Vec::with_capacity(batch_size);
 
     let mut count: usize = 0;
-    let mut record_num = 0;
     let mut batch_num = 0;
     let schema = Arc::clone(&schema_ref);
     let file_path = file_path.clone();
@@ -210,7 +197,6 @@ async fn get_local_vcf(file_path: String, schema_ref: SchemaRef,
         let iter_start_time = Instant::now();
         while count < batch_size {
             let record = records.next();
-            record_num += 1;
             if record.is_none() {
                 break;
             }
@@ -262,8 +248,8 @@ async fn get_local_vcf(file_path: String, schema_ref: SchemaRef,
 
 async fn get_remote_vcf_stream(file_path: String, schema: SchemaRef,
                                batch_size: usize,
-                               info_fields: Option<Vec<String>>, projection: Option<Vec<usize>>) -> datafusion::error::Result<AsyncStream<datafusion::error::Result<RecordBatch>, impl Future<Output=()> + Sized>> {
-    let mut reader = VcfRemoteReader::new(file_path.clone()).await;
+                               info_fields: Option<Vec<String>>, projection: Option<Vec<usize>>, chunk_size: usize, concurrent_fetches: usize) -> datafusion::error::Result<AsyncStream<datafusion::error::Result<RecordBatch>, impl Future<Output=()> + Sized>> {
+    let mut reader = VcfRemoteReader::new(file_path.clone(), chunk_size, concurrent_fetches).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());
@@ -370,7 +356,7 @@ fn set_info_builders(batch_size: usize, info_fields: Option<Vec<String>>, infos:
 
 async fn get_stream(file_path: String, schema_ref: SchemaRef, batch_size: usize,
                     thread_num: Option<usize>,
-                    info_fields: Option<Vec<String>>, projection: Option<Vec<usize>>) -> datafusion::error::Result<SendableRecordBatchStream> {
+                    info_fields: Option<Vec<String>>, projection: Option<Vec<usize>>, chunk_size: usize, concurrent_fetches: usize) -> datafusion::error::Result<SendableRecordBatchStream> {
     // Open the BGZF-indexed VCF using IndexedReader.
 
     let file_path = file_path.clone();
@@ -383,7 +369,7 @@ async fn get_stream(file_path: String, schema_ref: SchemaRef, batch_size: usize,
             Ok(Box::pin(RecordBatchStreamAdapter::new(schema_ref, stream)))
         },
         StorageType::GCS | StorageType::S3=> {
-            let stream = get_remote_vcf_stream(file_path.clone(), schema.clone(), batch_size, info_fields, projection).await?;
+            let stream = get_remote_vcf_stream(file_path.clone(), schema.clone(), batch_size, info_fields, projection, chunk_size, concurrent_fetches).await?;
             Ok(Box::pin(RecordBatchStreamAdapter::new(schema_ref, stream)))
         },
         _ => panic!("Unsupported storage type")
@@ -402,6 +388,8 @@ pub struct VcfExec {
     pub(crate) cache: PlanProperties,
     pub(crate) limit: Option<usize>,
     pub(crate) thread_num: Option<usize>,
+    pub(crate) chunk_size: Option<usize>,
+    pub(crate) concurrent_fetches: Option<usize>,
 }
 
 
@@ -415,6 +403,8 @@ impl VcfExec {
         cache: PlanProperties,
         limit: Option<usize>,
         thread_num: Option<usize>,
+        chunk_size: Option<usize>,
+        concurrent_fetches: Option<usize>,
     ) -> Self {
         debug!("VcfExec::new");
         Self {
@@ -426,21 +416,23 @@ impl VcfExec {
             cache,
             limit,
             thread_num,
+            chunk_size,
+            concurrent_fetches
         }
     }
 }
 
 
 impl Debug for VcfExec {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+    fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result {
         Ok(())
     }
 }
 
 
 
 impl DisplayAs for VcfExec {
-    fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
+    fn fmt_as(&self, _t: DisplayFormatType, _f: &mut Formatter) -> std::fmt::Result {
         Ok(())
     }
 
@@ -464,7 +456,7 @@ impl ExecutionPlan for VcfExec {
         vec![]
     }
 
-    fn with_new_children(self: Arc<Self>, children: Vec<Arc<dyn ExecutionPlan>>) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
+    fn with_new_children(self: Arc<Self>, _children: Vec<Arc<dyn ExecutionPlan>>) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
         Ok(self)
     }
 
@@ -479,7 +471,9 @@ impl ExecutionPlan for VcfExec {
                              schema.clone(),
                              batch_size, self.thread_num,
                              self.info_fields.clone(),
-                             self.projection.clone());
+                             self.projection.clone(),
+                             self.chunk_size.unwrap_or(64),
+                             self.concurrent_fetches.unwrap_or(8));
         let stream = futures::stream::once(fut).try_flatten();
         Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream)))
 
diff --git a/datafusion/vcf/src/storage.rs b/datafusion/vcf/src/storage.rs
@@ -1,10 +1,8 @@
 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::StreamExt;
 use futures::stream::BoxStream;
 use log::debug;
 use noodles::{bgzf, vcf};
@@ -14,7 +12,7 @@ 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::io::BufReader;
 use tokio_util::io::StreamReader;
 
 
@@ -95,8 +93,8 @@ pub fn get_compression_type(file_path: String) -> CompressionType {
 
 
 
-pub async fn get_remote_stream_bgzf(file_path: String) ->  Result<AsyncReader<StreamReader<FuturesBytesStream, bytes::Bytes>>, opendal::Error> {
-    let remote_stream = StreamReader::new(get_remote_stream(file_path.clone()).await?);
+pub async fn get_remote_stream_bgzf(file_path: String, chunk_size: usize, concurrent_fetches: usize) ->  Result<AsyncReader<StreamReader<FuturesBytesStream, bytes::Bytes>>, opendal::Error> {
+    let remote_stream = StreamReader::new(get_remote_stream(file_path.clone(), chunk_size, concurrent_fetches).await?);
     Ok(bgzf::r#async::Reader::new(remote_stream))
 
 }
@@ -127,7 +125,7 @@ fn get_bucket_name(file_path: String) -> String {
     bucket_name.to_string()
 }
 
-pub async fn get_remote_stream(file_path: String) ->  Result<FuturesBytesStream, opendal::Error> {
+pub async fn get_remote_stream(file_path: String, chunk_size: usize, concurrent_fetches: usize) ->  Result<FuturesBytesStream, opendal::Error> {
     let storage_type = get_storage_type(file_path.clone());
     let bucket_name = get_bucket_name(file_path.clone());
     let file_path = get_file_path(file_path.clone());
@@ -142,7 +140,10 @@ pub async fn get_remote_stream(file_path: String) ->  Result<FuturesBytesStream,
                 .layer(RetryLayer::new().with_max_times(3))
                 .layer(LoggingLayer::default())
                 .finish();
-            operator.reader_with(file_path.as_str()).concurrent(1).await?.into_bytes_stream(..).await
+            operator.reader_with(file_path.as_str())
+                .chunk(chunk_size * 1024 * 1024)
+                .concurrent(concurrent_fetches)
+                .await?.into_bytes_stream(..).await
         }
         StorageType::S3 => {
             let builder = S3::default()
@@ -162,15 +163,15 @@ pub async fn get_remote_stream(file_path: String) ->  Result<FuturesBytesStream,
     }
 }
 
-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);
+pub async fn get_remote_vcf_bgzf_reader(file_path: String, chunk_size: usize, concurrent_fetches: usize) -> vcf::r#async::io::Reader<AsyncReader<StreamReader<FuturesBytesStream, Bytes>>> {
+    let inner = get_remote_stream_bgzf(file_path.clone(), chunk_size, concurrent_fetches).await.unwrap();
+    let 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);
+pub async fn get_remote_vcf_reader(file_path: String, chunk_size: usize, concurrent_fetches: usize) -> vcf::r#async::io::Reader<StreamReader<FuturesBytesStream, Bytes>> {
+    let inner = StreamReader::new(get_remote_stream(file_path.clone(), chunk_size, concurrent_fetches).await.unwrap());
+    let reader = vcf::r#async::io::Reader::new(inner);
     reader
 }
 
@@ -187,7 +188,7 @@ pub fn get_local_vcf_bgzf_reader(file_path: String, thread_num: usize) -> Result
 
 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")
+    let reader = tokio::fs::File::open(file_path)
         .await
         .map(BufReader::new)
         .map(vcf::r#async::io::Reader::new)?;
@@ -210,15 +211,15 @@ pub async fn get_local_vcf_header(file_path: String, thread_num: usize) -> Resul
     Ok(header)
 }
 
-pub async fn get_remote_vcf_header(file_path: String) -> Result<vcf::Header, Error> {
+pub async fn get_remote_vcf_header(file_path: String, chunk_size: usize, concurrent_fetches: 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_remote_vcf_bgzf_reader(file_path).await;
+            let mut reader = get_remote_vcf_bgzf_reader(file_path, chunk_size, concurrent_fetches).await;
             reader.read_header().await?
         }
         CompressionType::NONE => {
-            let mut reader = get_remote_vcf_reader(file_path).await;
+            let mut reader = get_remote_vcf_reader(file_path, chunk_size, concurrent_fetches).await;
             reader.read_header().await?
         }
     };
@@ -232,7 +233,7 @@ pub async fn get_header(file_path: String) -> Result<vcf::Header, Error> {
             get_local_vcf_header(file_path, 1).await?
         }
         _ => {
-            get_remote_vcf_header(file_path).await?
+            get_remote_vcf_header(file_path, 64,1).await?
         }
     };
     Ok(header)
@@ -244,15 +245,15 @@ pub enum VcfRemoteReader {
 }
 
 impl VcfRemoteReader {
-    pub async fn new(file_path: String) -> Self {
+    pub async fn new(file_path: String, chunk_size: usize, concurrent_fetches: usize) -> 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;
+                let reader = get_remote_vcf_bgzf_reader(file_path, chunk_size, concurrent_fetches).await;
                 VcfRemoteReader::BGZF(reader)
             }
             CompressionType::NONE => {
-                let reader = get_remote_vcf_reader(file_path).await;
+                let reader = get_remote_vcf_reader(file_path, chunk_size, concurrent_fetches).await;
                 VcfRemoteReader::PLAIN(reader)
             }
         }
diff --git a/datafusion/vcf/src/table_provider.rs b/datafusion/vcf/src/table_provider.rs
