feat: extract TimeType to int/decimal

native/spark-expr/src/comet_scalar_funcs.rs

@@ -28,7 +28,7 @@ use crate::{
     spark_round, spark_rpad, spark_to_time, spark_unhex, spark_unscaled_value, EvalMode,
     SparkArrayCompact, SparkArrayPositionFunc, SparkArraySlice, SparkArraysOverlap, SparkContains,
     SparkDateDiff, SparkDateFromUnixDate, SparkDateTrunc, SparkFlatten, SparkMakeDate,
-    SparkMakeTime, SparkNextDay, SparkSecondsToTimestamp, SparkSizeFunc,
+    SparkMakeTime, SparkNextDay, SparkSecondWithFraction, SparkSecondsToTimestamp, SparkSizeFunc,
 };
 use arrow::datatypes::DataType;
 use datafusion::common::{DataFusionError, Result as DataFusionResult};
@@ -256,6 +256,7 @@ fn all_scalar_functions() -> Vec<Arc<ScalarUDF>> {
         Arc::new(ScalarUDF::new_from_impl(SparkMakeTime::default())),
         Arc::new(ScalarUDF::new_from_impl(SparkNextDay::default())),
         Arc::new(ScalarUDF::new_from_impl(SparkSecondsToTimestamp::default())),
+        Arc::new(ScalarUDF::new_from_impl(SparkSecondWithFraction::default())),
         Arc::new(ScalarUDF::new_from_impl(SparkSizeFunc::default())),
         Arc::new(ScalarUDF::new_from_impl(JsonArrayLength::default())),
     ]

native/spark-expr/src/datetime_funcs/extract_date_part.rs

@@ -16,13 +16,21 @@
 // under the License.
 
 use crate::utils::array_with_timezone;
+use arrow::array::{Array, Decimal128Builder, Int32Array, Time64NanosecondArray};
 use arrow::compute::{date_part, DatePart};
 use arrow::datatypes::{DataType, TimeUnit::Microsecond};
-use datafusion::common::{internal_datafusion_err, DataFusionError};
+use datafusion::common::{
+    internal_datafusion_err, utils::take_function_args, DataFusionError, Result,
+};
 use datafusion::logical_expr::{
     ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility,
 };
-use std::fmt::Debug;
+use std::{fmt::Debug, sync::Arc};
+
+const NANOS_PER_SECOND: i64 = 1_000_000_000;
+const MICROS_PER_SECOND: i128 = 1_000_000;
+const SECOND_DECIMAL_PRECISION: u8 = 8;
+const SECOND_DECIMAL_SCALE: i8 = 6;
 
 /// Returns true when the type is a timestamp without a timezone (Spark's TimestampNTZType),
 /// including when wrapped in a dictionary. Such values store local wall-clock time and must not
@@ -117,10 +125,113 @@ extract_date_part!(SparkHour, "hour", Hour);
 extract_date_part!(SparkMinute, "minute", Minute);
 extract_date_part!(SparkSecond, "second", Second);
 
