Spark date part (#19823)

## Which issue does this PR close?

<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax. For example
`Closes #123` indicates that this PR will close issue #123.
-->

- Closes #19822 .
- Part of #15914 

## Rationale for this change

The current date_part function in datafusion have a few differences with
the spark implementation:

- day of week parts are 1 indexed in spark but 0 indexed in datafusion
- spark supports a few more aliases for certain parts

Full list of spark supported aliases:
https://github.com/apache/spark/blob/a03bedb6c1281c5263a42bfd20608d2ee005ab05/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datetimeExpressions.scala#L3356-L3371

## What changes are included in this PR?

New date_part function in spark crate.

## Are these changes tested?

Yes with SLT

## Are there any user-facing changes?
yes

Author: cht42

datafusion/functions/src/datetime/date_part.rs

@@ -152,6 +152,7 @@ impl ScalarUDFImpl for DatePartFunc {
 
     fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
         let [field, _] = take_function_args(self.name(), args.scalar_arguments)?;
+        let nullable = args.arg_fields[1].is_nullable();
 
         field
             .and_then(|sv| {
@@ -160,9 +161,9 @@ impl ScalarUDFImpl for DatePartFunc {
                     .filter(|s| !s.is_empty())
                     .map(|part| {
                         if is_epoch(part) {
-                            Field::new(self.name(), DataType::Float64, true)
+                            Field::new(self.name(), DataType::Float64, nullable)
                         } else {
-                            Field::new(self.name(), DataType::Int32, true)
+                            Field::new(self.name(), DataType::Int32, nullable)
                         }
                     })
             })

datafusion/spark/src/function/datetime/date_part.rs

@@ -0,0 +1,139 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use arrow::datatypes::{DataType, Field, FieldRef};
+use datafusion_common::types::logical_date;
+use datafusion_common::{
+    Result, ScalarValue, internal_err, types::logical_string, utils::take_function_args,
+};
+use datafusion_expr::expr::ScalarFunction;
+use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyContext};
+use datafusion_expr::{
+    Coercion, ColumnarValue, Expr, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl,
+    Signature, TypeSignature, TypeSignatureClass, Volatility,
+};
+use std::{any::Any, sync::Arc};
+
+/// Wrapper around datafusion date_part function to handle
+/// Spark behavior returning day of the week 1-indexed instead of 0-indexed and different part aliases.
+/// <https://spark.apache.org/docs/latest/api/sql/index.html#date_part>
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SparkDatePart {
+    signature: Signature,
+    aliases: Vec<String>,
+}
+
+impl Default for SparkDatePart {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl SparkDatePart {
+    pub fn new() -> Self {
+        Self {
+            signature: Signature::one_of(
+                vec![
+                    TypeSignature::Coercible(vec![
+                        Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
+                        Coercion::new_exact(TypeSignatureClass::Timestamp),
+                    ]),
+                    TypeSignature::Coercible(vec![
+                        Coercion::new_exact(TypeSignatureClass::Native(logical_string())),
+                        Coercion::new_exact(TypeSignatureClass::Native(logical_date())),
+                    ]),
+                ],
+                Volatility::Immutable,
+            ),
+            aliases: vec![String::from("datepart")],
+        }
+    }
+}
+
+impl ScalarUDFImpl for SparkDatePart {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn name(&self) -> &str {
+        "date_part"
+    }
+
+    fn aliases(&self) -> &[String] {
+        &self.aliases
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+        internal_err!("Use return_field_from_args in this case instead.")
+    }
+
+    fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
+        let nullable = args.arg_fields.iter().any(|f| f.is_nullable());
+
+        Ok(Arc::new(Field::new(self.name(), DataType::Int32, nullable)))
+    }
+
+    fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        internal_err!("spark date_part should have been simplified to standard date_part")
+    }
+
+    fn simplify(
+        &self,
+        args: Vec<Expr>,
+        _info: &SimplifyContext,
+    ) -> Result<ExprSimplifyResult> {
+        let [part_expr, date_expr] = take_function_args(self.name(), args)?;
+
+        let part = match part_expr.as_literal() {
+            Some(ScalarValue::Utf8(Some(v)))
+            | Some(ScalarValue::Utf8View(Some(v)))
+            | Some(ScalarValue::LargeUtf8(Some(v))) => v.to_lowercase(),
+            _ => {
+                return internal_err!(
+                    "First argument of `DATE_PART` must be non-null scalar Utf8"
+                );
+            }
+        };
+
+        // Map Spark-specific date part aliases to datafusion ones
+        let part = match part.as_str() {
+            "yearofweek" | "year_iso" => "isoyear",
+            "dayofweek" => "dow",
+            "dayofweek_iso" | "dow_iso" => "isodow",
+            other => other,
+        };
+
+        let part_expr = Expr::Literal(ScalarValue::new_utf8(part), None);
+
+        let date_part_expr = Expr::ScalarFunction(ScalarFunction::new_udf(
+            datafusion_functions::datetime::date_part(),
+            vec![part_expr, date_expr],
+        ));
+
+        match part {
+            // Add 1 for day-of-week parts to convert 0-indexed to 1-indexed
+            "dow" | "isodow" => Ok(ExprSimplifyResult::Simplified(
+                date_part_expr + Expr::Literal(ScalarValue::Int32(Some(1)), None),
+            )),
+            _ => Ok(ExprSimplifyResult::Simplified(date_part_expr)),
+        }
+    }
+}

datafusion/spark/src/function/datetime/mod.rs

@@ -16,6 +16,7 @@
 // under the License.
 
 pub mod date_add;
+pub mod date_part;
 pub mod date_sub;
 pub mod extract;
 pub mod last_day;
@@ -36,6 +37,7 @@ make_udf_function!(last_day::SparkLastDay, last_day);
 make_udf_function!(make_dt_interval::SparkMakeDtInterval, make_dt_interval);
 make_udf_function!(make_interval::SparkMakeInterval, make_interval);
 make_udf_function!(next_day::SparkNextDay, next_day);
+make_udf_function!(date_part::SparkDatePart, date_part);
 
 pub mod expr_fn {
     use datafusion_functions::export_functions;
@@ -83,6 +85,11 @@ pub mod expr_fn {
         "Returns the first date which is later than start_date and named as indicated. The function returns NULL if at least one of the input parameters is NULL.",
         arg1 arg2
     ));
+    export_functions!((
+        date_part,
+        "Extracts a part of the date or time from a date, time, or timestamp expression.",
+        arg1 arg2
+    ));
 }
 
 pub fn functions() -> Vec<Arc<ScalarUDF>> {
@@ -96,5 +103,6 @@ pub fn functions() -> Vec<Arc<ScalarUDF>> {
         make_dt_interval(),
         make_interval(),
         next_day(),
+        date_part(),
     ]
 }

