datafusion_comet_spark_expr/datetime_funcs/
minute.rsuse crate::utils::array_with_timezone;
use arrow::{
compute::{date_part, DatePart},
record_batch::RecordBatch,
};
use arrow_schema::{DataType, Schema, TimeUnit::Microsecond};
use datafusion::logical_expr::ColumnarValue;
use datafusion_common::DataFusionError;
use datafusion_physical_expr::PhysicalExpr;
use std::hash::Hash;
use std::{
any::Any,
fmt::{Debug, Display, Formatter},
sync::Arc,
};
#[derive(Debug, Eq)]
pub struct MinuteExpr {
child: Arc<dyn PhysicalExpr>,
timezone: String,
}
impl Hash for MinuteExpr {
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
self.child.hash(state);
self.timezone.hash(state);
}
}
impl PartialEq for MinuteExpr {
fn eq(&self, other: &Self) -> bool {
self.child.eq(&other.child) && self.timezone.eq(&other.timezone)
}
}
impl MinuteExpr {
pub fn new(child: Arc<dyn PhysicalExpr>, timezone: String) -> Self {
MinuteExpr { child, timezone }
}
}
impl Display for MinuteExpr {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(
f,
"Minute [timezone:{}, child: {}]",
self.timezone, self.child
)
}
}
impl PhysicalExpr for MinuteExpr {
fn as_any(&self) -> &dyn Any {
self
}
fn data_type(&self, input_schema: &Schema) -> datafusion_common::Result<DataType> {
match self.child.data_type(input_schema).unwrap() {
DataType::Dictionary(key_type, _) => {
Ok(DataType::Dictionary(key_type, Box::new(DataType::Int32)))
}
_ => Ok(DataType::Int32),
}
}
fn nullable(&self, _: &Schema) -> datafusion_common::Result<bool> {
Ok(true)
}
fn evaluate(&self, batch: &RecordBatch) -> datafusion_common::Result<ColumnarValue> {
let arg = self.child.evaluate(batch)?;
match arg {
ColumnarValue::Array(array) => {
let array = array_with_timezone(
array,
self.timezone.clone(),
Some(&DataType::Timestamp(
Microsecond,
Some(self.timezone.clone().into()),
)),
)?;
let result = date_part(&array, DatePart::Minute)?;
Ok(ColumnarValue::Array(result))
}
_ => Err(DataFusionError::Execution(
"Minute(scalar) should be fold in Spark JVM side.".to_string(),
)),
}
}
fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
vec![&self.child]
}
fn with_new_children(
self: Arc<Self>,
children: Vec<Arc<dyn PhysicalExpr>>,
) -> Result<Arc<dyn PhysicalExpr>, DataFusionError> {
Ok(Arc::new(MinuteExpr::new(
Arc::clone(&children[0]),
self.timezone.clone(),
)))
}
}