use std::sync::Arc;
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use object_store::ObjectStore;
pub struct FASTAConfig {
pub batch_size: usize,
pub file_schema: SchemaRef,
pub object_store: Arc<dyn ObjectStore>,
pub projection: Option<Vec<usize>>,
pub fasta_sequence_buffer_capacity: usize,
}
impl FASTAConfig {
pub fn new(object_store: Arc<dyn ObjectStore>, file_schema: SchemaRef) -> Self {
Self {
object_store,
file_schema,
batch_size: 8192,
projection: None,
fasta_sequence_buffer_capacity: 384, }
}
pub fn with_batch_size(mut self, batch_size: usize) -> Self {
self.batch_size = batch_size;
self
}
pub fn with_projection(mut self, projection: Vec<usize>) -> Self {
let file_projection = projection
.iter()
.filter(|f| **f < self.file_schema.fields().len())
.cloned()
.collect::<Vec<_>>();
self.projection = Some(file_projection);
self
}
pub fn with_fasta_sequence_buffer_capacity(
mut self,
fasta_sequence_buffer_capacity: usize,
) -> Self {
self.fasta_sequence_buffer_capacity = fasta_sequence_buffer_capacity;
self
}
}
impl Default for FASTAConfig {
fn default() -> Self {
let file_schema = FASTASchemaBuilder::default().build();
Self::new(
Arc::new(object_store::local::LocalFileSystem::new()),
file_schema,
)
}
}
pub struct FASTASchemaBuilder {
fields: Vec<Field>,
}
impl FASTASchemaBuilder {
pub fn extend(mut self, fields: Vec<Field>) -> Self {
self.fields.extend(fields);
self
}
pub fn build(self) -> SchemaRef {
Arc::new(Schema::new(self.fields))
}
}
impl Default for FASTASchemaBuilder {
fn default() -> Self {
let fields = vec![
Field::new("id", DataType::Utf8, false),
Field::new("description", DataType::Utf8, true),
Field::new("sequence", DataType::Utf8, false),
];
Self { fields }
}
}