feat: route MapSort fallback through the codegen dispatcher

docs/source/user-guide/latest/compatibility/expressions/_category_template/map.md

@@ -23,12 +23,15 @@ under the License.
 
 Spark 4.0 inserts `MapSort` to normalize map values when they appear in shuffle hash partitioning
 keys, in `try_element_at`, and in other contexts where map ordering must be deterministic. Comet
-runs `MapSort` natively, so map shuffle and group-by-on-map stay on Comet under Spark 4.0.
+runs `MapSort` natively for supported scalar key types. Other orderable key types use Spark's own
+generated code through the JVM codegen dispatcher, so the enclosing operator stays in the Comet
+pipeline.
 
-When `spark.comet.exec.strictFloatingPoint=true`, `MapSort` falls back to Spark for maps whose
-keys contain `Float` or `Double` (consistent with `SortOrder` and `SortArray`). Arrow's sort uses
-IEEE total ordering for floating-point, which differs from Spark's `Double.compare` semantics for
-`NaN` and `-0.0`.
+When `spark.comet.exec.strictFloatingPoint=true`, maps whose keys contain `Float` or `Double` also
+use the dispatcher (consistent with `SortOrder` and `SortArray`). Arrow's sort uses IEEE total
+ordering for floating-point, which differs from Spark's `Double.compare` semantics for `NaN` and
+`-0.0`. If the dispatcher is disabled or cannot handle an expression, Comet safely falls back to
+Spark.
 
 <!--BEGIN:EXPR_COMPAT[map]-->
 <!--END:EXPR_COMPAT-->

spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala

@@ -452,9 +452,10 @@ object CometShuffleExchangeExec
       case MapType(keyType, valueType, _) if nestedHashPartitioningEnabled =>
         // Map entry order is not semantically meaningful, so two equal maps must hash alike.
         // Spark 4.0+ normalizes a map shuffle key by wrapping it in `mapsort(...)`, which is
-        // gated separately by CometMapSort (scalar map keys only) and, when unsupported, fails
-        // the expression check below. Earlier Spark versions insert no such normalization, so
-        // Comet would hash physical entry order and could route equal maps differently.
+        // gated separately by CometMapSort. Scalar keys use native map_sort; other orderable key
+        // types can use Spark's generated MapSort code through the JVM dispatcher. Earlier Spark
+        // versions insert no such normalization, so Comet would hash physical entry order and
+        // could route equal maps differently.
         isSpark40Plus &&
         supportedHashPartitioningDataType(keyType) &&
         supportedHashPartitioningDataType(valueType)

spark/src/main/spark-4.x/org/apache/comet/serde/CometMapSort.scala

@@ -25,7 +25,8 @@ import org.apache.spark.sql.types.MapType
 import org.apache.comet.CometConf
 import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, scalarFunctionExprToProtoWithReturnType, supportedScalarSortElementType}
 
