fix doy add dow

oap-project · Jun 25, 2021 · 48baa17 · 48baa17
1 parent a82b3e1
commit 48baa17
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 14 deletions.
diff --git a/...sql-engine/core/src/main/scala/com/intel/oap/expression/ColumnarDateTimeExpressions.scala b/...sql-engine/core/src/main/scala/com/intel/oap/expression/ColumnarDateTimeExpressions.scala
@@ -31,6 +31,7 @@ import org.apache.spark.sql.catalyst.expressions.CurrentDate
 import org.apache.spark.sql.catalyst.expressions.CurrentTimestamp
 import org.apache.spark.sql.catalyst.expressions.DateDiff
 import org.apache.spark.sql.catalyst.expressions.DayOfMonth
+import org.apache.spark.sql.catalyst.expressions.DayOfWeek
 import org.apache.spark.sql.catalyst.expressions.DayOfYear
 import org.apache.spark.sql.catalyst.expressions.Expression
 import org.apache.spark.sql.catalyst.expressions.Hour
@@ -151,6 +152,19 @@ object ColumnarDateTimeExpressions {
     }
   }
 
+  class ColumnarDayOfWeek(child: Expression) extends DayOfWeek(child) with
+      ColumnarExpression {
+    override def doColumnarCodeGen(args: Object): (TreeNode, ArrowType) = {
+      val (childNodeUtc, childTypeUtc) =
+        child.asInstanceOf[ColumnarExpression].doColumnarCodeGen(args)
+      val (childNode, childType) = ConverterUtils.toGandivaTimestamp(childNodeUtc, childTypeUtc)
+      val outType = ArrowUtils.toArrowType(LongType, null)
+      val funcNode = TreeBuilder.makeFunction(
+        "extractDow", Lists.newArrayList(childNode), outType)
+      ConverterUtils.toInt32(funcNode, outType)
+    }
+  }
+
   class ColumnarMonth(child: Expression) extends Month(child) with
       ColumnarExpression {
     override def doColumnarCodeGen(args: Object): (TreeNode, ArrowType) = {

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
@@ -33,6 +33,7 @@ import org.apache.spark.sql.types._
 import scala.collection.mutable.ListBuffer
 
 import com.intel.oap.expression.ColumnarDateTimeExpressions.ColumnarDayOfMonth
+import com.intel.oap.expression.ColumnarDateTimeExpressions.ColumnarDayOfWeek
 import com.intel.oap.expression.ColumnarDateTimeExpressions.ColumnarDayOfYear
 import com.intel.oap.expression.ColumnarDateTimeExpressions.ColumnarHour
 import com.intel.oap.expression.ColumnarDateTimeExpressions.ColumnarHour
@@ -808,19 +809,25 @@ object ColumnarUnaryOperator {
     case a: MicrosToTimestamp =>
       new ColumnarMicrosToTimestamp(child)
     case other =>
-      if (child.dataType.isInstanceOf[TimestampType]) other match {
-        case a: Hour =>
-          new ColumnarHour(child)
-        case a: Minute =>
-          new ColumnarMinute(child)
-        case a: Second =>
-          new ColumnarSecond(child)
-        case a: DayOfYear =>
-          new ColumnarDayOfYear(child)
-        case other =>
+      child.dataType match {
+        case _: DateType => other match {
+          case a: DayOfYear =>
+            new ColumnarDayOfYear(new ColumnarCast(child, TimestampType, None, null))
+          case a: DayOfWeek =>
+            new ColumnarDayOfWeek(new ColumnarCast(child, TimestampType, None, null))
+        }
+        case _: TimestampType => other match {
+          case a: Hour =>
+            new ColumnarHour(child)
+          case a: Minute =>
+            new ColumnarMinute(child)
+          case a: Second =>
+            new ColumnarSecond(child)
+          case other =>
+            throw new UnsupportedOperationException(s"not currently supported: $other.")
+        }
+        case _ =>
           throw new UnsupportedOperationException(s"not currently supported: $other.")
-      } else {
-        throw new UnsupportedOperationException(s"not currently supported: $other.")
       }
   }
 }
diff --git a/native-sql-engine/core/src/test/scala/com/intel/oap/misc/DateTimeSuite.scala b/native-sql-engine/core/src/test/scala/com/intel/oap/misc/DateTimeSuite.scala
@@ -468,8 +468,29 @@ class DateTimeSuite extends QueryTest with SharedSparkSession {
     }
   }
 
-  // FIXME this is falling back. Requiring date input support
-  ignore("datetime function - dayofyear") {
+  test("datetime function - dayofweek") {
+    withTempView("timestamps") {
+      val dates = Seq(0L, 1L, 2L, 3L, 4L, 2L, 3L)
+          .map(i => Tuple1(new Timestamp(i))).toDF("time")
+      dates.createOrReplaceTempView("timestamps")
+      val frame = sql("SELECT DAYOFWEEK(time) FROM timestamps")
+      frame.explain()
+      frame.show()
+      assert(frame.queryExecution.executedPlan.find(p => p
+          .isInstanceOf[ColumnarConditionProjectExec]).isDefined)
+      checkAnswer(
+        frame,
+        Seq(Row(Integer.valueOf(4)),
+          Row(Integer.valueOf(4)),
+          Row(Integer.valueOf(4)),
+          Row(Integer.valueOf(4)),
+          Row(Integer.valueOf(4)),
+          Row(Integer.valueOf(4)),
+          Row(Integer.valueOf(4))))
+    }
+  }
+
+  test("datetime function - dayofyear") {
     withTempView("timestamps") {
       val dates = Seq(0L, 1L, 2L, 3L, 4L, 2L, 3L)
           .map(i => Tuple1(new Timestamp(i))).toDF("time")