datafusion/spark/src/lib.rs

@@ -92,6 +92,24 @@
 //! let expr = sha2(col("my_data"), lit(256));
 //! ```
 //!
+//! # Example: using the Spark expression planner
+//!
+//! The [`planner::SparkFunctionPlanner`] provides Spark-compatible expression
+//! planning, such as mapping SQL `EXTRACT` expressions to Spark's `date_part`
+//! function. To use it, register it with your session context:
+//!
+//! ```ignore
+//! use std::sync::Arc;
+//! use datafusion::prelude::SessionContext;
+//! use datafusion_spark::planner::SparkFunctionPlanner;
+//!
+//! let mut ctx = SessionContext::new();
+//! // Register the Spark expression planner
+//! ctx.register_expr_planner(Arc::new(SparkFunctionPlanner))?;
+//! // Now EXTRACT expressions will use Spark semantics
+//! let df = ctx.sql("SELECT EXTRACT(YEAR FROM timestamp_col) FROM my_table").await?;
+//! ```
+//!
 //![`Expr`]: datafusion_expr::Expr
 
 pub mod function;

datafusion/spark/src/planner.rs

@@ -23,6 +23,15 @@ use datafusion_expr::planner::{ExprPlanner, PlannerResult};
 pub struct SparkFunctionPlanner;
 
 impl ExprPlanner for SparkFunctionPlanner {
+    fn plan_extract(
+        &self,
+        args: Vec<Expr>,
+    ) -> datafusion_common::Result<PlannerResult<Vec<Expr>>> {
+        Ok(PlannerResult::Planned(Expr::ScalarFunction(
+            ScalarFunction::new_udf(crate::function::datetime::date_part(), args),
+        )))
+    }
+
     fn plan_substring(
         &self,
         args: Vec<Expr>,