ARROW-11426: [Rust][DataFusion] EXTRACT support

rust/datafusion/src/logical_plan/expr.rs

@@ -1144,9 +1144,9 @@ fn create_name(e: &Expr, input_schema: &DFSchema) -> Result<String> {
             let expr = create_name(expr, input_schema)?;
             let list = list.iter().map(|expr| create_name(expr, input_schema));
             if *negated {
-                Ok(format!("{:?} NOT IN ({:?})", expr, list))
+                Ok(format!("{} NOT IN ({:?})", expr, list))
             } else {
-                Ok(format!("{:?} IN ({:?})", expr, list))
+                Ok(format!("{} IN ({:?})", expr, list))
             }
         }
         other => Err(DataFusionError::NotImplemented(format!(

rust/datafusion/src/logical_plan/mod.rs

@@ -29,7 +29,6 @@ mod extension;
 mod operators;
 mod plan;
 mod registry;
-
 pub use builder::LogicalPlanBuilder;
 pub use dfschema::{DFField, DFSchema, DFSchemaRef, ToDFSchema};
 pub use display::display_schema;

rust/datafusion/src/physical_plan/datetime_expressions.rs

@@ -16,27 +16,30 @@
 // under the License.
 
 //! DateTime expressions
-
 use std::sync::Arc;
 
+use super::ColumnarValue;
 use crate::{
     error::{DataFusionError, Result},
     scalar::{ScalarType, ScalarValue},
 };
-use arrow::temporal_conversions::timestamp_ns_to_datetime;
+use arrow::{
+    array::{Array, ArrayRef, GenericStringArray, PrimitiveArray, StringOffsetSizeTrait},
+    datatypes::{ArrowPrimitiveType, DataType, TimestampNanosecondType},
+};
 use arrow::{
     array::{
-        Array, GenericStringArray, PrimitiveArray, StringOffsetSizeTrait,
-        TimestampNanosecondArray,
+        Date32Array, Date64Array, TimestampMicrosecondArray, TimestampMillisecondArray,
+        TimestampNanosecondArray, TimestampSecondArray,
     },
-    datatypes::{ArrowPrimitiveType, DataType, TimestampNanosecondType},
+    compute::kernels::temporal,
+    datatypes::TimeUnit,
+    temporal_conversions::timestamp_ns_to_datetime,
 };
 use chrono::prelude::*;
 use chrono::Duration;
 use chrono::LocalResult;
 
-use super::ColumnarValue;
-
 #[inline]
 /// Accepts a string in RFC3339 / ISO8601 standard format and some
 /// variants and converts it to a nanosecond precision timestamp.
@@ -344,6 +347,98 @@ pub fn date_trunc(args: &[ColumnarValue]) -> Result<ColumnarValue> {
     })
 }
 
+macro_rules! extract_date_part {
+    ($ARRAY: expr, $FN:expr) => {
+        match $ARRAY.data_type() {
+            DataType::Date32 => {
+                let array = $ARRAY.as_any().downcast_ref::<Date32Array>().unwrap();
+                Ok($FN(array)?)
+            }
+            DataType::Date64 => {
+                let array = $ARRAY.as_any().downcast_ref::<Date64Array>().unwrap();
+                Ok($FN(array)?)
+            }
+            DataType::Timestamp(time_unit, None) => match time_unit {
+                TimeUnit::Second => {
+                    let array = $ARRAY
+                        .as_any()
+                        .downcast_ref::<TimestampSecondArray>()
+                        .unwrap();
+                    Ok($FN(array)?)
+                }
+                TimeUnit::Millisecond => {
+                    let array = $ARRAY
+                        .as_any()
+                        .downcast_ref::<TimestampMillisecondArray>()
+                        .unwrap();
+                    Ok($FN(array)?)
+                }
+                TimeUnit::Microsecond => {
+                    let array = $ARRAY
+                        .as_any()
+                        .downcast_ref::<TimestampMicrosecondArray>()
+                        .unwrap();
+                    Ok($FN(array)?)
+                }
+                TimeUnit::Nanosecond => {
+                    let array = $ARRAY
+                        .as_any()
+                        .downcast_ref::<TimestampNanosecondArray>()
+                        .unwrap();
+                    Ok($FN(array)?)
+                }
+            },
+            datatype => Err(DataFusionError::Internal(format!(
+                "Extract does not support datatype {:?}",
+                datatype
+            ))),
+        }
+    };
+}
+
+/// DATE_PART SQL function
+pub fn date_part(args: &[ColumnarValue]) -> Result<ColumnarValue> {
+    if args.len() != 2 {
+        return Err(DataFusionError::Execution(
+            "Expected two arguments in DATE_PART".to_string(),
+        ));
+    }
+    let (date_part, array) = (&args[0], &args[1]);
+
+    let date_part = if let ColumnarValue::Scalar(ScalarValue::Utf8(Some(v))) = date_part {
+        v
+    } else {
+        return Err(DataFusionError::Execution(
+            "First argument of `DATE_PART` must be non-null scalar Utf8".to_string(),
+        ));
+    };
+
+    let is_scalar = matches!(array, ColumnarValue::Scalar(_));
+
+    let array = match array {
+        ColumnarValue::Array(array) => array.clone(),
+        ColumnarValue::Scalar(scalar) => scalar.to_array(),
+    };
+
+    let arr = match date_part.to_lowercase().as_str() {
+        "hour" => extract_date_part!(array, temporal::hour),
+        "year" => extract_date_part!(array, temporal::year),
+        _ => Err(DataFusionError::Execution(format!(
+            "Date part '{}' not supported",
+            date_part
+        ))),
+    }?;
+
+    Ok(if is_scalar {
+        ColumnarValue::Scalar(ScalarValue::try_from_array(
+            &(Arc::new(arr) as ArrayRef),
+            0,
+        )?)
+    } else {
+        ColumnarValue::Array(Arc::new(arr))
+    })
+}
+
 #[cfg(test)]
 mod tests {
     use std::sync::Arc;

rust/datafusion/src/physical_plan/expressions/mod.rs

@@ -58,7 +58,6 @@ pub use negative::{negative, NegativeExpr};
 pub use not::{not, NotExpr};
 pub use nullif::{nullif_func, SUPPORTED_NULLIF_TYPES};
 pub use sum::{sum_return_type, Sum};
-
 /// returns the name of the state
 pub fn format_state_name(name: &str, state_name: &str) -> String {
     format!("{}[{}]", name, state_name)

rust/datafusion/src/physical_plan/functions.rs

@@ -71,6 +71,8 @@ pub enum Signature {
     Exact(Vec<DataType>),
     /// fixed number of arguments of arbitrary types
     Any(usize),
+    /// One of a list of signatures
+    OneOf(Vec<Signature>),
 }
 
 /// Scalar function
@@ -138,6 +140,8 @@ pub enum BuiltinScalarFunction {
     NullIf,
     /// Date truncate
     DateTrunc,
+    /// Date part
+    DatePart,
     /// MD5
     MD5,
     /// SHA224
@@ -192,6 +196,7 @@ impl FromStr for BuiltinScalarFunction {
             "upper" => BuiltinScalarFunction::Upper,
             "to_timestamp" => BuiltinScalarFunction::ToTimestamp,
             "date_trunc" => BuiltinScalarFunction::DateTrunc,
+            "date_part" => BuiltinScalarFunction::DatePart,
             "array" => BuiltinScalarFunction::Array,
             "nullif" => BuiltinScalarFunction::NullIf,
             "md5" => BuiltinScalarFunction::MD5,
@@ -294,6 +299,7 @@ pub fn return_type(
         BuiltinScalarFunction::DateTrunc => {
             Ok(DataType::Timestamp(TimeUnit::Nanosecond, None))
         }
+        BuiltinScalarFunction::DatePart => Ok(DataType::Int32),
         BuiltinScalarFunction::Array => Ok(DataType::FixedSizeList(
             Box::new(Field::new("item", arg_types[0].clone(), true)),
             arg_types.len() as i32,
@@ -463,6 +469,7 @@ pub fn create_physical_expr(
                 _ => unreachable!(),
             },
         },
+        BuiltinScalarFunction::DatePart => datetime_expressions::date_part,
     });
     // coerce
     let args = coerce(args, input_schema, &signature(fun))?;
@@ -507,6 +514,26 @@ fn signature(fun: &BuiltinScalarFunction) -> Signature {
             DataType::Utf8,
             DataType::Timestamp(TimeUnit::Nanosecond, None),
         ]),
+        BuiltinScalarFunction::DatePart => Signature::OneOf(vec![
+            Signature::Exact(vec![DataType::Utf8, DataType::Date32]),
+            Signature::Exact(vec![DataType::Utf8, DataType::Date64]),
+            Signature::Exact(vec![
+                DataType::Utf8,
+                DataType::Timestamp(TimeUnit::Second, None),
+            ]),
+            Signature::Exact(vec![
+                DataType::Utf8,
+                DataType::Timestamp(TimeUnit::Microsecond, None),
+            ]),
+            Signature::Exact(vec![
+                DataType::Utf8,
+                DataType::Timestamp(TimeUnit::Millisecond, None),
+            ]),
+            Signature::Exact(vec![
+                DataType::Utf8,
+                DataType::Timestamp(TimeUnit::Nanosecond, None),
+            ]),
+        ]),
         BuiltinScalarFunction::Array => {
             Signature::Variadic(array_expressions::SUPPORTED_ARRAY_TYPES.to_vec())
         }

rust/datafusion/src/physical_plan/type_coercion.rs

@@ -29,7 +29,7 @@
 //! i64. However, i64 -> i32 is never performed as there are i64
 //! values which can not be represented by i32 values.
 
-use std::sync::Arc;
+use std::{sync::Arc, vec};
 
 use arrow::datatypes::{DataType, Schema, TimeUnit};
 
@@ -68,6 +68,32 @@ pub fn data_types(
     current_types: &[DataType],
     signature: &Signature,
 ) -> Result<Vec<DataType>> {
+    let valid_types = get_valid_types(signature, current_types)?;
+
+    if valid_types
+        .iter()
+        .any(|data_type| data_type == current_types)
+    {
+        return Ok(current_types.to_vec());
+    }
+
+    for valid_types in valid_types {
+        if let Some(types) = maybe_data_types(&valid_types, &current_types) {
+            return Ok(types);
+        }
+    }
+
+    // none possible -> Error
+    Err(DataFusionError::Plan(format!(
+        "Coercion from {:?} to the signature {:?} failed.",
+        current_types, signature
+    )))
+}
+
+fn get_valid_types(
+    signature: &Signature,
+    current_types: &[DataType],
+) -> Result<Vec<Vec<DataType>>> {
     let valid_types = match signature {
         Signature::Variadic(valid_types) => valid_types
             .iter()
@@ -95,23 +121,16 @@ pub fn data_types(
             }
             vec![(0..*number).map(|i| current_types[i].clone()).collect()]
         }
-    };
-
-    if valid_types.contains(&current_types.to_owned()) {
-        return Ok(current_types.to_vec());
-    }
-
-    for valid_types in valid_types {
-        if let Some(types) = maybe_data_types(&valid_types, &current_types) {
-            return Ok(types);
+        Signature::OneOf(types) => {
+            let mut r = vec![];
+            for s in types {
+                r.extend(get_valid_types(s, current_types)?);
+            }
+            r
         }
-    }
+    };
 
-    // none possible -> Error
-    Err(DataFusionError::Plan(format!(
-        "Coercion from {:?} to the signature {:?} failed.",
-        current_types, signature
-    )))
+    Ok(valid_types)
 }
 
 /// Try to coerce current_types into valid_types.

rust/datafusion/src/sql/planner.rs

@@ -726,6 +726,13 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
             SQLExpr::Value(Value::Boolean(n)) => Ok(lit(*n)),
 
             SQLExpr::Value(Value::Null) => Ok(Expr::Literal(ScalarValue::Utf8(None))),
+            SQLExpr::Extract { field, expr } => Ok(Expr::ScalarFunction {
+                fun: functions::BuiltinScalarFunction::DatePart,
+                args: vec![
+                    Expr::Literal(ScalarValue::Utf8(Some(format!("{}", field)))),
+                    self.sql_expr_to_logical_expr(expr)?,
+                ],
+            }),
 
             SQLExpr::Value(Value::Interval {
                 value,

rust/datafusion/tests/sql.rs

@@ -1717,7 +1717,7 @@ fn make_timestamp_nano_table() -> Result<Arc<MemTable>> {
 }
 
 #[tokio::test]
-async fn to_timstamp() -> Result<()> {
+async fn to_timestamp() -> Result<()> {
     let mut ctx = ExecutionContext::new();
     ctx.register_table("ts_data", make_timestamp_nano_table()?);
 
@@ -2134,6 +2134,24 @@ async fn crypto_expressions() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn extract_date_part() -> Result<()> {
+    let mut ctx = ExecutionContext::new();
+    let sql = "SELECT
+        date_part('hour', CAST('2020-01-01' AS DATE)) AS hr1,
+        EXTRACT(HOUR FROM CAST('2020-01-01' AS DATE)) AS hr2,
+        EXTRACT(HOUR FROM to_timestamp('2020-09-08T12:00:00+00:00')) AS hr3,
+        date_part('YEAR', CAST('2000-01-01' AS DATE)) AS year1,
+        EXTRACT(year FROM to_timestamp('2020-09-08T12:00:00+00:00')) AS year2
+    ";
+
+    let actual = execute(&mut ctx, sql).await;
+
+    let expected = vec![vec!["0", "0", "12", "2000", "2020"]];
+    assert_eq!(expected, actual);
+    Ok(())
+}
+
 #[tokio::test]
 async fn in_list_array() -> Result<()> {
     let mut ctx = ExecutionContext::new();