revert the changes for TimestampNTZType and PartitionIdPassthrough

apache · Jun 17, 2024 · 99f2d5c · 99f2d5c
1 parent 9cc0b3e
commit 99f2d5c
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 1 deletion.
diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -61,8 +61,9 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
   def supportedDataType(dt: DataType): Boolean = dt match {
     case _: ByteType | _: ShortType | _: IntegerType | _: LongType | _: FloatType |
         _: DoubleType | _: StringType | _: BinaryType | _: TimestampType | _: DecimalType |
-        _: DateType | _: BooleanType | _: NullType | _: TimestampNTZType =>
+        _: DateType | _: BooleanType | _: NullType =>
       true
+    case dt if dt.typeName == "timestamp_ntz" => true
     case dt =>
       emitWarning(s"unsupported Spark data type: $dt")
       false

diff --git a/...rc/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/...rc/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
@@ -581,3 +581,18 @@ class CometShuffleWriteProcessor(
     }
   }
 }
+
+/**
+ * Copied from Spark `PartitionIdPassthrough` as it is private in Spark 3.2.
+ */
+private[spark] class PartitionIdPassthrough(override val numPartitions: Int) extends Partitioner {
+  override def getPartition(key: Any): Int = key.asInstanceOf[Int]
+}
+
+/**
+ * Copied from Spark `ConstantPartitioner` as it doesn't exist in Spark 3.2.
+ */
+private[spark] class ConstantPartitioner extends Partitioner {
+  override def numPartitions: Int = 1
+  override def getPartition(key: Any): Int = 0
+}