Github user HyukjinKwon commented on a diff in the pull request:
https://github.com/apache/spark/pull/20015#discussion_r157925994
--- Diff:
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/datetimeExpressions.scala
---
@@ -1295,87 +1295,181 @@ case class ParseToTimestamp(left: Expression,
format: Option[Expression], child:
override def dataType: DataType = TimestampType
}
-/**
- * Returns date truncated to the unit specified by the format.
- */
-// scalastyle:off line.size.limit
-@ExpressionDescription(
- usage = "_FUNC_(date, fmt) - Returns `date` with the time portion of the
day truncated to the unit specified by the format model `fmt`.",
- examples = """
- Examples:
- > SELECT _FUNC_('2009-02-12', 'MM');
- 2009-02-01
- > SELECT _FUNC_('2015-10-27', 'YEAR');
- 2015-01-01
- """,
- since = "1.5.0")
-// scalastyle:on line.size.limit
-case class TruncDate(date: Expression, format: Expression)
- extends BinaryExpression with ImplicitCastInputTypes {
- override def left: Expression = date
- override def right: Expression = format
-
- override def inputTypes: Seq[AbstractDataType] = Seq(DateType,
StringType)
- override def dataType: DataType = DateType
+trait TruncInstant extends BinaryExpression with ImplicitCastInputTypes {
+ val instant: Expression
+ val format: Expression
override def nullable: Boolean = true
- override def prettyName: String = "trunc"
private lazy val truncLevel: Int =
DateTimeUtils.parseTruncLevel(format.eval().asInstanceOf[UTF8String])
- override def eval(input: InternalRow): Any = {
+ /**
+ * @param input internalRow (time)
+ * @param maxLevel Maximum level that can be used for truncation (e.g
MONTH for Date input)
+ * @param truncFunc function: (time, level) => time
+ */
+ protected def evalHelper(input: InternalRow, maxLevel: Int)(
+ truncFunc: (Any, Int) => Any): Any = {
val level = if (format.foldable) {
truncLevel
} else {
DateTimeUtils.parseTruncLevel(format.eval().asInstanceOf[UTF8String])
}
- if (level == -1) {
- // unknown format
+ if (level == DateTimeUtils.TRUNC_INVALID || level > maxLevel) {
+ // unknown format or too large level
null
} else {
- val d = date.eval(input)
- if (d == null) {
+ val t = instant.eval(input)
+ if (t == null) {
null
} else {
- DateTimeUtils.truncDate(d.asInstanceOf[Int], level)
+ truncFunc(t, level)
}
}
}
- override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ protected def codeGenHelper(
+ ctx: CodegenContext,
+ ev: ExprCode,
+ maxLevel: Int,
+ orderReversed: Boolean = false)(
+ truncFunc: (String, String) => String)
+ : ExprCode = {
val dtu = DateTimeUtils.getClass.getName.stripSuffix("$")
if (format.foldable) {
- if (truncLevel == -1) {
+ if (truncLevel == DateTimeUtils.TRUNC_INVALID || truncLevel >
maxLevel) {
ev.copy(code = s"""
boolean ${ev.isNull} = true;
${ctx.javaType(dataType)} ${ev.value} =
${ctx.defaultValue(dataType)};""")
} else {
- val d = date.genCode(ctx)
+ val t = instant.genCode(ctx)
+ val truncFuncStr = truncFunc(t.value, truncLevel.toString)
ev.copy(code = s"""
- ${d.code}
- boolean ${ev.isNull} = ${d.isNull};
+ ${t.code}
+ boolean ${ev.isNull} = ${t.isNull};
${ctx.javaType(dataType)} ${ev.value} =
${ctx.defaultValue(dataType)};
if (!${ev.isNull}) {
- ${ev.value} = $dtu.truncDate(${d.value}, $truncLevel);
+ ${ev.value} = $dtu.$truncFuncStr;
}""")
}
} else {
- nullSafeCodeGen(ctx, ev, (dateVal, fmt) => {
+ nullSafeCodeGen(ctx, ev, (left, right) => {
val form = ctx.freshName("form")
+ val (dateVal, fmt) = if (orderReversed) {
+ (right, left)
+ } else {
+ (left, right)
+ }
+ val truncFuncStr = truncFunc(dateVal, form)
s"""
int $form = $dtu.parseTruncLevel($fmt);
- if ($form == -1) {
+ if ($form == -1 || $form > $maxLevel) {
${ev.isNull} = true;
} else {
- ${ev.value} = $dtu.truncDate($dateVal, $form);
+ ${ev.value} = $dtu.$truncFuncStr
}
"""
})
}
}
}
+/**
+ * Returns date truncated to the unit specified by the format.
+ */
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+ usage = """
+ _FUNC_(date, fmt) - Returns `date` with the time portion of the day
truncated to the unit specified by the format model `fmt`.
+ `fmt` should be one of ["year", "yyyy", "yy", "mon", "month", "mm"]
+ """,
+ examples = """
+ Examples:
+ > SELECT _FUNC_('2009-02-12', 'MM');
+ 2009-02-01
+ > SELECT _FUNC_('2015-10-27', 'YEAR');
+ 2015-01-01
+ """,
+ since = "1.5.0")
+// scalastyle:on line.size.limit
+case class TruncDate(date: Expression, format: Expression)
+ extends TruncInstant {
+ override def left: Expression = date
+ override def right: Expression = format
+
+ override def inputTypes: Seq[AbstractDataType] = Seq(DateType,
StringType)
+ override def dataType: DataType = DateType
+ override def prettyName: String = "trunc"
+ override val instant = date
+
+ override def eval(input: InternalRow): Any = {
+ evalHelper(input, maxLevel = DateTimeUtils.TRUNC_TO_MONTH) { (d: Any,
level: Int) =>
+ DateTimeUtils.truncDate(d.asInstanceOf[Int], level)
+ }
+ }
+
+ override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ codeGenHelper(ctx, ev, maxLevel = DateTimeUtils.TRUNC_TO_MONTH) {
(date: String, fmt: String) =>
+ s"truncDate($date, $fmt);"
+ }
+ }
+}
+
+/**
+ * Returns timestamp truncated to the unit specified by the format.
+ */
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+ usage = """
+ _FUNC_(fmt, date) - Returns timestamp `ts` truncated to the unit
specified by the format model `fmt`.
--- End diff --
`date` -> `ts`.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
