fix(spark): date_part SECOND returns Decimal(8,6)

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

@@ -2,18 +2,50 @@
 // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
 // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
 
-use arrow::datatypes::{DataType, Field, FieldRef};
+use arrow::array::{Array, ArrayRef, AsArray, Decimal128Array};
+use arrow::compute::{DatePart, cast, date_part};
+use arrow::datatypes::{DataType, Field, FieldRef, Int32Type, TimeUnit};
 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,
+    Coercion, ColumnarValue, Expr, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF,
+    ScalarUDFImpl, Signature, TypeSignature, TypeSignatureClass, Volatility,
 };
-use std::sync::Arc;
+use datafusion_functions::utils::make_scalar_function;
+use std::sync::{Arc, OnceLock};
+
+/// Precision and scale of the `DECIMAL(8, 6)` Spark returns for the seconds
+/// part, which carries the fraction down to microseconds.
+const SECONDS_PRECISION: u8 = 8;
+const SECONDS_SCALE: i8 = 6;
+
+/// Returns true when `part` is a Spark alias for the seconds field.
+///
+/// Spark returns `DECIMAL(8, 6)` with the fractional seconds for these instead
+/// of an integer.
+pub fn is_seconds_part(part: &str) -> bool {
+    matches!(
+        part.to_lowercase().as_str(),
+        "second" | "s" | "sec" | "seconds" | "secs"
+    )
+}
+
+/// UDF that implements Spark `date_part('SECOND', …)` / `EXTRACT(SECOND FROM …)`.
+///
+/// Always returns `DECIMAL(8, 6)`. Rewritten onto call sites by
+/// [`crate::planner::SparkFunctionPlanner`] so the return type is known before
+/// analysis (user-written casts on the part literal hide it from
+/// `return_field_from_args`).
+pub fn date_part_seconds() -> Arc<ScalarUDF> {
+    static INSTANCE: OnceLock<Arc<ScalarUDF>> = OnceLock::new();
+    INSTANCE
+        .get_or_init(|| Arc::new(ScalarUDF::from(SparkDatePartSeconds::new())))
+        .clone()
+}
 
 /// 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.
@@ -69,9 +101,24 @@ impl ScalarUDFImpl for SparkDatePart {
     }
 
     // Nullable, like the integer parts of `date_part`, which this simplifies to,
-    // because they can be NULL for a non-null input
-    fn return_field_from_args(&self, _args: ReturnFieldArgs) -> Result<FieldRef> {
-        Ok(Arc::new(Field::new(self.name(), DataType::Int32, true)))
+    // because they can be NULL for a non-null input. Seconds calls are rewritten
+    // to [`date_part_seconds`] (DECIMAL(8, 6)) by the Spark expression planner.
+    fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
+        let is_seconds = matches!(
+            args.scalar_arguments.first(),
+            Some(Some(
+                ScalarValue::Utf8(Some(part))
+                    | ScalarValue::Utf8View(Some(part))
+                    | ScalarValue::LargeUtf8(Some(part))
+            )) if is_seconds_part(part)
+        );
+        let data_type = if is_seconds {
+            DataType::Decimal128(SECONDS_PRECISION, SECONDS_SCALE)
+        } else {
+            DataType::Int32
+        };
+
+        Ok(Arc::new(Field::new(self.name(), data_type, true)))
     }
 
     fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
@@ -96,6 +143,17 @@ impl ScalarUDFImpl for SparkDatePart {
             }
         };
 
+        // Prefer the planner rewrite; keep this path for DataFrame callers that
+        // pass a bare string literal (return type already Decimal via
+        // `return_field_from_args`).
+        if is_seconds_part(&part) {
+            let seconds_expr = Expr::ScalarFunction(ScalarFunction::new_udf(
+                date_part_seconds(),
+                vec![date_expr],
+            ));
+            return Ok(ExprSimplifyResult::Simplified(seconds_expr));
+        }
+
         // Map Spark-specific date part aliases to DataFusion ones.
         let part = match part.as_str() {
             "yearofweek" | "year_iso" => "isoyear",
@@ -114,3 +172,87 @@ impl ScalarUDFImpl for SparkDatePart {
         Ok(ExprSimplifyResult::Simplified(date_part_expr))
     }
 }
