build: Enable spark-4.0 Spark tests

apache · Jun 7, 2024 · c1d90aa · c1d90aa
1 parent 5988fff
commit c1d90aa
Show file tree

Hide file tree

Showing 6 changed files with 31 additions and 4 deletions.
diff --git a/.github/workflows/spark_sql_test_ansi.yml b/.github/workflows/spark_sql_test_ansi.yml
@@ -76,7 +76,7 @@ jobs:
         run: |
           cd apache-spark
           rm -rf /root/.m2/repository/org/apache/parquet # somehow parquet cache requires cleanups
-          ENABLE_COMET=true ENABLE_COMET_ANSI_MODE=true build/sbt ${{ matrix.module.args1 }} "${{ matrix.module.args2 }}"
+          RUST_BACKTRACE=1 ENABLE_COMET=true ENABLE_COMET_ANSI_MODE=true build/sbt ${{ matrix.module.args1 }} "${{ matrix.module.args2 }}"
         env:
           LC_ALL: "C.UTF-8"
 
diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -66,7 +66,7 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
         _: DateType | _: BooleanType | _: NullType =>
       true
     // `TimestampNTZType` is private in Spark 3.2.
-    case dt if dt.typeName == "timestamp_ntz" => true
+    case dt if isTimestampNTZType(dt) => true
     case dt =>
       emitWarning(s"unsupported Spark data type: $dt")
       false
@@ -90,7 +90,7 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
       case _: BinaryType => 8
       case _: TimestampType => 9
       case _: DecimalType => 10
-      case dt if dt.typeName == "timestamp_ntz" => 11
+      case dt if isTimestampNTZType(dt) => 11
       case _: DateType => 12
       case _: NullType => 13
       case _: ArrayType => 14
@@ -1047,6 +1047,8 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
                   com.google.protobuf.ByteString.copyFrom(value.asInstanceOf[Array[Byte]])
                 exprBuilder.setBytesVal(byteStr)
               case _: DateType => exprBuilder.setIntVal(value.asInstanceOf[Int])
+              case dt if isTimestampNTZType(dt) =>
+                exprBuilder.setLongVal(value.asInstanceOf[Long])
               case dt =>
                 logWarning(s"Unexpected date type '$dt' for literal value '$value'")
             }
@@ -2258,7 +2260,7 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
         _: DoubleType | _: StringType | _: DateType | _: DecimalType | _: BooleanType =>
       true
     // `TimestampNTZType` is private in Spark 3.2/3.3.
-    case dt if dt.typeName == "timestamp_ntz" => true
+    case dt if isTimestampNTZType(dt) => true
     case _ => false
   }
 

diff --git a/spark/src/main/spark-3.2/org/apache/comet/shims/CometExprShim.scala b/spark/src/main/spark-3.2/org/apache/comet/shims/CometExprShim.scala
@@ -19,6 +19,7 @@
 package org.apache.comet.shims
 
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.types.DataType
 
 /**
  * `CometExprShim` acts as a shim for for parsing expressions from different Spark versions.
@@ -30,4 +31,10 @@ trait CometExprShim {
     def unhexSerde(unhex: Unhex): (Expression, Expression) = {
         (unhex.child, Literal(false))
     }
+
+    protected def isTimestampNTZType(dt: DataType): Boolean = dt match {
+        // `TimestampNTZType` is private in Spark 3.2.
+        case dt if dt.typeName == "timestamp_ntz" => true
+        case _ => false
+    }
 }
diff --git a/spark/src/main/spark-3.3/org/apache/comet/shims/CometExprShim.scala b/spark/src/main/spark-3.3/org/apache/comet/shims/CometExprShim.scala
@@ -19,6 +19,7 @@
 package org.apache.comet.shims
 
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.types.{DataType, TimestampNTZType}
 
 /**
  * `CometExprShim` acts as a shim for for parsing expressions from different Spark versions.
@@ -30,4 +31,9 @@ trait CometExprShim {
     def unhexSerde(unhex: Unhex): (Expression, Expression) = {
         (unhex.child, Literal(false))
     }
+
+    protected def isTimestampNTZType(dt: DataType): Boolean = dt match {
+        case _: TimestampNTZType => true
+        case _ => false
+    }
 }
diff --git a/spark/src/main/spark-3.4/org/apache/comet/shims/CometExprShim.scala b/spark/src/main/spark-3.4/org/apache/comet/shims/CometExprShim.scala
@@ -19,6 +19,7 @@
 package org.apache.comet.shims
 
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.types.{DataType, TimestampNTZType}
 
 /**
  * `CometExprShim` acts as a shim for for parsing expressions from different Spark versions.
@@ -30,4 +31,9 @@ trait CometExprShim {
     def unhexSerde(unhex: Unhex): (Expression, Expression) = {
         (unhex.child, Literal(unhex.failOnError))
     }
+
+    protected def isTimestampNTZType(dt: DataType): Boolean = dt match {
+        case _: TimestampNTZType => true
+        case _ => false
+    }
 }
diff --git a/spark/src/main/spark-4.0/org/apache/comet/shims/CometExprShim.scala b/spark/src/main/spark-4.0/org/apache/comet/shims/CometExprShim.scala
@@ -19,6 +19,7 @@
 package org.apache.comet.shims
 
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.types.{DataType, TimestampNTZType}
 
 /**
  * `CometExprShim` acts as a shim for for parsing expressions from different Spark versions.
@@ -30,4 +31,9 @@ trait CometExprShim {
     protected def unhexSerde(unhex: Unhex): (Expression, Expression) = {
         (unhex.child, Literal(unhex.failOnError))
     }
+
+    protected def isTimestampNTZType(dt: DataType): Boolean = dt match {
+        case _: TimestampNTZType => true
+        case _ => false
+    }
 }