use std::any::Any;
use arrow::datatypes::DataType;
use arrow::datatypes::DataType::Date32;
use chrono::{Datelike, NaiveDate};
use datafusion_common::{internal_err, Result, ScalarValue};
use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyInfo};
use datafusion_expr::{ColumnarValue, Expr, ScalarUDFImpl, Signature, Volatility};
#[derive(Debug)]
pub(super) struct CurrentDateFunc {
signature: Signature,
aliases: Vec<String>,
}
impl CurrentDateFunc {
pub fn new() -> Self {
Self {
signature: Signature::uniform(0, vec![], Volatility::Stable),
aliases: vec![String::from("today")],
}
}
}
impl ScalarUDFImpl for CurrentDateFunc {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
"current_date"
}
fn signature(&self) -> &Signature {
&self.signature
}
fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
Ok(Date32)
}
fn invoke(&self, _args: &[ColumnarValue]) -> Result<ColumnarValue> {
internal_err!(
"invoke should not be called on a simplified current_date() function"
)
}
fn aliases(&self) -> &[String] {
&self.aliases
}
fn simplify(
&self,
_args: Vec<Expr>,
info: &dyn SimplifyInfo,
) -> Result<ExprSimplifyResult> {
let now_ts = info.execution_props().query_execution_start_time;
let days = Some(
now_ts.num_days_from_ce()
- NaiveDate::from_ymd_opt(1970, 1, 1)
.unwrap()
.num_days_from_ce(),
);
Ok(ExprSimplifyResult::Simplified(Expr::Literal(
ScalarValue::Date32(days),
)))
}
}