fix regressions

Author: andygrove

docs/source/contributor-guide/jvm_shuffle.md

@@ -49,9 +49,9 @@ JVM shuffle (`CometColumnarExchange`) is used instead of native shuffle (`CometE
 3. **Unsupported partitioning type**: Native shuffle only supports `HashPartitioning`, `RangePartitioning`,
    and `SinglePartition`. JVM shuffle additionally supports `RoundRobinPartitioning`.
 
-4. **Unsupported partition key types**: For `HashPartitioning` and `RangePartitioning`, native shuffle
-   only supports primitive types as partition keys. Complex types (struct, array, map) cannot be used
-   as partition keys in native shuffle, though they are fully supported as data columns in both implementations.
+4. **Unsupported partition key types**: For `RangePartitioning`, native shuffle only supports primitive
+   types as partition keys. Complex types (struct, array, map) are supported as hash partition keys in
+   native shuffle.
 
 ## Input Handling
 

docs/source/contributor-guide/native_shuffle.md

@@ -55,9 +55,9 @@ Native shuffle (`CometExchange`) is selected when all of the following condition
 
    `RoundRobinPartitioning` requires JVM shuffle.
 
-4. **Supported partition key types**: For `HashPartitioning` and `RangePartitioning`, partition
-   keys must be primitive types. Complex types (struct, array, map) as partition keys require
-   JVM shuffle. Note that complex types are fully supported as data columns in native shuffle.
+4. **Supported partition key types**: For `HashPartitioning`, both primitive and complex types
+   (struct, array, map) are supported as partition keys. For `RangePartitioning`, only primitive
+   types are supported as partition keys.
 
 ## Architecture
 

docs/source/user-guide/latest/tuning.md

@@ -141,8 +141,8 @@ back to Spark for shuffle operations.
 #### Native Shuffle
 
 Comet provides a fully native shuffle implementation, which generally provides the best performance. Native shuffle
-supports `HashPartitioning`, `RangePartitioning` and `SinglePartitioning` but currently only supports primitive type
-partitioning keys. Columns that are not partitioning keys may contain complex types like maps, structs, and arrays.
+supports `HashPartitioning`, `RangePartitioning` and `SinglePartitioning`. Complex types (structs, arrays, and maps)
+are supported as hash partition keys. Range partitioning only supports primitive types as partition keys.
 
 #### Columnar (JVM) Shuffle
 

spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala

@@ -179,7 +179,7 @@ class CometFuzzTestSuite extends CometFuzzTestBase {
     df.createOrReplaceTempView("t1")
     val columns = df.schema.fields.filter(f => isComplexType(f.dataType)).map(_.name)
     for (col <- columns) {
-      // DISTRIBUTE BY is equivalent to df.repartition($col) and uses
+      // DISTRIBUTE BY is equivalent to df.repartition($col)
       val sql = s"SELECT $col FROM t1 DISTRIBUTE BY $col"
       val df = spark.sql(sql)
       df.collect()
@@ -191,13 +191,7 @@ class CometFuzzTestSuite extends CometFuzzTestBase {
           // native_comet does not support reading complex types
           0
         case _ =>
-          CometConf.COMET_SHUFFLE_MODE.get() match {
-            case "jvm" =>
-              1
-            case "native" =>
-              // native shuffle does not support complex types as partitioning keys
-              0
-          }
+          1
       }
       assert(cometShuffleExchanges.length == expectedNumCometShuffles)
     }

spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala

@@ -19,6 +19,9 @@
 
 package org.apache.comet
 
+import org.scalactic.source.Position
+import org.scalatest.Tag
+
 import org.apache.spark.sql.CometTestBase
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
 
@@ -30,6 +33,15 @@ import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
  */
 class CometHashExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
 
+  override protected def test(testName: String, testTags: Tag*)(testFun: => Any)(implicit
+      pos: Position): Unit = {
+    super.test(testName, testTags: _*) {
+      withSQLConf(CometConf.COMET_NATIVE_SCAN_IMPL.key -> CometConf.SCAN_AUTO) {
+        testFun
+      }
+    }
+  }
+
   test("hash - boolean") {
     withTable("t") {
       sql("CREATE TABLE t(c BOOLEAN) USING parquet")