Skip to content

feat: Support Spark expression hours_of_time #3615

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
0lai0 wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

feat: Support Spark expression hours_of_time #3615

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0b59dac
Support HourOfTime expression via version-specific shims
0lai0 Mar 1, 2026
45b2236
modify strings.scala
0lai0 Mar 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 46 additions & 0 deletions spark/src/main/scala/org/apache/comet/serde/strings.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -408,4 +408,50 @@ trait CommonStringExprs {
None
}
}

def hoursOfTimeToProto(
Copy link
Member

@andygrove andygrove Mar 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It might make more sense to place hoursOfTimeToProto alongside the existing CometHour handler in datetime.scala, since that's where the Hour serde already lives.

Copy link
Contributor Author

@0lai0 0lai0 Mar 4, 2026 •
edited by martin-g
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @andygrove for the feedback. I'll prioritize completing secondsoftime first since there are components we can share between the two.

expr: Expression,
inputs: Seq[Attribute],
binding: Boolean): Option[Expr] = {
val childOpt = expr.children.headOption.orElse {
withInfo(expr, "HoursOfTime has no child expression")
None
}

childOpt.flatMap { child =>
val timeZoneId = {
val exprClass = expr.getClass
try {
val timeZoneIdMethod = exprClass.getMethod("timeZoneId")
timeZoneIdMethod.invoke(expr).asInstanceOf[Option[String]]
} catch {
case _: NoSuchMethodException =>
try {
val timeZoneIdField = exprClass.getField("timeZoneId")
timeZoneIdField.get(expr).asInstanceOf[Option[String]]
} catch {
case _: NoSuchFieldException | _: SecurityException => None
}
}
}

exprToProtoInternal(child, inputs, binding)
.map { childExpr =>
val builder = ExprOuterClass.Hour.newBuilder()
builder.setChild(childExpr)

val timeZone = timeZoneId.getOrElse("UTC")
builder.setTimezone(timeZone)

ExprOuterClass.Expr
.newBuilder()
.setHour(builder)
.build()
}
.orElse {
withInfo(expr, child)
None
}
}
}
}
5 changes: 5 additions & 0 deletions spark/src/main/spark-3.4/org/apache/comet/shims/CometExprShim.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,11 @@ package org.apache.comet.shims

import org.apache.spark.sql.catalyst.expressions._

import org.apache.comet.CometSparkSessionExtensions.withInfo
import org.apache.comet.expressions.CometEvalMode
import org.apache.comet.serde.CommonStringExprs
import org.apache.comet.serde.ExprOuterClass.{BinaryOutputStyle, Expr}
import org.apache.comet.serde.QueryPlanSerde.exprToProtoInternal

/**
* `CometExprShim` acts as a shim for parsing expressions from different Spark versions.
Expand All @@ -43,6 +45,9 @@ trait CometExprShim extends CommonStringExprs {
// Right child is the encoding expression.
stringDecode(expr, s.charset, s.bin, inputs, binding)

case _ if expr.getClass.getSimpleName == "HoursOfTime" =>
hoursOfTimeToProto(expr, inputs, binding)

case _ => None
}
}
Expand Down
3 changes: 3 additions & 0 deletions spark/src/main/spark-3.5/org/apache/comet/shims/CometExprShim.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,9 @@ trait CometExprShim extends CommonStringExprs {
// val optExpr = scalarFunctionExprToProto("width_bucket", childExprs: _*)
// optExprWithInfo(optExpr, wb, wb.children: _*)

case _ if expr.getClass.getSimpleName == "HoursOfTime" =>
hoursOfTimeToProto(expr, inputs, binding)

case _ => None
}
}
Expand Down
20 changes: 20 additions & 0 deletions spark/src/main/spark-4.0/org/apache/comet/shims/CometExprShim.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,26 @@ trait CometExprShim extends CommonStringExprs {
// val optExpr = scalarFunctionExprToProto("width_bucket", childExprs: _*)
// optExprWithInfo(optExpr, wb, wb.children: _*)

case h: HoursOfTime =>
exprToProtoInternal(h.child, inputs, binding) match {
case Some(childExpr) =>
val builder = ExprOuterClass.Hour.newBuilder()
builder.setChild(childExpr)
val timeZone = h.timeZoneId.getOrElse("UTC")
builder.setTimezone(timeZone)
Some(
ExprOuterClass.Expr
.newBuilder()
.setHour(builder)
.build())
case None =>
withInfo(h, h.child)
None
}

case _ if expr.getClass.getSimpleName == "HoursOfTime" =>
hoursOfTimeToProto(expr, inputs, binding)

case _ => None
}
}
Expand Down
17 changes: 17 additions & 0 deletions spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -601,6 +601,23 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
}
}

test("hourOfTime expression support") {
// This test verifies that hour() function works correctly with timestamp columns.
// If Spark generates HoursOfTime expression (a RuntimeReplaceable expression),
// it will be handled by the version-specific shim and converted to Hour proto.
Seq(true, false).foreach { dictionaryEnabled =>
withTempDir { dir =>
val path = new Path(dir.toURI.toString, "part-r-0.parquet")
makeRawTimeParquetFile(path, dictionaryEnabled = dictionaryEnabled, 10000)
readParquetFile(path.toString) { df =>
val query = df.select(expr("hour(_1)"))

checkSparkAnswerAndOperator(query)
}
}
}
}

test("cast timestamp and timestamp_ntz") {
withSQLConf(
SESSION_LOCAL_TIMEZONE.key -> "Asia/Kathmandu",
Expand Down
Loading