Add format support for Time

Author: vinodkc

python/pyspark/sql/connect/functions/builtin.py

@@ -3657,6 +3657,13 @@ def time_diff(unit: "ColumnOrName", start: "ColumnOrName", end: "ColumnOrName")
 time_diff.__doc__ = pysparkfuncs.time_diff.__doc__
 
 
+def time_format(time: "ColumnOrName", format: str) -> Column:
+    return _invoke_function("time_format", _to_col(time), lit(format))
+
+
+time_format.__doc__ = pysparkfuncs.time_format.__doc__
+
+
 def time_trunc(unit: "ColumnOrName", time: "ColumnOrName") -> Column:
     return _invoke_function_over_columns("time_trunc", unit, time)
 

python/pyspark/sql/functions/__init__.py

@@ -249,6 +249,7 @@
     "timestamp_millis",
     "timestamp_seconds",
     "time_diff",
+    "time_format",
     "time_from_micros",
     "time_from_millis",
     "time_from_seconds",

python/pyspark/sql/functions/builtin.py

@@ -12829,6 +12829,78 @@ def time_diff(unit: "ColumnOrName", start: "ColumnOrName", end: "ColumnOrName")
     return _invoke_function_over_columns("time_diff", unit, start, end)
 
 
+@_try_remote_functions
+def time_format(time: "ColumnOrName", format: str) -> Column:
+    """
+    Converts a time to a value of string in the format specified by the time
+    format given by the second argument.
+
+    A pattern could be for instance `HH:mm:ss` and could return a string like '14:30:45'.
+    Time-related pattern letters of `datetime pattern`_ can be used.
+
+    .. versionadded:: 4.2.0
+
+    Parameters
+    ----------
+    time : :class:`~pyspark.sql.Column` or column name
+        input column of TIME values to format.
+    format: literal string
+        format to use to represent time values.
+
+    Returns
+    -------
+    :class:`~pyspark.sql.Column`
+        string value representing formatted time.
+
+    See Also
+    --------
+    :meth:`pyspark.sql.functions.to_time`
+    :meth:`pyspark.sql.functions.try_to_time`
+    :meth:`pyspark.sql.functions.date_format`
+
+    Examples
+    --------
+    Example 1: Basic 24-hour format
+
+    >>> from pyspark.sql import functions as sf
+    >>> df = spark.sql("SELECT TIME'14:30:45' as time_col")
+    >>> df.select("*", sf.time_format('time_col', 'HH:mm:ss')).show()
+    +--------+-------------------------------+
+    |time_col|time_format(time_col, HH:mm:ss)|
+    +--------+-------------------------------+
+    |14:30:45|                       14:30:45|
+    +--------+-------------------------------+
+    <BLANKLINE>
+
+    Example 2: 12-hour format with AM/PM
+
+    >>> from pyspark.sql import functions as sf
+    >>> df = spark.sql("SELECT TIME'14:30:45' as time_col")
+    >>> df.select("*", sf.time_format('time_col', 'hh:mm:ss a')).show()
+    +--------+---------------------------------+
+    |time_col|time_format(time_col, hh:mm:ss a)|
+    +--------+---------------------------------+
+    |14:30:45|                      02:30:45 PM|
+    +--------+---------------------------------+
+    <BLANKLINE>
+
+    Example 3: With microseconds
+
+    >>> from pyspark.sql import functions as sf
+    >>> df = spark.sql("SELECT TIME'14:30:45.123456' as time_col")
+    >>> df.select("*", sf.time_format('time_col', 'HH:mm:ss.SSSSSS')).show()
+    +---------------+--------------------------------------+
+    |       time_col|time_format(time_col, HH:mm:ss.SSSSSS)|
+    +---------------+--------------------------------------+
+    |14:30:45.123456|                       14:30:45.123456|
+    +---------------+--------------------------------------+
+    <BLANKLINE>
+    """
+    from pyspark.sql.classic.column import _to_java_column
+
+    return _invoke_function("time_format", _to_java_column(time), _enum_to_value(format))
+
+
 @_try_remote_functions
 def time_trunc(unit: "ColumnOrName", time: "ColumnOrName") -> Column:
     """

sql/api/src/main/scala/org/apache/spark/sql/functions.scala

@@ -6889,6 +6889,22 @@ object functions {
     Column.fn("time_diff", unit, start, end)
   }
 
+  /**
+   * Converts a time to a string in the specified format.
+   *
+   * @param time
+   *   A column of time values to be formatted.
+   * @param format
+   *   A time format string. for valid patterns.
+   * @return
+   *   String representation of the time in the specified format.
+   * @group datetime_funcs
+   * @since 4.2.0
+   */
+  def time_format(time: Column, format: String): Column = {
+    Column.fn("time_format", time, lit(format))
+  }
+
   /**
    * Returns `time` truncated to the `unit`.
    *

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala

@@ -680,6 +680,7 @@ object FunctionRegistry {
     expression[ParseToTimestamp]("to_timestamp"),
     expression[ParseToDate]("to_date"),
     expression[TimeDiff]("time_diff"),
+    expression[TimeFormat]("time_format"),
     expression[ToTime]("to_time"),
     expression[ToBinary]("to_binary"),
     expression[ToUnixTimestamp]("to_unix_timestamp"),

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/timeExpressions.scala

@@ -35,7 +35,7 @@ import org.apache.spark.sql.catalyst.util.TimeFormatter
 import org.apache.spark.sql.catalyst.util.TypeUtils.ordinalNumber
 import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors}
 import org.apache.spark.sql.internal.types.StringTypeWithCollation
-import org.apache.spark.sql.types.{AbstractDataType, AnyTimeType, ByteType, DataType, DayTimeIntervalType, Decimal, DecimalType, DoubleType, FloatType, IntegerType, IntegralType, LongType, NumericType, ObjectType, TimeType}
+import org.apache.spark.sql.types.{AbstractDataType, AnyTimeType, ByteType, DataType, DayTimeIntervalType, Decimal, DecimalType, DoubleType, FloatType, IntegerType, IntegralType, LongType, NumericType, ObjectType, StringType, TimeType}
 import org.apache.spark.sql.types.DayTimeIntervalType.{HOUR, SECOND}
 import org.apache.spark.unsafe.types.UTF8String
 
@@ -1036,3 +1036,82 @@ case class TimeToMicros(child: Expression)
   override protected def withNewChildInternal(newChild: Expression): TimeToMicros =
     copy(child = newChild)
 }
+
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+  usage = "_FUNC_(time, format) - Converts a time to a value of string in the format specified by the date format given by the second argument.",
+  arguments = """
+    Arguments:
+      * time - A time value to be converted to string.
+      * format - Time format pattern to follow. See <a href="https://spark.apache.org/docs/latest/sql-ref-datetime-pattern.html">Datetime Patterns</a> for valid
+                 time format patterns. Note: Only time-related patterns (H, h, m, s, S, a) are meaningful for TIME values.
+  """,
+  examples = """
+    Examples:
+      > SELECT _FUNC_(TIME'14:30:45', 'HH:mm:ss');
+       14:30:45
+      > SELECT _FUNC_(TIME'14:30:45', 'hh:mm:ss a');
+       02:30:45 PM
+      > SELECT _FUNC_(TIME'14:30:45.123456', 'HH:mm:ss.SSSSSS');
+       14:30:45.123456
+      > SELECT _FUNC_(TIME'09:05:00', 'h:mm a');
+       9:05 AM
+  """,
+  group = "datetime_funcs",
+  since = "4.2.0")
+// scalastyle:on line.size.limit
+case class TimeFormat(left: Expression, right: Expression)
+  extends BinaryExpression
+  with ImplicitCastInputTypes {
+
+  override def nullIntolerant: Boolean = true
+
+  override def inputTypes: Seq[AbstractDataType] =
+    Seq(AnyTimeType, StringTypeWithCollation(supportsTrimCollation = true))
+
+  override def dataType: DataType = StringType
+
+  // Cache the formatter if the format string is a foldable expression
+  @transient private lazy val formatterOption: Option[TimeFormatter] =
+    if (right.foldable) {
+      Option(right.eval()).map { format =>
+        TimeFormatter(format.toString, TimeFormatter.defaultLocale, isParsing = false)
+      }
+    } else {
+      None
+    }
+
+  override protected def nullSafeEval(time: Any, format: Any): Any = {
+    val nanos = time.asInstanceOf[Long]
+    val formatter = formatterOption.getOrElse {
+      TimeFormatter(format.toString, TimeFormatter.defaultLocale, isParsing = false)
+    }
+    UTF8String.fromString(formatter.format(nanos))
+  }
+
+  override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+    formatterOption.map { tf =>
+      val timeFormatter = ctx.addReferenceObj("timeFormatter", tf)
+      defineCodeGen(ctx, ev, (time, _) => {
+        s"""UTF8String.fromString($timeFormatter.format($time))"""
+      })
+    }.getOrElse {
+      val tfClass = TimeFormatter.getClass.getName.stripSuffix("$")
+      val locale = ctx.addReferenceObj(
+        "locale", TimeFormatter.defaultLocale, classOf[Locale].getName)
+      defineCodeGen(ctx, ev, (time, format) => {
+        s"""|UTF8String.fromString($tfClass$$.MODULE$$.apply(
+            |  $format.toString(),
+            |  $locale,
+            |  false)
+            |.format($time))""".stripMargin
+      })
+    }
+  }
+
+  override def prettyName: String = "time_format"
+
+  override protected def withNewChildrenInternal(
+      newLeft: Expression, newRight: Expression): TimeFormat =
+    copy(left = newLeft, right = newRight)
+}

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/TimeExpressionsSuite.scala

@@ -666,4 +666,66 @@ class TimeExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
     val timeVal = TimeFromSeconds(Literal(secondsValue))
     checkEvaluation(TimeToMicros(timeVal), 14500000L)
   }
+
+  test("TimeFormat - foldable and non-foldable formats with codegen") {
+    val time = localTime(14, 30, 45, 123456)
+    val timeLit = Literal(time, TimeType())
+
+    // Foldable format - formatter cached at planning time
+    val foldableExpr = TimeFormat(timeLit, Literal("HH:mm:ss"))
+    checkEvaluation(foldableExpr, "14:30:45")
+    checkEvaluation(
+      TimeFormat(Literal(localTime(9, 5, 0), TimeType()), Literal("hh:mm:ss a")),
+      "09:05:00 AM")
+    checkEvaluation(
+      TimeFormat(timeLit, Literal("HH:mm:ss.SSSSSS")),
+      "14:30:45.123456")
+
+    // Non-foldable format - formatter created per evaluation
+    val nonFoldableExpr = TimeFormat(timeLit, NonFoldableLiteral("HH:mm:ss", StringType))
+    checkEvaluation(nonFoldableExpr, "14:30:45")
+    checkEvaluation(
+      TimeFormat(timeLit, NonFoldableLiteral("hh:mm:ss a", StringType)),
+      "02:30:45 PM")
+    checkEvaluation(
+      TimeFormat(timeLit, NonFoldableLiteral("HH:mm:ss.SSSSSS", StringType)),
+      "14:30:45.123456")
+
+    // Multiple formats on same time
+    Seq("HH:mm:ss", "hh:mm a", "HH:mm", "HH").zip(
+      Seq("14:30:45", "02:30 PM", "14:30", "14")
+    ).foreach { case (format, expected) =>
+      checkEvaluation(
+        TimeFormat(timeLit, NonFoldableLiteral(format, StringType)),
+        expected)
+    }
+
+    // Edge cases
+    checkEvaluation(
+      TimeFormat(Literal(localTime(0, 0, 0), TimeType()), Literal("HH:mm:ss")),
+      "00:00:00")
+    checkEvaluation(
+      TimeFormat(Literal(localTime(23, 59, 59, 999999), TimeType()),
+        Literal("HH:mm:ss.SSSSSS")),
+      "23:59:59.999999")
+
+    // Null handling
+    checkEvaluation(
+      TimeFormat(Literal.create(null, TimeType()), Literal("HH:mm:ss")),
+      null)
+    checkEvaluation(
+      TimeFormat(timeLit, Literal.create(null, StringType)),
+      null)
+    checkEvaluation(
+      TimeFormat(Literal.create(null, TimeType()), NonFoldableLiteral("HH:mm:ss", StringType)),
+      null)
+    checkEvaluation(
+      TimeFormat(timeLit, NonFoldableLiteral(null, StringType)),
+      null)
+
+    // Verify foldable and non-foldable produce same result
+    val foldableResult = foldableExpr.eval()
+    val nonFoldableResult = nonFoldableExpr.eval()
+    assert(foldableResult == nonFoldableResult)
+  }
 }

sql/core/src/test/resources/sql-functions/sql-expression-schema.md

@@ -363,6 +363,7 @@
 | org.apache.spark.sql.catalyst.expressions.ThetaSketchEstimate | theta_sketch_estimate | SELECT theta_sketch_estimate(theta_sketch_agg(col)) FROM VALUES (1), (1), (2), (2), (3) tab(col) | struct<theta_sketch_estimate(theta_sketch_agg(col, 12)):bigint> |
 | org.apache.spark.sql.catalyst.expressions.ThetaUnion | theta_union | SELECT theta_sketch_estimate(theta_union(theta_sketch_agg(col1), theta_sketch_agg(col2))) FROM VALUES (1, 4), (1, 4), (2, 5), (2, 5), (3, 6) tab(col1, col2) | struct<theta_sketch_estimate(theta_union(theta_sketch_agg(col1, 12), theta_sketch_agg(col2, 12), 12)):bigint> |
 | org.apache.spark.sql.catalyst.expressions.TimeDiff | time_diff | SELECT time_diff('HOUR', TIME'20:30:29', TIME'21:30:28') | struct<time_diff(HOUR, TIME '20:30:29', TIME '21:30:28'):bigint> |
+| org.apache.spark.sql.catalyst.expressions.TimeFormat | time_format | SELECT time_format(TIME'14:30:45', 'HH:mm:ss') | struct<time_format(TIME '14:30:45', HH:mm:ss):string> |
 | org.apache.spark.sql.catalyst.expressions.TimeFromMicros | time_from_micros | SELECT time_from_micros(0) | struct<time_from_micros(0):time(6)> |
 | org.apache.spark.sql.catalyst.expressions.TimeFromMillis | time_from_millis | SELECT time_from_millis(0) | struct<time_from_millis(0):time(6)> |
 | org.apache.spark.sql.catalyst.expressions.TimeFromSeconds | time_from_seconds | SELECT time_from_seconds(0) | struct<time_from_seconds(0):time(6)> |