+
+/// `date_part('SECOND', ts)` for Spark: seconds including the microsecond
+/// fraction as `DECIMAL(8, 6)`.
+#[derive(Debug, PartialEq, Eq, Hash)]
+struct SparkDatePartSeconds {
+    signature: Signature,
+}
+
+impl SparkDatePartSeconds {
+    fn new() -> Self {
+        Self {
+            signature: Signature::one_of(
+                vec![
+                    TypeSignature::Coercible(vec![Coercion::new_exact(
+                        TypeSignatureClass::Timestamp,
+                    )]),
+                    TypeSignature::Coercible(vec![Coercion::new_exact(
+                        TypeSignatureClass::Native(logical_date()),
+                    )]),
+                ],
+                Volatility::Immutable,
+            ),
+        }
+    }
+}
+
+impl ScalarUDFImpl for SparkDatePartSeconds {
+    fn name(&self) -> &str {
+        // Keep the display name as date_part so EXPLAIN/DESCRIBE stay familiar.
+        "date_part"
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+        Ok(DataType::Decimal128(SECONDS_PRECISION, SECONDS_SCALE))
+    }
+
+    fn return_field_from_args(&self, _args: ReturnFieldArgs) -> Result<FieldRef> {
+        Ok(Arc::new(Field::new(
+            self.name(),
+            DataType::Decimal128(SECONDS_PRECISION, SECONDS_SCALE),
+            true,
+        )))
+    }
+
+    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        make_scalar_function(spark_date_part_seconds, vec![])(&args.args)
+    }
+}
+
+/// Seconds of a timestamp or date including the fraction, as `DECIMAL(8, 6)`.
+fn spark_date_part_seconds(args: &[ArrayRef]) -> Result<ArrayRef> {
+    let [date_arg] = take_function_args("date_part", args)?;
+
+    // A date has no time of day; Spark casts it to a timestamp at midnight.
+    let date_arg = match date_arg.data_type() {
+        DataType::Date32 | DataType::Date64 => {
+            cast(date_arg, &DataType::Timestamp(TimeUnit::Microsecond, None))?
+        }
+        _ => Arc::clone(date_arg),
+    };
+
+    let secs = date_part(date_arg.as_ref(), DatePart::Second)?;
+    let micros = date_part(date_arg.as_ref(), DatePart::Microsecond)?;
+    let secs = secs.as_primitive::<Int32Type>();
+    let micros = micros.as_primitive::<Int32Type>();
+
+    // `DatePart::Microsecond` is the fraction of the current second, so the
+    // unscaled decimal value is the whole seconds times 10^6 plus the fraction.
+    let result: Decimal128Array = secs
+        .iter()
+        .zip(micros.iter())
+        .map(|(s, us)| match (s, us) {
+            (Some(s), Some(us)) => Some(s as i128 * 1_000_000 + us as i128),
+            _ => None,
+        })
+        .collect();
+    let result = result.with_precision_and_scale(SECONDS_PRECISION, SECONDS_SCALE)?;
+
+    Ok(Arc::new(result))
+}

datafusion/spark/src/planner.rs

@@ -2,9 +2,11 @@
 // SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
 // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
 
+use crate::function::datetime::date_part::{date_part_seconds, is_seconds_part};
+use datafusion_common::ScalarValue;
 use datafusion_expr::Expr;
 use datafusion_expr::expr::ScalarFunction;
-use datafusion_expr::planner::{ExprPlanner, PlannerResult};
+use datafusion_expr::planner::{ExprPlanner, PlannerResult, RawScalarExpr};
 
 #[derive(Default, Debug)]
 pub struct SparkFunctionPlanner;
@@ -14,6 +16,20 @@ impl ExprPlanner for SparkFunctionPlanner {
         &self,
         args: Vec<Expr>,
     ) -> datafusion_common::Result<PlannerResult<Vec<Expr>>> {
+        // EXTRACT(SECOND FROM ts) must return DECIMAL(8, 6) like date_part('SECOND', ts).
+        if let Some(part) = args.first().and_then(expr_string_literal)
+            && is_seconds_part(&part)
+        {
+            let date_expr = args.into_iter().nth(1).ok_or_else(|| {
+                datafusion_common::DataFusionError::Internal(
+                    "EXTRACT expects a source expression".to_string(),
+                )
+            })?;
+            return Ok(PlannerResult::Planned(Expr::ScalarFunction(
+                ScalarFunction::new_udf(date_part_seconds(), vec![date_expr]),
+            )));
+        }
+
         Ok(PlannerResult::Planned(Expr::ScalarFunction(
             ScalarFunction::new_udf(crate::function::datetime::date_part(), args),
         )))
@@ -27,4 +43,44 @@ impl ExprPlanner for SparkFunctionPlanner {
             ScalarFunction::new_udf(crate::function::string::substring(), args),
         )))
     }
