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

oap-project#199)
zhouyuan · Mar 29, 2021 · d45cc87 · d45cc87
1 parent caf1546
commit d45cc87
Showing 1 changed file with 81 additions and 3 deletions.
diff --git a/native-sql-engine/core/src/main/scala/com/intel/oap/expression/ColumnarUnaryOperator.scala b/native-sql-engine/core/src/main/scala/com/intel/oap/expression/ColumnarUnaryOperator.scala
@@ -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 {
 
@@ -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.")
     }
   }
 
@@ -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",
@@ -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 =>