+fn second_with_fraction(nanos: i64, precision: i32) -> Result<i128> {
+    if !(0..=SECOND_DECIMAL_SCALE as i32).contains(&precision) {
+        return Err(DataFusionError::Execution(format!(
+            "second_with_fraction: precision must be between 0 and 6, found {precision}"
+        )));
+    }
+
+    let whole_seconds = (nanos / NANOS_PER_SECOND) % 60;
+    let fraction_nanos = nanos % NANOS_PER_SECOND;
+    let scale_factor: i64 = 10_i64.pow(precision as u32);
+    let fraction_digits = (fraction_nanos * scale_factor) / NANOS_PER_SECOND;
+    let output_scale_factor = 10_i128.pow((SECOND_DECIMAL_SCALE as i32 - precision) as u32);
+    let result =
+        whole_seconds as i128 * MICROS_PER_SECOND + fraction_digits as i128 * output_scale_factor;
+
+    Ok(result)
+}
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct SparkSecondWithFraction {
+    signature: Signature,
+}
+
+impl SparkSecondWithFraction {
+    pub fn new() -> Self {
+        Self {
+            signature: Signature::any(2, Volatility::Immutable),
+        }
+    }
+}
+
+impl Default for SparkSecondWithFraction {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl ScalarUDFImpl for SparkSecondWithFraction {
+    fn name(&self) -> &str {
+        "second_with_fraction"
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, _: &[DataType]) -> Result<DataType> {
+        Ok(DataType::Decimal128(
+            SECOND_DECIMAL_PRECISION,
+            SECOND_DECIMAL_SCALE,
+        ))
+    }
+
+    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        let [time, precision] = take_function_args(self.name(), args.args)?;
+
+        let num_rows = [&time, &precision]
+            .iter()
+            .find_map(|arg| match arg {
+                ColumnarValue::Array(array) => Some(array.len()),
+                ColumnarValue::Scalar(_) => None,
+            })
+            .unwrap_or(1);
+
+        let time = time.into_array(num_rows)?;
+        let precision = precision.into_array(num_rows)?;
+        let time = time
+            .as_any()
+            .downcast_ref::<Time64NanosecondArray>()
+            .ok_or_else(|| {
+                DataFusionError::Execution(
+                    "second_with_fraction: expected Time64(Nanosecond) input".to_string(),
+                )
+            })?;
+        let precision = precision
+            .as_any()
+            .downcast_ref::<Int32Array>()
+            .ok_or_else(|| {
+                DataFusionError::Execution(
+                    "second_with_fraction: expected Int32 precision".to_string(),
+                )
+            })?;
+
+        let mut result = Decimal128Builder::with_capacity(num_rows);
+
+        for index in 0..num_rows {
+            if time.is_null(index) || precision.is_null(index) {
+                result.append_null();
+            } else {
+                let value = second_with_fraction(time.value(index), precision.value(index))?;
+
+                result.append_value(value);
+            }
+        }
+
+        let result = result
+            .with_precision_and_scale(SECOND_DECIMAL_PRECISION, SECOND_DECIMAL_SCALE)?
+            .finish();
+
+        Ok(ColumnarValue::Array(Arc::new(result)))
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
-    use arrow::array::{ArrayRef, Int32Array, TimestampMicrosecondArray};
+    use arrow::array::{ArrayRef, Int32Array, Time64NanosecondArray, TimestampMicrosecondArray};
     use arrow::datatypes::{Field, TimeUnit};
     use datafusion::config::ConfigOptions;
     use std::sync::Arc;
@@ -256,4 +367,34 @@ mod tests {
             .value(key);
         assert_eq!(extracted, 18);
     }
+
+    #[test]
+    fn time64_extracts_time_parts_without_timezone_conversion() {
+        let nanos = (12 * 60 * 60 + 30 * 60 + 45) * 1_000_000_000;
+        let array = Arc::new(Time64NanosecondArray::from(vec![Some(nanos)])) as ArrayRef;
+
+        for timezone in ["UTC", "America/Los_Angeles", "Asia/Tokyo"] {
+            assert_eq!(
+                invoke(&SparkHour::new(timezone.to_string()), Arc::clone(&array)),
+                12
+            );
+            assert_eq!(
+                invoke(&SparkMinute::new(timezone.to_string()), Arc::clone(&array)),
+                30
+            );
+            assert_eq!(
+                invoke(&SparkSecond::new(timezone.to_string()), Arc::clone(&array)),
+                45
+            );
+        }
+    }
+
+    #[test]
+    fn time64_extract_seconds_with_fraction() {
+        let nanos = (12 * 60 * 60 + 30 * 60 + 45) * NANOS_PER_SECOND + 123_456_000;
+
+        assert_eq!(second_with_fraction(nanos, 0).unwrap(), 45_000_000);
+        assert_eq!(second_with_fraction(nanos, 3).unwrap(), 45_123_000);
+        assert_eq!(second_with_fraction(nanos, 6).unwrap(), 45_123_456);
+    }
 }

native/spark-expr/src/datetime_funcs/mod.rs

@@ -35,7 +35,7 @@ pub use date_trunc::SparkDateTrunc;
 pub use day_month_name::{spark_day_name, spark_month_name};
 pub use extract_date_part::SparkHour;
 pub use extract_date_part::SparkMinute;
-pub use extract_date_part::SparkSecond;
+pub use extract_date_part::{SparkSecond, SparkSecondWithFraction};
 pub use hours::SparkHoursTransform;
 pub use make_date::SparkMakeDate;
 pub use make_time::SparkMakeTime;

native/spark-expr/src/lib.rs

@@ -78,7 +78,8 @@ pub use csv_funcs::*;
 pub use datetime_funcs::{
     spark_day_name, spark_month_name, spark_to_time, SparkDateDiff, SparkDateFromUnixDate,
     SparkDateTrunc, SparkHour, SparkHoursTransform, SparkMakeDate, SparkMakeTime, SparkMinute,
-    SparkNextDay, SparkSecond, SparkSecondsToTimestamp, SparkUnixTimestamp, TimestampTruncExpr,
+    SparkNextDay, SparkSecond, SparkSecondWithFraction, SparkSecondsToTimestamp,
+    SparkUnixTimestamp, TimestampTruncExpr,
 };
 pub use error::{decimal_overflow_error, SparkError, SparkErrorWithContext, SparkResult};
 pub use hash_funcs::*;

spark/src/main/spark-4.1+/org/apache/comet/shims/CometExprShim.scala

@@ -24,9 +24,10 @@ import org.apache.spark.sql.catalyst.expressions.aggregate.Sum
 import org.apache.spark.sql.catalyst.expressions.objects.{Invoke, StaticInvoke}
 import org.apache.spark.sql.catalyst.util.DateTimeUtils
 import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.types.TimeType
+import org.apache.spark.sql.types.{DecimalType, IntegerType, TimeType}
 
 import org.apache.comet.expressions.CometEvalMode
+import org.apache.comet.serde.ExprOuterClass
 import org.apache.comet.serde.ExprOuterClass.{BinaryOutputStyle, Expr}
 import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, optExprWithFallbackReason, scalarFunctionExprToProtoWithReturnType}
 
