Skip to content

Commit

Permalink
[SPARK-23905][SQL] Add UDF weekday
Browse files Browse the repository at this point in the history
## What changes were proposed in this pull request?

Add UDF weekday

## How was this patch tested?

A new test

Author: yucai <yyu1@ebay.com>

Closes #21009 from yucai/SPARK-23905.
  • Loading branch information
yucai authored and gatorsmile committed Apr 13, 2018
1 parent 4b07036 commit 0323e61
Show file tree
Hide file tree
Showing 5 changed files with 67 additions and 11 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -395,6 +395,7 @@ object FunctionRegistry {
expression[TruncTimestamp]("date_trunc"),
expression[UnixTimestamp]("unix_timestamp"),
expression[DayOfWeek]("dayofweek"),
expression[WeekDay]("weekday"),
expression[WeekOfYear]("weekofyear"),
expression[Year]("year"),
expression[TimeWindow]("window"),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -426,36 +426,71 @@ case class DayOfMonth(child: Expression) extends UnaryExpression with ImplicitCa
""",
since = "2.3.0")
// scalastyle:on line.size.limit
case class DayOfWeek(child: Expression) extends UnaryExpression with ImplicitCastInputTypes {
case class DayOfWeek(child: Expression) extends DayWeek {

override def inputTypes: Seq[AbstractDataType] = Seq(DateType)

override def dataType: DataType = IntegerType
override protected def nullSafeEval(date: Any): Any = {
cal.setTimeInMillis(date.asInstanceOf[Int] * 1000L * 3600L * 24L)
cal.get(Calendar.DAY_OF_WEEK)
}

@transient private lazy val c = {
Calendar.getInstance(DateTimeUtils.getTimeZone("UTC"))
override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
nullSafeCodeGen(ctx, ev, time => {
val cal = classOf[Calendar].getName
val dtu = DateTimeUtils.getClass.getName.stripSuffix("$")
val c = "calDayOfWeek"
ctx.addImmutableStateIfNotExists(cal, c,
v => s"""$v = $cal.getInstance($dtu.getTimeZone("UTC"));""")
s"""
$c.setTimeInMillis($time * 1000L * 3600L * 24L);
${ev.value} = $c.get($cal.DAY_OF_WEEK);
"""
})
}
}

// scalastyle:off line.size.limit
@ExpressionDescription(
usage = "_FUNC_(date) - Returns the day of the week for date/timestamp (0 = Monday, 1 = Tuesday, ..., 6 = Sunday).",
examples = """
Examples:
> SELECT _FUNC_('2009-07-30');
3
""",
since = "2.4.0")
// scalastyle:on line.size.limit
case class WeekDay(child: Expression) extends DayWeek {

override protected def nullSafeEval(date: Any): Any = {
c.setTimeInMillis(date.asInstanceOf[Int] * 1000L * 3600L * 24L)
c.get(Calendar.DAY_OF_WEEK)
cal.setTimeInMillis(date.asInstanceOf[Int] * 1000L * 3600L * 24L)
(cal.get(Calendar.DAY_OF_WEEK) + 5 ) % 7
}

override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
nullSafeCodeGen(ctx, ev, time => {
val cal = classOf[Calendar].getName
val dtu = DateTimeUtils.getClass.getName.stripSuffix("$")
val c = "calDayOfWeek"
val c = "calWeekDay"
ctx.addImmutableStateIfNotExists(cal, c,
v => s"""$v = $cal.getInstance($dtu.getTimeZone("UTC"));""")
s"""
$c.setTimeInMillis($time * 1000L * 3600L * 24L);
${ev.value} = $c.get($cal.DAY_OF_WEEK);
${ev.value} = ($c.get($cal.DAY_OF_WEEK) + 5) % 7;
"""
})
}
}

abstract class DayWeek extends UnaryExpression with ImplicitCastInputTypes {

override def inputTypes: Seq[AbstractDataType] = Seq(DateType)

override def dataType: DataType = IntegerType

@transient protected lazy val cal: Calendar = {
Calendar.getInstance(DateTimeUtils.getTimeZone("UTC"))
}
}

// scalastyle:off line.size.limit
@ExpressionDescription(
usage = "_FUNC_(date) - Returns the week of the year of the given date. A week is considered to start on a Monday and week 1 is the first week with >3 days.",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -211,6 +211,17 @@ class DateExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
checkConsistencyBetweenInterpretedAndCodegen(DayOfWeek, DateType)
}

test("WeekDay") {
checkEvaluation(WeekDay(Literal.create(null, DateType)), null)
checkEvaluation(WeekDay(Literal(d)), 2)
checkEvaluation(WeekDay(Cast(Literal(sdfDate.format(d)), DateType, gmtId)), 2)
checkEvaluation(WeekDay(Cast(Literal(ts), DateType, gmtId)), 4)
checkEvaluation(WeekDay(Cast(Literal("2011-05-06"), DateType, gmtId)), 4)
checkEvaluation(WeekDay(Literal(new Date(sdf.parse("2017-05-27 13:10:15").getTime))), 5)
checkEvaluation(WeekDay(Literal(new Date(sdf.parse("1582-10-15 13:10:15").getTime))), 4)
checkConsistencyBetweenInterpretedAndCodegen(WeekDay, DateType)
}

test("WeekOfYear") {
checkEvaluation(WeekOfYear(Literal.create(null, DateType)), null)
checkEvaluation(WeekOfYear(Literal(d)), 15)
Expand Down
2 changes: 2 additions & 0 deletions sql/core/src/test/resources/sql-tests/inputs/datetime.sql
Original file line number Diff line number Diff line change
Expand Up @@ -25,3 +25,5 @@ create temporary view ttf2 as select * from values
select current_date = current_date(), current_timestamp = current_timestamp(), a, b from ttf2;

select a, b from ttf2 order by a, current_date;

select weekday('2007-02-03'), weekday('2009-07-30'), weekday('2017-05-27'), weekday(null), weekday('1582-10-15 13:10:15');
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
-- Automatically generated by SQLQueryTestSuite
-- Number of queries: 9
-- Number of queries: 10


-- !query 0
Expand Down Expand Up @@ -81,3 +81,10 @@ struct<a:int,b:int>
-- !query 8 output
1 2
2 3

-- !query 9
select weekday('2007-02-03'), weekday('2009-07-30'), weekday('2017-05-27'), weekday(null), weekday('1582-10-15 13:10:15')
-- !query 3 schema
struct<weekday(CAST(2007-02-03 AS DATE)):int,weekday(CAST(2009-07-30 AS DATE)):int,weekday(CAST(2017-05-27 AS DATE)):int,weekday(CAST(NULL AS DATE)):int,weekday(CAST(1582-10-15 13:10:15 AS DATE)):int>
-- !query 3 output
5 3 5 NULL 4

0 comments on commit 0323e61

Please sign in to comment.