style: fix clippy

wheretrue · May 28, 2024 · af872d3 · af872d3
1 parent 4be456a
commit af872d3
Show file tree

Hide file tree

Showing 6 changed files with 27 additions and 20 deletions.
diff --git a/src/bam_reader.rs b/src/bam_reader.rs
@@ -15,10 +15,9 @@
 use arrow::ffi_stream::ArrowArrayStreamReader;
 use arrow::ffi_stream::FFI_ArrowArrayStream;
 use arrow::pyarrow::IntoPyArrow;
-use datafusion::prelude::SessionContext;
 use exon::ffi::DataFrameRecordBatchStream;
 use exon::new_exon_config;
-use exon::ExonSessionExt;
+use exon::ExonSession;
 use pyo3::prelude::*;
 
 use tokio::runtime::Runtime;
@@ -72,7 +71,7 @@ impl BamIndexedReader {
             config = config.with_batch_size(batch_size);
         }
 
-        let ctx = SessionContext::with_config_exon(config);
+        let ctx = ExonSession::with_config_exon(config);
 
         let df = self._runtime.block_on(async {
             ctx.sql(&format!(

diff --git a/src/bcf_reader.rs b/src/bcf_reader.rs
@@ -14,15 +14,14 @@
 
 use arrow::ffi_stream::{ArrowArrayStreamReader, FFI_ArrowArrayStream};
 use arrow::pyarrow::IntoPyArrow;
-use datafusion::prelude::{SessionConfig, SessionContext};
+use datafusion::prelude::SessionConfig;
 use exon::datasources::bcf::table_provider::ListingBCFTableOptions;
 use exon::ffi::DataFrameRecordBatchStream;
+use exon::ExonSession;
 use noodles::core::Region;
 use pyo3::prelude::*;
 use tokio::runtime::Runtime;
 
-use exon::ExonSessionExt;
-
 use std::io;
 use std::str::FromStr;
 use std::sync::Arc;
@@ -69,7 +68,7 @@ impl BCFIndexedReader {
             config = config.with_batch_size(batch_size);
         }
 
-        let ctx = SessionContext::new_with_config(config);
+        let ctx = ExonSession::with_config_exon(config);
 
         let region = Region::from_str(region).map_err(|e| {
             io::Error::new(io::ErrorKind::Other, format!("Error parsing region: {e}"))

diff --git a/src/datasources/fasta.rs b/src/datasources/fasta.rs
@@ -13,7 +13,7 @@
 // limitations under the License.
 
 use crate::{error::BioBearResult, file_compression_type::FileCompressionType};
-use exon::datasources::fasta::table_provider::ListingFASTATableOptions;
+use exon::datasources::fasta::{table_provider::ListingFASTATableOptions, SequenceDataType};
 use pyo3::{pyclass, pymethods};
 
 const DEFAULT_FASTA_FILE_EXTENSION: &str = "fasta";
@@ -27,6 +27,17 @@ pub enum FastaSequenceDataType {
     IntegerEncodeProtein,
 }
 
+impl From<FastaSequenceDataType> for SequenceDataType {
+    fn from(data_type: FastaSequenceDataType) -> Self {
+        match data_type {
+            FastaSequenceDataType::Utf8 => SequenceDataType::Utf8,
+            FastaSequenceDataType::LargeUtf8 => SequenceDataType::LargeUtf8,
+            FastaSequenceDataType::IntegerEncodeDNA => SequenceDataType::IntegerEncodeDNA,
+            FastaSequenceDataType::IntegerEncodeProtein => SequenceDataType::IntegerEncodeProtein,
+        }
+    }
+}
+
 #[pyclass]
 #[derive(Debug, Clone)]
 /// Options for reading FASTA files.

diff --git a/src/exon_reader.rs b/src/exon_reader.rs
@@ -18,10 +18,9 @@ use std::sync::Arc;
 use arrow::ffi_stream::{ArrowArrayStreamReader, FFI_ArrowArrayStream};
 use arrow::pyarrow::IntoPyArrow;
 use datafusion::datasource::file_format::file_compression_type::FileCompressionType;
-use datafusion::prelude::SessionContext;
 use exon::datasources::ExonFileType;
 use exon::ffi::DataFrameRecordBatchStream;
-use exon::{new_exon_config, ExonRuntimeEnvExt, ExonSessionExt};
+use exon::{new_exon_config, ExonRuntimeEnvExt, ExonSession};
 use pyo3::prelude::*;
 use tokio::runtime::Runtime;
 
@@ -49,10 +48,11 @@ impl ExonReader {
             config = config.with_batch_size(batch_size);
         }
 
-        let ctx = SessionContext::with_config_exon(config);
+        let ctx = ExonSession::with_config_exon(config);
 
         let df = rt.block_on(async {
-            ctx.runtime_env()
+            ctx.session
+                .runtime_env()
                 .exon_register_object_store_uri(path)
                 .await
                 .map_err(BioBearError::from)?;

diff --git a/src/session_context.rs b/src/session_context.rs
@@ -12,9 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use datafusion::prelude::SessionContext;
 use exon::datasources::bigwig;
-use exon::{ExonRuntimeEnvExt, ExonSessionExt};
+use exon::{ExonRuntimeEnvExt, ExonSession};
 
 use pyo3::prelude::*;
 
@@ -30,13 +29,13 @@ use crate::runtime::wait_for_future;
 
 #[pyclass]
 pub struct BioBearSessionContext {
-    ctx: SessionContext,
+    ctx: ExonSession,
 }
 
 impl Default for BioBearSessionContext {
     fn default() -> Self {
         Self {
-            ctx: SessionContext::new_exon(),
+            ctx: ExonSession::new_exon(),
         }
     }
 }
@@ -291,7 +290,7 @@ impl BioBearSessionContext {
 
     /// Register an object store with the given URL.
     fn register_object_store_from_url(&mut self, url: &str, py: Python) -> PyResult<()> {
-        let runtime = self.ctx.runtime_env();
+        let runtime = self.ctx.session.runtime_env();
         let registration = runtime.exon_register_object_store_uri(url);
         wait_for_future(py, registration).map_err(error::BioBearError::from)?;
 

diff --git a/src/vcf_reader.rs b/src/vcf_reader.rs
@@ -15,14 +15,13 @@
 use arrow::ffi_stream::{ArrowArrayStreamReader, FFI_ArrowArrayStream};
 use arrow::pyarrow::IntoPyArrow;
 use datafusion::datasource::file_format::file_compression_type::FileCompressionType;
-use datafusion::prelude::SessionContext;
 use exon::datasources::vcf::ListingVCFTableOptions;
 use exon::ffi::DataFrameRecordBatchStream;
 use noodles::core::Region;
 use pyo3::prelude::*;
 use tokio::runtime::Runtime;
 
-use exon::{new_exon_config, ExonSessionExt};
+use exon::{new_exon_config, ExonSession};
 
 use std::io;
 use std::str::FromStr;
@@ -70,7 +69,7 @@ impl VCFIndexedReader {
             config = config.with_batch_size(batch_size);
         }
 
-        let ctx = SessionContext::with_config_exon(config);
+        let ctx = ExonSession::with_config_exon(config);
 
         let region = Region::from_str(region).map_err(|e| {
             io::Error::new(