fix: Unsupported types for SinglePartition should fallback to Spark (a…

…pache#765)
viirya · Aug 2, 2024 · 39e030b · 39e030b
1 parent 65f22a2
commit 39e030b
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -2931,8 +2931,8 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
           msg = s"unsupported Spark partitioning expressions: $expressions"
         }
         supported
-      case SinglePartition => true
-      case RoundRobinPartitioning(_) => true
+      case SinglePartition => inputs.forall(attr => supportedDataType(attr.dataType))
+      case RoundRobinPartitioning(_) => inputs.forall(attr => supportedDataType(attr.dataType))
       case RangePartitioning(orderings, _) =>
         val supported =
           orderings.map(QueryPlanSerde.exprToProto(_, inputs)).forall(_.isDefined) &&
@@ -2980,7 +2980,7 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
           msg = s"unsupported Spark partitioning expressions: $expressions"
         }
         supported
-      case SinglePartition => true
+      case SinglePartition => inputs.forall(attr => supportedDataType(attr.dataType))
       case _ =>
         msg = s"unsupported Spark partitioning: ${partitioning.getClass.getName}"
         false

diff --git a/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala
@@ -66,6 +66,19 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar
 
   setupTestData()
 
+  test("Unsupported types for SinglePartition should fallback to Spark") {
+    checkSparkAnswer(spark.sql("""
+                  |SELECT
+                  |  AVG(null),
+                  |  COUNT(null),
+                  |  FIRST(null),
+                  |  LAST(null),
+                  |  MAX(null),
+                  |  MIN(null),
+                  |  SUM(null)
+        """.stripMargin))
+  }
+
   test("Disable Comet shuffle with AQE coalesce partitions enabled") {
     Seq(true, false).foreach { coalescePartitionsEnabled =>
       withSQLConf(