use crate::{
file::FileSource, file_scan_config::FileScanConfig, file_stream::FileOpener,
schema_adapter::SchemaAdapterFactory,
};
use std::sync::Arc;
use arrow::datatypes::{Schema, SchemaRef};
use datafusion_common::{Result, Statistics};
use datafusion_physical_expr::{expressions::Column, PhysicalExpr};
use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
use object_store::ObjectStore;
#[derive(Clone, Default)]
pub(crate) struct MockSource {
metrics: ExecutionPlanMetricsSet,
projected_statistics: Option<Statistics>,
schema_adapter_factory: Option<Arc<dyn SchemaAdapterFactory>>,
}
impl FileSource for MockSource {
fn create_file_opener(
&self,
_object_store: Arc<dyn ObjectStore>,
_base_config: &FileScanConfig,
_partition: usize,
) -> Arc<dyn FileOpener> {
unimplemented!()
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
fn with_batch_size(&self, _batch_size: usize) -> Arc<dyn FileSource> {
Arc::new(Self { ..self.clone() })
}
fn with_schema(&self, _schema: SchemaRef) -> Arc<dyn FileSource> {
Arc::new(Self { ..self.clone() })
}
fn with_projection(&self, _config: &FileScanConfig) -> Arc<dyn FileSource> {
Arc::new(Self { ..self.clone() })
}
fn with_statistics(&self, statistics: Statistics) -> Arc<dyn FileSource> {
let mut source = self.clone();
source.projected_statistics = Some(statistics);
Arc::new(source)
}
fn metrics(&self) -> &ExecutionPlanMetricsSet {
&self.metrics
}
fn statistics(&self) -> Result<Statistics> {
Ok(self
.projected_statistics
.as_ref()
.expect("projected_statistics must be set")
.clone())
}
fn file_type(&self) -> &str {
"mock"
}
fn with_schema_adapter_factory(
&self,
schema_adapter_factory: Arc<dyn SchemaAdapterFactory>,
) -> Result<Arc<dyn FileSource>> {
Ok(Arc::new(Self {
schema_adapter_factory: Some(schema_adapter_factory),
..self.clone()
}))
}
fn schema_adapter_factory(&self) -> Option<Arc<dyn SchemaAdapterFactory>> {
self.schema_adapter_factory.clone()
}
}
pub(crate) fn col(name: &str, schema: &Schema) -> Result<Arc<dyn PhysicalExpr>> {
Ok(Arc::new(Column::new_with_schema(name, schema)?))
}