From 4bae3338befbbc623e440fe0bba77086d2b59d26 Mon Sep 17 00:00:00 2001 From: Vincent <52391318+ump45nose@users.noreply.github.com> Date: Fri, 30 Oct 2026 05:19:01 +0800 Subject: [PATCH] fix(spark): date_part SECOND returns Decimal(8,6) Spark's date_part('SECOND', ts) / EXTRACT(SECOND FROM ts) return DECIMAL(8, 6) with the microsecond fraction. Rewrite those calls to a dedicated seconds UDF at plan time so the return type is fixed before analysis, and keep standalone second(ts) as Int32. --- .../spark/src/function/datetime/date_part.rs | 156 +++++++++++++++++- datafusion/spark/src/planner.rs | 58 ++++++- .../test_files/spark/datetime/date_part.slt | 63 +++++-- 3 files changed, 255 insertions(+), 22 deletions(-) diff --git a/datafusion/spark/src/function/datetime/date_part.rs b/datafusion/spark/src/function/datetime/date_part.rs index 672a6e060c8d7..87b9db83b2ca3 100644 --- a/datafusion/spark/src/function/datetime/date_part.rs +++ b/datafusion/spark/src/function/datetime/date_part.rs @@ -2,7 +2,9 @@ // 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, @@ -10,10 +12,40 @@ use datafusion_common::{ 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 { + static INSTANCE: OnceLock> = 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 { - 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 { + 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 { @@ -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 { + Ok(DataType::Decimal128(SECONDS_PRECISION, SECONDS_SCALE)) + } + + fn return_field_from_args(&self, _args: ReturnFieldArgs) -> Result { + Ok(Arc::new(Field::new( + self.name(), + DataType::Decimal128(SECONDS_PRECISION, SECONDS_SCALE), + true, + ))) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + 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 { + 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::(); + let micros = micros.as_primitive::(); + + // `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)) +} diff --git a/datafusion/spark/src/planner.rs b/datafusion/spark/src/planner.rs index e6908f9ef378d..493e4df2ad441 100644 --- a/datafusion/spark/src/planner.rs +++ b/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, ) -> datafusion_common::Result>> { + // 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> { + // 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 { + 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, + } } diff --git a/datafusion/sqllogictest/test_files/spark/datetime/date_part.slt b/datafusion/sqllogictest/test_files/spark/datetime/date_part.slt index 764f30d71ec6a..4024b77e370d1 100644 --- a/datafusion/sqllogictest/test_files/spark/datetime/date_part.slt +++ b/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