-object CometMapSort extends CometExpressionSerde[MapSort] {
+// Key types unsupported by the native kernel can still run Spark's generated code in-pipeline.
+object CometMapSort extends CometExpressionSerde[MapSort] with CodegenDispatchFallback {
 
   override def getIncompatibleReasons(): Seq[String] =
     Seq(

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

@@ -33,13 +33,23 @@ import org.apache.comet.vector.CometVector
 trait CometCodegenAssertions {
 
   /** Asserts the dispatcher actually ran during `f`, guarding against silent serde fallback. */
-  protected def assertCodegenRan(f: => Unit): Unit = {
+  protected def assertCodegenRan[T](f: => T): T = {
     CometScalaUDFCodegen.resetStats()
-    f
+    val result = f
     val after = CometScalaUDFCodegen.stats()
     assert(
       after.compileCount + after.cacheHitCount >= 1,
       s"expected codegen dispatcher activity, got $after")
+    result
+  }
+
+  /** Asserts the dispatcher did not run during `f`, guarding a native-path control case. */
+  protected def assertCodegenDidNotRun[T](f: => T): T = {
+    CometScalaUDFCodegen.resetStats()
+    val result = f
+    val after = CometScalaUDFCodegen.stats()
+    assert(after.totalLookups == 0, s"expected no codegen dispatcher activity, got $after")
+    result
   }
 
   /**

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

@@ -22,7 +22,7 @@ package org.apache.comet
 import scala.util.Random
 
 import org.apache.hadoop.fs.Path
-import org.apache.spark.sql.CometTestBase
+import org.apache.spark.sql.{CometTestBase, DataFrame}
 import org.apache.spark.sql.catalyst.expressions.ArrayContains
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.internal.SQLConf
@@ -31,7 +31,14 @@ import org.apache.spark.sql.types.BinaryType
 import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus
 import org.apache.comet.testing.{DataGenOptions, ParquetGenerator, SchemaGenOptions}
 
-class CometMapExpressionSuite extends CometTestBase {
+class CometMapExpressionSuite extends CometTestBase with CometCodegenAssertions {
+
+  private def assertMapSortInPlan(df: DataFrame): Unit = {
+    val plan = df.queryExecution.optimizedPlan
+    assert(
+      plan.exists(_.expressions.exists(_.exists(_.prettyName == "mapsort"))),
+      s"expected MapSort in optimized plan:\n$plan")
+  }
 
   test("read map[int, int] from parquet") {
 
@@ -247,6 +254,125 @@ class CometMapExpressionSuite extends CometTestBase {
     }
   }
 
+  test("mapsort routes array keys through codegen dispatcher") {
+    assume(isSpark40Plus, "Spark 4.0 inserts MapSort for group-by on map keys")
+    withTable("t_map_sort_array_key") {
+      sql("CREATE TABLE t_map_sort_array_key (m MAP<ARRAY<INT>, INT>) USING parquet")
+      sql("""INSERT INTO t_map_sort_array_key VALUES
+            |(map(array(2, 1), 20, array(1, 2), 10)),
+            |(map(array(1, 2), 10, array(2, 1), 20)),
+            |(map(array(3), 30)),
+            |(NULL)""".stripMargin)
+      val df = sql("SELECT m, count(*) FROM t_map_sort_array_key GROUP BY m")
+
+      assertMapSortInPlan(df)
+      val (_, cometPlan) = assertCodegenRan {
+        checkSparkAnswer(df)
+      }
+      val dispatched = new ExtendedExplainInfo().getCodegenDispatchExpressions(cometPlan)
+      assert(
+        dispatched.contains("mapsort"),
+        s"expected mapsort on codegen dispatch path, got $dispatched in:\n$cometPlan")
+    }
+  }
+
+  test("mapsort routes struct keys through codegen dispatcher") {
+    assume(isSpark40Plus, "Spark 4.0 inserts MapSort for group-by on map keys")
+    withTable("t_map_sort_struct_key") {
+      sql("""CREATE TABLE t_map_sort_struct_key (
+            |  m MAP<STRUCT<a: INT, b: STRING>, INT>) USING parquet""".stripMargin)
+      sql("""INSERT INTO t_map_sort_struct_key VALUES
+            |(map(named_struct('a', 2, 'b', 'b'), 20,
+            |     named_struct('a', 1, 'b', 'a'), 10)),
+            |(map(named_struct('a', 1, 'b', 'a'), 10,
+            |     named_struct('a', 2, 'b', 'b'), 20)),
+            |(map(named_struct('a', 3, 'b', 'c'), 30)),
+            |(NULL)""".stripMargin)
+      val df = sql("SELECT m, count(*) FROM t_map_sort_struct_key GROUP BY m")
+
+      assertMapSortInPlan(df)
+      val (_, cometPlan) = assertCodegenRan {
+        checkSparkAnswer(df)
+      }
+      val dispatched = new ExtendedExplainInfo().getCodegenDispatchExpressions(cometPlan)
+      assert(
+        dispatched.contains("mapsort"),
+        s"expected mapsort on codegen dispatch path, got $dispatched in:\n$cometPlan")
+    }
+  }
+
+  test("mapsort keeps scalar keys on the native path") {
+    assume(isSpark40Plus, "Spark 4.0 inserts MapSort for group-by on map keys")
+    withTable("t_map_sort_scalar_key") {
+      sql("CREATE TABLE t_map_sort_scalar_key (m MAP<INT, INT>) USING parquet")
+      sql("""INSERT INTO t_map_sort_scalar_key VALUES
+            |(map(2, 20, 1, 10)),
+            |(map(1, 10, 2, 20)),
+            |(map(3, 30)),
+            |(NULL)""".stripMargin)
+      val df = sql("SELECT m, count(*) FROM t_map_sort_scalar_key GROUP BY m")
+
+      assertMapSortInPlan(df)
+      val (_, cometPlan) = assertCodegenDidNotRun(checkSparkAnswer(df))
+      val explain = new ExtendedExplainInfo()
+      val nativeExpressions = explain.getNativeExpressions(cometPlan)
+      assert(
+        nativeExpressions.contains("mapsort"),
+        s"expected native mapsort expression, got $nativeExpressions in:\n$cometPlan")
+      assert(
+        !explain.getCodegenDispatchExpressions(cometPlan).contains("mapsort"),
+        s"scalar-key mapsort should not use codegen dispatch:\n$cometPlan")
+    }
+  }
+
+  test("mapsort with complex keys falls back when codegen dispatcher is disabled") {
+    assume(isSpark40Plus, "Spark 4.0 inserts MapSort for group-by on map keys")
+    withTable("t_map_sort_dispatch_disabled") {
+      sql("CREATE TABLE t_map_sort_dispatch_disabled (m MAP<ARRAY<INT>, INT>) USING parquet")
+      sql("""INSERT INTO t_map_sort_dispatch_disabled VALUES
+            |(map(array(2, 1), 20, array(1, 2), 10)),
+            |(map(array(1, 2), 10, array(2, 1), 20))""".stripMargin)
+      withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false") {
+        val df = sql("SELECT m, count(*) FROM t_map_sort_dispatch_disabled GROUP BY m")
+
+        assertMapSortInPlan(df)
+        assertCodegenDidNotRun {
+          checkSparkAnswerAndFallbackReason(
+            df,
+            CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key + "=false")
+        }
+      }
+    }
+  }
+
+  // Spark rejects collated strings as map keys (`UNSUPPORTED_FEATURE.COLLATIONS_IN_MAP_KEYS`), so
+  // `MapSort` never sees that shape. `supportedScalarSortElementType` still excludes them, and the
+  // same `Unsupported` → dispatcher path is covered by the array/struct cases above.
+
+  test("mapsort routes strict floating-point keys through codegen dispatcher") {
+    assume(isSpark40Plus, "Spark 4.0 inserts MapSort for group-by on map keys")
+    withSQLConf(CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key -> "true") {
+      withTable("t_map_sort_fp_key") {
+        sql("CREATE TABLE t_map_sort_fp_key (m MAP<DOUBLE, INT>) USING parquet")
+        sql("""INSERT INTO t_map_sort_fp_key VALUES
+              |(map(CAST('NaN' AS DOUBLE), 1, CAST('-0.0' AS DOUBLE), 2, 1.0, 3)),
+              |(map(1.0, 3, CAST('-0.0' AS DOUBLE), 2, CAST('NaN' AS DOUBLE), 1)),
+              |(map(0.0, 4)),
+              |(NULL)""".stripMargin)
+        val df = sql("SELECT m, count(*) FROM t_map_sort_fp_key GROUP BY m")
+
+        assertMapSortInPlan(df)
+        val (_, cometPlan) = assertCodegenRan {
+          checkSparkAnswer(df)
+        }
+        val dispatched = new ExtendedExplainInfo().getCodegenDispatchExpressions(cometPlan)
+        assert(
+          dispatched.contains("mapsort"),
+          s"expected mapsort on codegen dispatch path, got $dispatched in:\n$cometPlan")
+      }
+    }
+  }
+
   test("map_from_entries - binary type routes through codegen dispatcher") {
     val table = "t2"
     withTable(table) {

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

@@ -38,7 +38,6 @@ import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
 
 import org.apache.comet.CometConf
-import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus
 
 abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper {
   protected val adaptiveExecutionEnabled: Boolean
@@ -142,10 +141,8 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar
   }
 
   test("columnar shuffle on array/struct map key/value") {
-    // Spark 4.0 normalizes maps used as shuffle keys with mapsort(...). Comet's map_sort
-    // relies on Arrow's sort_to_indices, which only supports scalar key types, so a map
-    // with array or struct keys cannot be sorted natively and the shuffle falls back.
-    val complexKeyShuffles = if (isSpark40Plus) 0 else 1
+    // Spark 4.0+ normalizes maps used as shuffle keys with mapsort(...). Native map_sort only
+    // supports scalar keys, so array and struct keys use Spark codegen in the Comet pipeline.
     Seq("false", "true").foreach { execEnabled =>
       Seq(10, 201).foreach { numPartitions =>
         Seq("1.0", "10.0").foreach { ratio =>
@@ -158,7 +155,7 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar
                 .repartition(numPartitions, $"_1", $"_2")
                 .sortWithinPartitions($"_2")
 
-              checkShuffleAnswer(df, complexKeyShuffles)
+              checkShuffleAnswer(df, 1)
             }
 
             withParquetTable((0 until 50).map(i => (Map(i -> Seq(i, i + 1)), i + 1)), "tbl") {
@@ -176,7 +173,7 @@ abstract class CometColumnarShuffleSuite extends CometTestBase with AdaptiveSpar
                 .repartition(numPartitions, $"_1", $"_2")
                 .sortWithinPartitions($"_2")
 
-              checkShuffleAnswer(df, complexKeyShuffles)
+              checkShuffleAnswer(df, 1)
             }
 
             withParquetTable((0 until 50).map(i => (Map(i -> ((i, i.toString))), i + 1)), "tbl") {

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

@@ -43,12 +43,15 @@ import org.apache.spark.sql.functions.{col, count, sum}
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
 
-import org.apache.comet.{CometConf, CometExecIterator, CometShuffleBlockIterator, CometShuffleSizeLimitException, Native}
+import org.apache.comet.{CometCodegenAssertions, CometConf, CometExecIterator, CometShuffleBlockIterator, CometShuffleSizeLimitException, Native}
 import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus
 import org.apache.comet.serde.{OperatorOuterClass, PartitioningOuterClass}
 import org.apache.comet.shuffle.ShufflePartitionPusher
 
-class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper {
+class CometNativeShuffleSuite
+    extends CometTestBase
+    with AdaptiveSparkPlanHelper
+    with CometCodegenAssertions {
   override protected def test(testName: String, testTags: Tag*)(testFun: => Any)(implicit
       pos: Position): Unit = {
     super.test(testName, testTags: _*) {
@@ -792,7 +795,8 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper
       // Map entry order carries no meaning, so equal maps must hash alike. Spark 4.0+ normalizes a
       // map shuffle key with `mapsort(...)`; earlier versions do not, so Comet must not hash a raw
       // map there. The gate therefore only admits map keys on Spark 4.0+, and only when the
-      // `mapsort` itself is convertible (CometMapSort supports scalar map keys only).
+      // `mapsort` itself can stay in the Comet pipeline, either natively or through the JVM
+      // codegen dispatcher.
       withParquetTable((0 until 50).map(i => (i, Map(i % 7 -> (i % 5)))), "tbl") {
         val df = sql("SELECT * FROM tbl").repartition(10, $"_2").sortWithinPartitions($"_1")
 
@@ -801,15 +805,28 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper
     }
   }
 
-  test("native shuffle on map hash partitioning key with non-scalar map key falls back") {
-    // A map whose own key is nested cannot be `mapsort`ed by Comet (Arrow's sort_to_indices
-    // handles scalar keys only), so the normalization Spark 4.0+ requires is unavailable and the
-    // shuffle must fall back rather than hash an unnormalized map.
+  test("native shuffle on map hash partitioning key with non-scalar map key uses dispatcher") {
+    // Arrow's sort_to_indices only handles scalar keys, so Spark's MapSort.doGenCode performs the
+    // normalization through the JVM codegen dispatcher while the shuffle remains native.
     assume(isSpark40Plus, "map shuffle keys are only normalized with mapsort on Spark 4.0+")
-    withParquetTable((0 until 50).map(i => (i, Map(Seq(i % 7) -> (i % 5)))), "tbl") {
-      val df = sql("SELECT * FROM tbl").repartition(10, $"_2").sortWithinPartitions($"_1")
+    withNestedHashPartitioning {
+      withParquetTable((0 until 50).map(i => (i, Map(Seq(i % 7) -> (i % 5)))), "tbl") {
+        val df = sql("SELECT * FROM tbl").repartition(10, $"_2").sortWithinPartitions($"_1")
 
-      checkShuffleAnswer(df, 0)
+        assertCodegenRan {
+          checkShuffleAnswer(df, 1)
+        }
+      }
+
+      withParquetTable(
+        (0 until 50).map(i => (i, Map((i % 7, (i % 5).toString) -> (i % 3)))),
+        "tbl") {
+        val df = sql("SELECT * FROM tbl").repartition(10, $"_2").sortWithinPartitions($"_1")
+
+        assertCodegenRan {
+          checkShuffleAnswer(df, 1)
+        }
+      }
     }
   }