use std::sync::Arc;
use crate::{
datasources::{exon_listing_table_options::ExonListingConfig, ScanFunction},
ExonRuntimeEnvExt,
};
use datafusion::{
datasource::{function::TableFunctionImpl, TableProvider},
error::Result,
execution::context::SessionContext,
logical_expr::Expr,
};
use exon_bed::BEDSchemaBuilder;
use super::table_provider::{ListingBEDTable, ListingBEDTableOptions};
pub struct BEDScanFunction {
ctx: SessionContext,
}
impl std::fmt::Debug for BEDScanFunction {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("BEDScanFunction").finish()
}
}
impl BEDScanFunction {
pub fn new(ctx: SessionContext) -> Self {
Self { ctx }
}
}
impl TableFunctionImpl for BEDScanFunction {
fn call(&self, exprs: &[Expr]) -> Result<Arc<dyn TableProvider>> {
let listing_scan_function = ScanFunction::try_from(exprs)?;
futures::executor::block_on(async {
self.ctx
.runtime_env()
.exon_register_object_store_url(listing_scan_function.listing_table_url.as_ref())
.await
})?;
let schema = BEDSchemaBuilder::default().build();
let listing_table_options =
ListingBEDTableOptions::new(listing_scan_function.file_compression_type);
let listing_table_config = ExonListingConfig::new_with_options(
listing_scan_function.listing_table_url,
listing_table_options,
);
let listing_table = ListingBEDTable::new(listing_table_config, schema);
Ok(Arc::new(listing_table))
}
}