@@ -46,16 +47,50 @@ trait CometExprShim extends Spark4xCometExprShim {
     }
   }
 
-  // Spark 4.1 introduced TimeType and the make_time / to_time / try_to_time functions.
-  // Their planner forms differ from the shared 4.x patterns (DateTimeUtils.makeTime
-  // StaticInvoke and ToTimeParser Invoke / TryEval(Invoke)), so they live here rather
-  // than in the shared Spark4xCometExprShim parent trait. Spark 4.0 lacks TimeType, which
-  // is why this shim is shared by 4.1 and later rather than all 4.x.
+  private val integralTimePartMethods =
+    Set("getHoursOfTime", "getMinutesOfTime", "getSecondsOfTime")
+
+  // Spark 4.1 introduced TimeType and its extraction, make_time, to_time, and try_to_time
+  // expressions. Their planner forms differ from the shared 4.x patterns (DateTimeUtils
+  // StaticInvoke and ToTimeParser Invoke / TryEval(Invoke)), so they live here rather than in the
+  // shared Spark4xCometExprShim parent trait. Spark 4.0 lacks TimeType, which is why this shim is
+  // shared by 4.1 and later rather than all 4.x.
   override def sparkVersionSpecificExprToProtoInternal(
       expr: Expression,
       inputs: Seq[Attribute],
       binding: Boolean): Option[Expr] = {
+
     expr match {
+      case s: StaticInvoke
+          if s.staticObject == classOf[DateTimeUtils.type] &&
+            integralTimePartMethods.contains(s.functionName) &&
+            s.arguments.size == 1 &&
+            s.dataType == IntegerType &&
+            s.arguments.head.dataType.isInstanceOf[TimeType] =>
+        val child = s.arguments.head
+        val childExpr = exprToProtoInternal(child, inputs, binding)
+
+        val optExpr = childExpr.map { childProto =>
+          timePartToProto(s.functionName, childProto)
+        }
+
+        optExprWithFallbackReason(optExpr, s, child)
+
+      case s: StaticInvoke
+          if s.staticObject == classOf[DateTimeUtils.type] &&
+            s.functionName == "getSecondsOfTimeWithFraction" &&
+            s.arguments.size == 2 &&
+            s.dataType == DecimalType(8, 6) &&
+            s.arguments.head.dataType.isInstanceOf[TimeType] &&
+            isValidTimePrecision(s.arguments(1)) =>
+        val childExprs = s.arguments.map(exprToProtoInternal(_, inputs, binding))
+        val optExpr = scalarFunctionExprToProtoWithReturnType(
+          "second_with_fraction",
+          s.dataType,
+          false,
+          childExprs: _*)
+        optExprWithFallbackReason(optExpr, s, s.arguments: _*)
+
       case s: StaticInvoke
           if s.staticObject == classOf[DateTimeUtils.type] &&
             s.functionName == "makeTime" &&
@@ -97,6 +132,37 @@ trait CometExprShim extends Spark4xCometExprShim {
       case _ => super.sparkVersionSpecificExprToProtoInternal(expr, inputs, binding)
     }
   }
+
+  private def timePartToProto(functionName: String, childProto: Expr): Expr = {
+    functionName match {
+      case "getHoursOfTime" =>
+        val builder = ExprOuterClass.Hour.newBuilder()
+        builder.setChild(childProto)
+        builder.setTimezone("UTC")
+        Expr.newBuilder().setHour(builder).build()
+
+      case "getMinutesOfTime" =>
+        val builder = ExprOuterClass.Minute.newBuilder()
+        builder.setChild(childProto)
+        builder.setTimezone("UTC")
+        Expr.newBuilder().setMinute(builder).build()
+
+      case "getSecondsOfTime" =>
+        val builder = ExprOuterClass.Second.newBuilder()
+        builder.setChild(childProto)
+        builder.setTimezone("UTC")
+        Expr.newBuilder().setSecond(builder).build()
+
+      case other =>
+        throw new IllegalArgumentException(s"Unsupported integral TimeType part: $other")
+    }
+  }
+
+  private def isValidTimePrecision(expr: Expression): Boolean = expr match {
+    case Literal(precision: Int, IntegerType) =>
+      precision >= TimeType.MIN_PRECISION && precision <= TimeType.MAX_PRECISION
+    case _ => false
+  }
 }
 
 object CometEvalModeUtil {

spark/src/test/resources/sql-tests/expressions/datetime/time_extract.sql

@@ -0,0 +1,49 @@
+-- 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.
+
+-- MinSparkVersion: 4.1
+-- Config: spark.sql.timeType.enabled=true
+
+statement
+CREATE TABLE test_time_extract(hours int, minutes int, secs DECIMAL(16, 6)) USING parquet
+
+statement
+INSERT INTO test_time_extract VALUES
+    (0, 0, 0.000000),
+    (12, 30, 45.123456),
+    (23, 59, 59.999999),
+    (NULL, NULL, NULL)
+
+query
+SELECT hour(make_time(hours, minutes, secs))
+FROM test_time_extract
+
+query
+SELECT minute(make_time(hours, minutes, secs))
+FROM test_time_extract
+
+query
+SELECT second(make_time(hours, minutes, secs))
+FROM test_time_extract
+
+query
+SELECT extract(SECOND FROM make_time(hours, minutes, secs))
+FROM test_time_extract
+
+query
+SELECT date_part('SECOND', make_time(hours, minutes, secs))
+FROM test_time_extract