feat: Enable complex types for columnar shuffle (#2573)

Author: mbutrovich

spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala

@@ -877,9 +877,9 @@ case class CometExecRule(session: SparkSession) extends Rule[SparkPlan] {
         true
       case RangePartitioning(orderings, _) =>
         if (!CometConf.COMET_EXEC_SHUFFLE_WITH_RANGE_PARTITIONING_ENABLED.get(conf)) {
-          // do not encourage the users to enable the config because we know that
-          // the experimental implementation is not correct yet
-          withInfo(s, "Range partitioning is not supported by native shuffle")
+          withInfo(
+            s,
+            s"${CometConf.COMET_EXEC_SHUFFLE_WITH_RANGE_PARTITIONING_ENABLED.key} is disabled")
           return false
         }
         var supported = true
@@ -924,22 +924,10 @@ case class CometExecRule(session: SparkSession) extends Rule[SparkPlan] {
           _: TimestampNTZType | _: DecimalType | _: DateType =>
         true
       case StructType(fields) =>
-        fields.forall(f => supportedSerializableDataType(f.dataType)) &&
+        fields.nonEmpty && fields.forall(f => supportedSerializableDataType(f.dataType)) &&
         // Java Arrow stream reader cannot work on duplicate field name
         fields.map(f => f.name).distinct.length == fields.length &&
         fields.nonEmpty
-
-      // TODO add support for nested complex types
-      // https://github.com/apache/datafusion-comet/issues/2199
-      case ArrayType(ArrayType(_, _), _) => false
-      case ArrayType(MapType(_, _, _), _) => false
-      case MapType(MapType(_, _, _), _, _) => false
-      case MapType(_, MapType(_, _, _), _) => false
-      case MapType(StructType(_), _, _) => false
-      case MapType(_, StructType(_), _) => false
-      case MapType(ArrayType(_, _), _, _) => false
-      case MapType(_, ArrayType(_, _), _) => false
-
       case ArrayType(elementType, _) =>
         supportedSerializableDataType(elementType)
       case MapType(keyType, valueType, _) =>

spark/src/test/scala/org/apache/comet/exec/CometColumnarShuffleSuite.scala

@@ -19,8 +19,7 @@
 
 package org.apache.comet.exec
 
-import java.nio.file.Files
-import java.nio.file.Paths
+import java.nio.file.{Files, Paths}
 
 import scala.reflect.runtime.universe._
 import scala.util.Random
@@ -168,8 +167,13 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar
                 .repartition(numPartitions, $"_1", $"_2")
                 .sortWithinPartitions($"_2")
 
-              // Array map key array element fallback to Spark shuffle for now
-              checkShuffleAnswer(df, 0)
+              if (isSpark40Plus) {
+                // https://github.com/apache/datafusion-comet/issues/1941
+                // Spark 4.0 introduces a mapsort which falls back
+                checkShuffleAnswer(df, 0)
+              } else {
+                checkShuffleAnswer(df, 1)
+              }
             }
 
             withParquetTable((0 until 50).map(i => (Map(i -> Seq(i, i + 1)), i + 1)), "tbl") {
@@ -178,8 +182,13 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar
                 .repartition(numPartitions, $"_1", $"_2")
                 .sortWithinPartitions($"_2")
 
-              // Array map value array element fallback to Spark shuffle for now
-              checkShuffleAnswer(df, 0)
+              if (isSpark40Plus) {
+                // https://github.com/apache/datafusion-comet/issues/1941
+                // Spark 4.0 introduces a mapsort which falls back
+                checkShuffleAnswer(df, 0)
+              } else {
+                checkShuffleAnswer(df, 1)
+              }
             }
 
             withParquetTable((0 until 50).map(i => (Map((i, i.toString) -> i), i + 1)), "tbl") {
@@ -188,8 +197,13 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar
                 .repartition(numPartitions, $"_1", $"_2")
                 .sortWithinPartitions($"_2")
 
-              // Struct map key array element fallback to Spark shuffle for now
-              checkShuffleAnswer(df, 0)
+              if (isSpark40Plus) {
+                // https://github.com/apache/datafusion-comet/issues/1941
+                // Spark 4.0 introduces a mapsort which falls back
+                checkShuffleAnswer(df, 0)
+              } else {
+                checkShuffleAnswer(df, 1)
+              }
             }
 
             withParquetTable((0 until 50).map(i => (Map(i -> (i, i.toString)), i + 1)), "tbl") {
@@ -198,8 +212,13 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar
                 .repartition(numPartitions, $"_1", $"_2")
                 .sortWithinPartitions($"_2")
 
-              // Struct map value array element fallback to Spark shuffle for now
-              checkShuffleAnswer(df, 0)
+              if (isSpark40Plus) {
+                // https://github.com/apache/datafusion-comet/issues/1941
+                // Spark 4.0 introduces a mapsort which falls back
+                checkShuffleAnswer(df, 0)
+              } else {
+                checkShuffleAnswer(df, 1)
+              }
             }
           }
         }
@@ -222,8 +241,13 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar
                 .repartition(numPartitions, $"_1", $"_2")
                 .sortWithinPartitions($"_2")
 
-              // Map array element fallback to Spark shuffle for now
-              checkShuffleAnswer(df, 0)
+              if (isSpark40Plus) {
+                // https://github.com/apache/datafusion-comet/issues/1941
+                // Spark 4.0 introduces a mapsort which falls back
+                checkShuffleAnswer(df, 0)
+              } else {
+                checkShuffleAnswer(df, 1)
+              }
             }
           }
         }
@@ -469,8 +493,7 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar
                 .repartition(numPartitions, $"_1", $"_2")
                 .sortWithinPartitions($"_1")
 
-              // Nested array fallback to Spark shuffle for now
-              checkShuffleAnswer(df, 0)
+              checkShuffleAnswer(df, 1)
             }
           }
         }
@@ -983,6 +1006,7 @@ class DisableAQECometAsyncShuffleSuite extends CometColumnarShuffleSuite {
 }
 
 class CometShuffleEncryptionSuite extends CometTestBase {
+
   import testImplicits._
 
   override protected def sparkConf: SparkConf = {
@@ -1034,6 +1058,7 @@ class CometShuffleManagerSuite extends CometTestBase {
         shuffleWriterProcessor = null,
         partitioner = new Partitioner {
           override def numPartitions: Int = 50
+
           override def getPartition(key: Any): Int = key.asInstanceOf[Int]
         },
         decodeTime = null)