Skip to content
This repository has been archived by the owner on Sep 18, 2023. It is now read-only.

[NSE-198] support the month() and dayofmonth() functions with DateType #199

Merged
merged 1 commit into from
Mar 29, 2021
Merged
Changes from all commits
Commits
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
Original file line number Diff line number Diff line change
Expand Up @@ -113,8 +113,8 @@ class ColumnarIsNull(child: Expression, original: Expression)
}
}

class ColumnarYear(child: Expression, original: Expression)
extends Year(child: Expression)
class ColumnarMonth(child: Expression, original: Expression)
extends Month(child: Expression)
with ColumnarExpression
with Logging {

Expand All @@ -124,7 +124,7 @@ class ColumnarYear(child: Expression, original: Expression)
val supportedTypes = List(LongType, StringType, DateType)
if (supportedTypes.indexOf(child.dataType) == -1) {
throw new UnsupportedOperationException(
s"${child.dataType} is not supported in ColumnarYear.")
s"${child.dataType} is not supported in ColumnarMonth.")
}
}

Expand All @@ -139,6 +139,80 @@ class ColumnarYear(child: Expression, original: Expression)
"castDATE",
Lists.newArrayList(child_node),
new ArrowType.Date(DateUnit.MILLISECOND))
val funcNode =
TreeBuilder.makeFunction(
"extractMonth",
Lists.newArrayList(cast_func),
new ArrowType.Int(64, true))
val castNode =
TreeBuilder.makeFunction("castINT", Lists.newArrayList(funcNode), resultType)
(castNode, resultType)
}
}

class ColumnarDayOfMonth(child: Expression, original: Expression)
extends DayOfMonth(child: Expression)
with ColumnarExpression
with Logging {

buildCheck()

def buildCheck(): Unit = {
val supportedTypes = List(LongType, StringType, DateType)
if (supportedTypes.indexOf(child.dataType) == -1) {
throw new UnsupportedOperationException(
s"${child.dataType} is not supported in ColumnarDayOfMonth.")
}
}

override def doColumnarCodeGen(args: java.lang.Object): (TreeNode, ArrowType) = {
val (child_node, childType): (TreeNode, ArrowType) =
child.asInstanceOf[ColumnarExpression].doColumnarCodeGen(args)

val resultType = new ArrowType.Int(32, true)
//FIXME(): requires utf8()/int64() as input
val cast_func =
TreeBuilder.makeFunction(
"castDATE",
Lists.newArrayList(child_node),
new ArrowType.Date(DateUnit.MILLISECOND))
val funcNode =
TreeBuilder.makeFunction(
"extractDay",
Lists.newArrayList(cast_func),
new ArrowType.Int(64, true))
val castNode =
TreeBuilder.makeFunction("castINT", Lists.newArrayList(funcNode), resultType)
(castNode, resultType)
}
}

class ColumnarYear(child: Expression, original: Expression)
extends Year(child: Expression)
with ColumnarExpression
with Logging {

buildCheck()

def buildCheck(): Unit = {
val supportedTypes = List(LongType, StringType, DateType)
if (supportedTypes.indexOf(child.dataType) == -1) {
throw new UnsupportedOperationException(
s"${child.dataType} is not supported in ColumnarYear.")
}
}

override def doColumnarCodeGen(args: java.lang.Object): (TreeNode, ArrowType) = {
val (child_node, childType): (TreeNode, ArrowType) =
child.asInstanceOf[ColumnarExpression].doColumnarCodeGen(args)

val resultType = new ArrowType.Int(32, true)
//FIXME(): requires utf8()/int64() as input
val cast_func =
TreeBuilder.makeFunction(
"castDATE",
Lists.newArrayList(child_node),
new ArrowType.Date(DateUnit.MILLISECOND))
val funcNode =
TreeBuilder.makeFunction(
"extractYear",
Expand Down Expand Up @@ -558,6 +632,10 @@ object ColumnarUnaryOperator {
new ColumnarIsNotNull(child, i)
case y: Year =>
new ColumnarYear(child, y)
case m: Month =>
new ColumnarMonth(child, m)
case d: DayOfMonth =>
new ColumnarDayOfMonth(child, d)
case n: Not =>
new ColumnarNot(child, n)
case a: Abs =>
Expand Down