refactor: use ExonArrayBuilder for bcf (#484)

wheretrue · Apr 22, 2024 · b7af849 · b7af849
1 parent 0b50da3
commit b7af849
Show file tree

Hide file tree

Showing 7 changed files with 97 additions and 34 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -38,7 +38,7 @@ arrow = { version = "51.0.0" }
 async-trait = "0.1.80"
 datafusion = { version = "37", features = ["compression", "parquet"] }
 futures = "0.3"
-noodles = { version = "0.69" }
+noodles = { version = "0.70" }
 object_store = { version = "0.9" }
 tokio = { version = "1", features = ["io-util"] }
 tokio-util = { version = "0.7.10", features = ["compat"] }

diff --git a/exon/exon-bcf/src/batch_reader.rs b/exon/exon-bcf/src/batch_reader.rs
@@ -14,6 +14,8 @@
 
 use std::sync::Arc;
 
+use exon_common::ExonArrayBuilder;
+
 use arrow::{error::ArrowError, record_batch::RecordBatch};
 
 use exon_vcf::VCFArrayBuilder;
@@ -50,7 +52,6 @@ where
             reader,
             config,
             header: Arc::new(header),
-            // string_maps,
         })
     }
 
@@ -136,7 +137,7 @@ impl BatchAdapter {
         let mut record_batch = VCFArrayBuilder::create(
             self.config.file_schema.clone(),
             self.config.batch_size,
-            self.config.projection.as_deref(),
+            self.config.projection.clone(),
             self.header.clone(),
         )?;
 
@@ -152,7 +153,8 @@ impl BatchAdapter {
             return Ok(None);
         }
 
-        let batch = RecordBatch::try_new(self.config.file_schema.clone(), record_batch.finish())?;
+        let schema = self.config.projected_schema()?;
+        let batch = record_batch.try_into_record_batch(schema)?;
 
         match &self.config.projection {
             Some(projection) => Ok(Some(batch.project(projection)?)),

diff --git a/exon/exon-bcf/src/config.rs b/exon/exon-bcf/src/config.rs
@@ -61,4 +61,18 @@ impl BCFConfig {
         self.projection = projection;
         self
     }
+
+    /// Get the projection, returning the identity projection if none is set.
+    pub fn projection(&self) -> Vec<usize> {
+        self.projection
+            .clone()
+            .unwrap_or_else(|| (0..self.file_schema.fields().len()).collect())
+    }
+
+    /// Get the projected schema.
+    pub fn projected_schema(&self) -> arrow::error::Result<SchemaRef> {
+        let schema = self.file_schema.project(&self.projection())?;
+
+        Ok(Arc::new(schema))
+    }
 }
diff --git a/exon/exon-core/src/datasources/bcf/table_provider.rs b/exon/exon-core/src/datasources/bcf/table_provider.rs
@@ -261,6 +261,16 @@ impl<T: ExonIndexedListingOptions + 'static> TableProvider for ListingBCFTable<T
         .limit_option(limit)
         .build();
 
+        if let Some(region) = self.config.options.regions().first() {
+            let plan = self
+                .config
+                .options
+                .create_physical_plan_with_regions(file_scan_config, vec![region.clone()])
+                .await?;
+
+            return Ok(plan);
+        }
+
         let plan = self
             .config
             .options

diff --git a/exon/exon-core/src/session_context/exon_context_ext.rs b/exon/exon-core/src/session_context/exon_context_ext.rs
@@ -663,7 +663,8 @@ mod tests {
 
     use crate::{
         datasources::{
-            bigwig, fasta::table_provider::ListingFASTATableOptions,
+            bcf::table_provider::ListingBCFTableOptions, bigwig,
+            fasta::table_provider::ListingFASTATableOptions,
             fastq::table_provider::ListingFASTQTableOptions,
         },
         session_context::ExonSessionExt,
@@ -833,4 +834,42 @@ mod tests {
 
         Ok(())
     }
+
+    #[tokio::test]
+    async fn test_bcf_file() -> Result<(), Box<dyn std::error::Error>> {
+        let ctx = SessionContext::new_exon();
+
+        let bcf_path = exon_test::test_path("bcf", "index.bcf");
+
+        let df = ctx
+            .read_bcf(
+                bcf_path.to_str().unwrap(),
+                ListingBCFTableOptions::default(),
+            )
+            .await?;
+
+        assert_eq!(df.count().await?, 621);
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_bcf_file_with_region() -> Result<(), Box<dyn std::error::Error>> {
+        let ctx = SessionContext::new_exon();
+
+        let bcf_path = exon_test::test_path("bcf", "index.bcf");
+
+        let region = "1".parse()?;
+
+        let df = ctx
+            .read_bcf(
+                bcf_path.to_str().unwrap(),
+                ListingBCFTableOptions::default().with_regions(vec![region]),
+            )
+            .await?;
+
+        assert_eq!(df.count().await?, 191);
+
+        Ok(())
+    }
 }
diff --git a/exon/exon-vcf/src/array_builder/eager_array_builder.rs b/exon/exon-vcf/src/array_builder/eager_array_builder.rs
@@ -15,13 +15,11 @@
 use std::sync::Arc;
 
 use arrow::{
-    array::{
-        ArrayBuilder, ArrayRef, Float32Builder, GenericListBuilder, GenericStringBuilder,
-        Int64Builder,
-    },
+    array::{ArrayRef, Float32Builder, GenericListBuilder, GenericStringBuilder, Int64Builder},
     datatypes::SchemaRef,
     error::ArrowError,
 };
+use exon_common::ExonArrayBuilder;
 use noodles::vcf::{
     variant::record::{AlternateBases, Filters, Ids},
     Header,
@@ -47,14 +45,16 @@ pub struct VCFArrayBuilder {
     header: Arc<Header>,
 
     projection: Vec<usize>,
+
+    n_rows: usize,
 }
 
 impl VCFArrayBuilder {
     /// Creates a new `VCFArrayBuilder` from a `Schema`.
     pub fn create(
         schema: SchemaRef,
         capacity: usize,
-        projection: Option<&[usize]>,
+        projection: Option<Vec<usize>>,
         header: Arc<Header>,
     ) -> Result<Self, ArrowError> {
         let info_field = schema.field_with_name("info")?;
@@ -66,6 +66,7 @@ impl VCFArrayBuilder {
         };
 
         Ok(Self {
+            n_rows: 0,
             chromosomes: GenericStringBuilder::<i32>::new(),
             positions: Int64Builder::new(),
             ids: GenericListBuilder::<i32, GenericStringBuilder<i32>>::new(GenericStringBuilder::<
@@ -89,16 +90,6 @@ impl VCFArrayBuilder {
         })
     }
 
-    /// Returns the number of records in the builder.
-    pub fn len(&self) -> usize {
-        self.chromosomes.len()
-    }
-
-    /// Returns whether the builder is empty.
-    pub fn is_empty(&self) -> bool {
-        self.len() == 0
-    }
-
     /// Appends a record to the builder.
     pub fn append<T>(&mut self, record: T) -> Result<(), ArrowError>
     where
@@ -169,11 +160,14 @@ impl VCFArrayBuilder {
             }
         }
 
+        self.n_rows += 1;
+
         Ok(())
     }
+}
 
-    /// Builds the `ArrayRef`.
-    pub fn finish(&mut self) -> Vec<ArrayRef> {
+impl ExonArrayBuilder for VCFArrayBuilder {
+    fn finish(&mut self) -> Vec<ArrayRef> {
         let mut arrays: Vec<ArrayRef> = vec![];
 
         for col_idx in self.projection.iter() {
@@ -193,4 +187,8 @@ impl VCFArrayBuilder {
 
         arrays
     }
+
+    fn len(&self) -> usize {
+        self.n_rows
+    }
 }