use std::sync::Arc;
use datafusion::{
datasource::{
file_format::file_compression_type::FileCompressionType,
physical_plan::{FileMeta, FileOpenFuture, FileOpener},
},
error::DataFusionError,
};
use exon_fcs::{BatchReader, FCSConfig, FcsReader};
use futures::{StreamExt, TryStreamExt};
use tokio_util::io::StreamReader;
pub struct FCSOpener {
config: Arc<FCSConfig>,
file_compression_type: FileCompressionType,
}
impl FCSOpener {
pub fn new(config: Arc<FCSConfig>, file_compression_type: FileCompressionType) -> Self {
Self {
config,
file_compression_type,
}
}
}
impl FileOpener for FCSOpener {
fn open(&self, file_meta: FileMeta) -> datafusion::error::Result<FileOpenFuture> {
let config = Arc::clone(&self.config);
let file_compression_type = self.file_compression_type;
Ok(Box::pin(async move {
let get_result = config.object_store.get(file_meta.location()).await?;
let stream_reader = Box::pin(get_result.into_stream().map_err(DataFusionError::from));
let new_reader = match file_compression_type.convert_stream(stream_reader) {
Ok(reader) => reader,
Err(e) => return Err(e),
};
let stream_reader = StreamReader::new(new_reader);
let reader = FcsReader::new(stream_reader).await?;
let fasta_batch_reader = BatchReader::new(reader, config).into_stream();
Ok(fasta_batch_reader.boxed())
}))
}
}