use std::sync::Arc;
use datafusion::{
datasource::{
file_format::file_compression_type::FileCompressionType,
physical_plan::{FileMeta, FileOpenFuture, FileOpener},
},
error::DataFusionError,
};
use exon_mzml::{BatchReader, MzMLConfig};
use futures::{StreamExt, TryStreamExt};
use tokio_util::io::StreamReader;
pub struct MzMLOpener {
config: Arc<MzMLConfig>,
file_compression_type: FileCompressionType,
}
impl MzMLOpener {
pub fn new(config: Arc<MzMLConfig>, file_compression_type: FileCompressionType) -> Self {
Self {
config,
file_compression_type,
}
}
}
impl FileOpener for MzMLOpener {
fn open(&self, file_meta: FileMeta) -> datafusion::error::Result<FileOpenFuture> {
let mzml_config = Arc::clone(&self.config);
let file_compression_type = self.file_compression_type;
Ok(Box::pin(async move {
let get_result = mzml_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 mzml_batch_reader = BatchReader::new(stream_reader, mzml_config).into_stream();
Ok(mzml_batch_reader.boxed())
}))
}
}