use std::sync::Arc;
use arrow::datatypes::SchemaRef;
use arrow::error::Result;
use exon_common::DEFAULT_BATCH_SIZE;
use object_store::ObjectStore;
pub struct SAMConfig {
pub batch_size: usize,
pub file_schema: SchemaRef,
pub object_store: Arc<dyn ObjectStore>,
pub projection: Option<Vec<usize>>,
}
impl SAMConfig {
pub fn new(object_store: Arc<dyn ObjectStore>, file_schema: SchemaRef) -> Self {
Self {
batch_size: DEFAULT_BATCH_SIZE,
file_schema,
object_store,
projection: None,
}
}
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 {
self.projection = Some(projection);
self
}
pub fn projection(&self) -> Vec<usize> {
self.projection
.clone()
.unwrap_or_else(|| (0..self.file_schema.fields().len()).collect())
}
pub fn projected_schema(&self) -> Result<SchemaRef> {
let schema = self.file_schema.project(&self.projection())?;
Ok(Arc::new(schema))
}
}