[VL] Support date_trunc function

apache · Oct 19, 2024 · 738b8f2 · 738b8f2
1 parent a5c1d26
commit 738b8f2
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 7 deletions.
diff --git a/docs/velox-backend-support-progress.md b/docs/velox-backend-support-progress.md
@@ -344,7 +344,7 @@ Gluten supports 199 functions. (Drag to right to see all data types)
 | date_from_unix_date          |                        |                       | S      |                          |         |      |       |     |      |       |        |      |           |        |         |      |        |          |       |     |        |     |
 | date_part                    |                        |                       |        |                          |         |      |       |     |      |       |        |      |           |        |         |      |        |          |       |     |        |     |
 | date_sub                     |                        |                       | S      |                          |         |      |       |     |      |       |        |      |           |        |         |      |        |          |       |     |        |     |
-| date_trunc                   | date_trunc             |                       |        |                          |         |      |       |     |      |       |        |      |           |        |         |      |        |          |       |     |        |     |
+| date_trunc                   | date_trunc             | date_trunc            | S      |                          |         |      |       |     |      |       |        |      |           |        |         |      |        |          |       |     |        |     |
 | datediff                     | date_diff              |                       | S      |                          |         |      |       |     |      |       |        | S    | S         |        |         |      |        |          |       |     |        |     |
 | day                          | day                    |                       | S      |                          |         |      |       |     |      |       |        | S    | S         |        |         |      |        |          |       |     |        |     |
 | dayofmonth                   | day_of_month           |                       | S      |                          |         |      |       |     |      |       |        | S    | S         |        |         |      |        |          |       |     |        |     |

diff --git a/ep/build-velox/src/get_velox.sh b/ep/build-velox/src/get_velox.sh
@@ -16,8 +16,8 @@
 
 set -exu
 
-VELOX_REPO=https://github.com/oap-project/velox.git
-VELOX_BRANCH=2024_10_19
+VELOX_REPO=https://github.com/zml1206/velox.git
+VELOX_BRANCH=date_trunc_intel
 VELOX_HOME=""
 
 OS=`uname -s`
@@ -101,7 +101,7 @@ function process_setup_centos9 {
   ensure_pattern_matched 'dnf_install' scripts/setup-centos9.sh
   sed -i 's/dnf_install ninja-build cmake curl ccache gcc-toolset-12 git/dnf_install ninja-build cmake curl ccache gcc-toolset-12/' scripts/setup-centos9.sh
   sed -i '/^.*dnf_install autoconf/a\  dnf_install libxml2-devel libgsasl-devel libuuid-devel' scripts/setup-centos9.sh
-  
+
   ensure_pattern_matched 'install_gflags' scripts/setup-centos9.sh
   sed -i '/^function install_gflags.*/i function install_openssl {\n  wget_and_untar https://github.com/openssl/openssl/releases/download/openssl-3.2.2/openssl-3.2.2.tar.gz openssl \n ( cd ${DEPENDENCY_DIR}/openssl \n  ./config no-shared && make depend && make && sudo make install ) \n}\n'     scripts/setup-centos9.sh
 

diff --git a/...ubstrait/src/main/scala/org/apache/gluten/expression/DateTimeExpressionsTransformer.scala b/...ubstrait/src/main/scala/org/apache/gluten/expression/DateTimeExpressionsTransformer.scala
@@ -17,8 +17,12 @@
 package org.apache.gluten.expression
 
 import org.apache.gluten.exception.GlutenNotSupportException
+import org.apache.gluten.substrait.expression.{ExpressionBuilder, ExpressionNode}
 
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.unsafe.types.UTF8String
+
+import java.util.Locale
 
 /** The extract trait for 'GetDateField' from Date */
 case class ExtractDateTransformer(
@@ -43,9 +47,45 @@ case class TruncTimestampTransformer(
     timestamp: ExpressionTransformer,
     original: TruncTimestamp)
   extends ExpressionTransformer {
-  override def children: Seq[ExpressionTransformer] = {
-    val timeZoneId = original.timeZoneId.map(timeZoneId => LiteralTransformer(timeZoneId))
-    Seq(format, timestamp) ++ timeZoneId
+  override def children: Seq[ExpressionTransformer] = Seq(format, timestamp)
+
+  override def doTransform(args: java.lang.Object): ExpressionNode = {
+    if (!original.format.foldable) {
+      throw new GlutenNotSupportException(s"The format ${original.format} must be constant string.")
+    }
+    val formatStr = original.format.eval().asInstanceOf[UTF8String]
+    if (formatStr == null) {
+      throw new GlutenNotSupportException("The format is null.")
+    }
+    val newFormatStr = formatStr.toString.toLowerCase(Locale.ROOT) match {
+      case "second" => "second"
+      case "minute" => "minute"
+      case "hour" => "hour"
+      case "day" | "dd" => "day"
+      case "week" => "week"
+      case "mon" | "month" | "mm" => "month"
+      case "quarter" => "quarter"
+      case "year" | "yyyy" | "yy" => "year"
+      // Can not support now.
+      // case "microsecond" => "microsecond"
+      // case "millisecond" => "millisecond"
+      case _ => throw new GlutenNotSupportException(s"The format $formatStr is invalidate.")
+    }
+
+    val functionMap = args.asInstanceOf[java.util.HashMap[String, java.lang.Long]]
+    val dataTypes = Seq(original.format.dataType, original.timestamp.dataType)
+    val functionId = ExpressionBuilder.newScalarFunction(
+      functionMap,
+      ConverterUtils.makeFuncName(substraitExprName, dataTypes))
+
+    val expressionNodes = new java.util.ArrayList[ExpressionNode]()
+    val timestampNode = timestamp.doTransform(args)
+    val lowerFormatNode = ExpressionBuilder.makeStringLiteral(newFormatStr)
+    expressionNodes.add(lowerFormatNode)
+    expressionNodes.add(timestampNode)
+
+    val typeNode = ConverterUtils.getTypeNode(original.dataType, original.nullable)
+    ExpressionBuilder.makeScalarFunction(functionId, expressionNodes, typeNode)
   }
 }