use std::str::FromStr;
use datafusion::{
datasource::{
file_format::file_compression_type::FileCompressionType, listing::ListingTableUrl,
},
error::{DataFusionError, Result},
logical_expr::Expr,
scalar::ScalarValue,
};
#[derive(Debug)]
pub(crate) struct ScanFunction {
pub listing_table_url: ListingTableUrl,
pub file_compression_type: FileCompressionType,
}
impl TryFrom<&[Expr]> for ScanFunction {
type Error = DataFusionError;
fn try_from(exprs: &[Expr]) -> Result<Self> {
let Some(Expr::Literal(ScalarValue::Utf8(Some(ref path)))) = exprs.first() else {
return Err(DataFusionError::Internal(
"this function requires the path to be specified as the first argument".into(),
));
};
let listing_table_url = ListingTableUrl::parse(path)?;
let passed_compression_type = exprs.get(1).and_then(|e| match e {
Expr::Literal(ScalarValue::Utf8(Some(ref compression_type))) => {
FileCompressionType::from_str(compression_type).ok()
}
_ => None,
});
let inferred_compression = listing_table_url
.prefix()
.extension()
.and_then(|ext| FileCompressionType::from_str(ext).ok());
let file_compression_type = passed_compression_type.or(inferred_compression);
Ok(Self {
listing_table_url,
file_compression_type: file_compression_type
.unwrap_or(FileCompressionType::UNCOMPRESSED),
})
}
}