Dandandan commented on a change in pull request #9359:
URL: https://github.com/apache/arrow/pull/9359#discussion_r567414429
##########
File path: rust/datafusion/src/physical_plan/expressions/extract.rs
##########
@@ -0,0 +1,115 @@
+use crate::error::Result;
+use core::fmt;
+use std::{any::Any, sync::Arc};
+
+use arrow::{
+ array::{
+ Date32Array, Date64Array, TimestampMicrosecondArray,
TimestampMillisecondArray,
+ TimestampNanosecondArray, TimestampSecondArray,
+ },
+ compute::hour,
+ datatypes::{DataType, Schema, TimeUnit},
+ record_batch::RecordBatch,
+};
+
+use crate::{
+ error::DataFusionError,
+ logical_plan::DatePart,
+ physical_plan::{ColumnarValue, PhysicalExpr},
+};
+
+impl fmt::Display for Extract {
+ fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+ write!(f, "EXTRACT({} AS {:?})", self.date_part, self.expr)
+ }
+}
+
+/// Extract
+#[derive(Debug)]
+pub struct Extract {
+ date_part: DatePart,
+ expr: Arc<dyn PhysicalExpr>,
+}
+
+impl Extract {
+ /// Create new Extract expression
+ pub fn new(date_part: DatePart, expr: Arc<dyn PhysicalExpr>) -> Self {
+ Self { date_part, expr }
+ }
+}
+
+impl PhysicalExpr for Extract {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
+ Ok(DataType::Int32)
+ }
+
+ fn nullable(&self, input_schema: &Schema) -> Result<bool> {
+ self.expr.nullable(input_schema)
+ }
+
+ fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
+ let value = self.expr.evaluate(batch)?;
+ let data_type = value.data_type();
+ let array = match value {
+ ColumnarValue::Array(array) => array,
+ ColumnarValue::Scalar(scalar) => scalar.to_array(),
+ };
+
+ match data_type {
+ DataType::Date32 => {
+ let array =
array.as_any().downcast_ref::<Date32Array>().unwrap();
+ Ok(ColumnarValue::Array(Arc::new(hour(array)?)))
Review comment:
There is no other datepart for now in the enum, so I guess that might
generate some clippy warnings. But we could match directly on the datepart. The
warning should be generated elsewhere already (when building the logical plan),
agree makes sense to add a test for that 👍.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