+
+    fn plan_scalar(
+        &self,
+        expr: RawScalarExpr,
+        _schema: &datafusion_common::DFSchema,
+    ) -> datafusion_common::Result<PlannerResult<RawScalarExpr>> {
+        // date_part('SECOND', ts) / datepart(...) → DECIMAL(8, 6) seconds UDF.
+        // Done here so the return type is fixed before analysis; a user-written
+        // cast on the part literal is invisible to return_field_from_args.
+        if expr.func.name() == "date_part"
+            && let Some(part) = expr.args.first().and_then(expr_string_literal)
+            && is_seconds_part(&part)
+        {
+            let date_expr = expr.args.into_iter().nth(1).ok_or_else(|| {
+                datafusion_common::DataFusionError::Internal(
+                    "date_part expects a date/timestamp argument".to_string(),
+                )
+            })?;
+            return Ok(PlannerResult::Planned(Expr::ScalarFunction(
+                ScalarFunction::new_udf(date_part_seconds(), vec![date_expr]),
+            )));
+        }
+
+        Ok(PlannerResult::Original(expr))
+    }
+}
+
+/// String value of a literal, or of a cast/try_cast wrapping a string literal.
+fn expr_string_literal(expr: &Expr) -> Option<String> {
+    match expr {
+        Expr::Literal(sv, _) => match sv {
+            ScalarValue::Utf8(Some(s))
+            | ScalarValue::Utf8View(Some(s))
+            | ScalarValue::LargeUtf8(Some(s)) => Some(s.clone()),
+            _ => None,
+        },
+        Expr::Cast(c) => expr_string_literal(c.expr.as_ref()),
+        Expr::TryCast(c) => expr_string_literal(c.expr.as_ref()),
+        _ => None,
+    }
 }

datafusion/sqllogictest/test_files/spark/datetime/date_part.slt

@@ -188,31 +188,67 @@ SELECT date_part('MINUTES'::string, '2000-01-01 12:30:45'::timestamp);
 ----
 30
 
-# SECOND
-query I
+# SECOND — Spark returns DECIMAL(8, 6) including the fraction
+query R
 SELECT date_part('SECOND'::string, '2000-01-01 12:30:45'::timestamp);
 ----
-45
+45.000000
 
-query I
+query R
 SELECT date_part('S'::string, '2000-01-01 12:30:45'::timestamp);
 ----
-45
+45.000000
 
-query I
+query R
 SELECT date_part('SEC'::string, '2000-01-01 12:30:45'::timestamp);
 ----
-45
+45.000000
 
-query I
+query R
 SELECT date_part('SECONDS'::string, '2000-01-01 12:30:45'::timestamp);
 ----
-45
+45.000000
 
-query I
+query R
 SELECT date_part('SECS'::string, '2000-01-01 12:30:45'::timestamp);
 ----
-45
+45.000000
+
+query R
+SELECT date_part('SECOND'::string, '2000-01-01 12:30:45.123456'::timestamp);
+----
+45.123456
+
+query R
+SELECT date_part('secs'::string, '2000-01-01 12:30:45.123456789'::timestamp);
+----
+45.123456
+
+query T
+SELECT arrow_typeof(date_part('SECOND'::string, '2000-01-01 12:30:45.123456'::timestamp));
+----
+Decimal128(8, 6)
+
+query R
+SELECT date_part('SECOND'::string, '2000-01-01'::date);
+----
+0.000000
+
+query R
+SELECT date_part('SECOND'::string, NULL::timestamp);
+----
+NULL
+
+# EXTRACT(SECOND FROM …) shares the DECIMAL(8, 6) path via the Spark planner
+query R
+SELECT extract(SECOND FROM timestamp '2000-01-01 12:30:45.123456');
+----
+45.123456
+
+query T
+SELECT arrow_typeof(extract(SECOND FROM timestamp '2000-01-01 12:30:45'));
+----
+Decimal128(8, 6)
 
 # NULL input
 query I
@@ -272,9 +308,8 @@ FROM VALUES
 2000
 NULL
 
-# Nullable even for a non-null input, like the integer parts of the date_part
-# this simplifies to
+# SECOND returns DECIMAL(8, 6); nullable even for a non-null input
 query TTT
 DESCRIBE SELECT date_part('SECOND'::string, ts) AS s FROM (VALUES (timestamp '2020-01-01')) AS t(ts)
 ----
-s Int32 YES
+s Decimal128(8, 6) YES