feat: route MapSort fallback through the codegen dispatcher

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

@@ -21,14 +21,23 @@ under the License.
 
 ## MapSort (Spark 4.0+)
 
-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.
-
-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`.
+Spark 4.0 inserts `MapSort` to normalize map values when they appear in grouping expressions or
+shuffle hash partitioning keys. Comet runs `MapSort` natively for supported scalar key types.
+Other dispatcher-eligible orderable key types use Spark's own generated JVM code through the
+codegen dispatcher, so the enclosing operator can stay in the Comet pipeline. This dispatcher
+route is not a native `MapSort` implementation.
+
+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.
+
+Set `spark.comet.expression.MapSort.enabled=false` to restore the previous behavior, where a
+`MapSort` without a native implementation causes its enclosing projection or shuffle to fall back
+to Spark. This expression-specific setting leaves the codegen dispatcher available to unrelated
+expressions. Retaining an enclosing operator in Comet is a functional routing benefit; by itself,
+it does not guarantee higher throughput.
 
 <!--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,20 +25,27 @@ 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 without a native implementation can still run Spark's generated code in-pipeline.
+// Spark rejects collated-string map keys by default, but they reach MapSort when
+// spark.sql.collation.allowInMapKeys=true and the map is built from dispatcher-supported inputs;
+// those expressions take the same Unsupported -> dispatcher route. A scan carrying a collated
+// map schema may still be rejected independently by the scan's schema support checks.
+object CometMapSort extends CometExpressionSerde[MapSort] with CodegenDispatchFallback {
 
   override def getIncompatibleReasons(): Seq[String] =
     Seq(
       "MapSort on floating-point keys is not 100% compatible with Spark when " +
         s"`${CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key}=true`.")
 
   override def getUnsupportedReasons(): Seq[String] =
-    Seq("MapSort is unsupported for non-scalar key types (struct, array, map, etc.).")
+    Seq(
+      "MapSort with an orderable key type outside native scalar coverage, including array, " +
+        "struct, interval, and non-default-collated string keys, has no native implementation.")
 
   override def getSupportLevel(expr: MapSort): SupportLevel = {
     val keyType = expr.dataType.asInstanceOf[MapType].keyType
     if (!supportedScalarSortElementType(keyType)) {
-      Unsupported(Some(s"MapSort on map with key type $keyType is not supported"))
+      Unsupported(Some(s"MapSort with key type $keyType has no native implementation"))
     } else {
       SupportLevel
         .strictFloatingPointReason(keyType, "MapSort on floating-point key")

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,16 +22,24 @@ 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
 import org.apache.spark.sql.types.BinaryType
 
 import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus
 import org.apache.comet.testing.{DataGenOptions, ParquetGenerator, SchemaGenOptions}
+import org.apache.comet.udf.codegen.CometScalaUDFCodegen
 
-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 +255,240 @@ 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 expression disable restores fallback while codegen dispatcher stays enabled") {
+    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 -> "true",
+        CometConf.getExprEnabledConfigKey("MapSort") -> "false") {
+        val df = sql("SELECT m, count(*) FROM t_map_sort_dispatch_disabled GROUP BY m")
+
+        assertMapSortInPlan(df)
+        assertCodegenDidNotRun {
+          checkSparkAnswerAndFallbackReason(
+            df,
+            "Expression support is disabled. Set " +
+              s"${CometConf.getExprEnabledConfigKey("MapSort")}=true to enable it.")
+        }
+      }
+    }
+  }
+
+  test("mapsort routes collated-string keys through codegen dispatcher") {
+    assume(isSpark40Plus, "collated map keys and MapSort require Spark 4.0+")
+    withSQLConf(
+      "spark.sql.collation.allowInMapKeys" -> "true",
+      CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") {
+      withTable("t_map_sort_collated_key") {
+        sql(
+          "CREATE TABLE t_map_sort_collated_key " +
+            "(k1 STRING, v1 INT, k2 STRING, v2 INT) USING parquet")
+        sql("""INSERT INTO t_map_sort_collated_key VALUES
+              |('b', 20, 'A', 10),
+              |('a', 10, 'B', 20),
+              |('c', 30, 'd', 40)""".stripMargin)
+        val query =
+          """SELECT m, count(*) FROM (
+            |  SELECT map(CAST(k1 AS STRING COLLATE UTF8_LCASE), v1,
+            |             CAST(k2 AS STRING COLLATE UTF8_LCASE), v2) AS m
+            |  FROM t_map_sort_collated_key)
+            |GROUP BY m""".stripMargin
+        val df = sql(query)
+
+        assertMapSortInPlan(df)
+        CometScalaUDFCodegen.resetStats()
+        val cometRows = df.collect()
+        val cometPlan = df.queryExecution.executedPlan
+        val dispatcherStats = CometScalaUDFCodegen.stats()
+        assert(
+          dispatcherStats.totalLookups >= 1,
+          s"expected codegen dispatcher activity, got $dispatcherStats; " +
+            s"fallback reasons: ${new ExtendedExplainInfo().getFallbackReasons(cometPlan)}\n" +
+            cometPlan)
+
+        // UTF8_LCASE considers the first two maps equal but Spark and Comet may retain different
+        // byte-level representatives for the grouped key ("A" versus "a"). Compare the collected
+        // answers after canonicalizing keys in the test, leaving the executed SQL plan untouched.
+        def canonicalize(rows: Array[org.apache.spark.sql.Row]) =
+          rows
+            .map { row =>
+              val entries =
+                if (row.isNullAt(0)) {
+                  None
+                } else {
+                  Some(
+                    row
+                      .getMap[String, Int](0)
+                      .toSeq
+                      .map { case (key, value) =>
+                        key.toLowerCase(java.util.Locale.ROOT) -> value
+                      }
+                      .sortBy(_._1))
+                }
+              entries -> row.getLong(1)
+            }
+            .sortBy(_.toString)
+            .toSeq
+
+        var sparkRows: Array[org.apache.spark.sql.Row] = Array.empty
+        withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+          sparkRows = sql(query).collect()
+        }
+        assert(canonicalize(cometRows) === canonicalize(sparkRows))
+
+        val explain = new ExtendedExplainInfo()
+        assert(
+          explain.getCodegenDispatchExpressions(cometPlan).contains("mapsort"),
+          s"expected collated-key mapsort on codegen dispatch path:\n$cometPlan")
+        assert(
+          !explain.getNativeExpressions(cometPlan).contains("mapsort"),
+          s"collated-key mapsort must not use native map_sort:\n$cometPlan")
+
+        withSQLConf(CometConf.getExprEnabledConfigKey("MapSort") -> "false") {
+          val fallback = sql(query)
+          assertMapSortInPlan(fallback)
+          val fallbackRows = assertCodegenDidNotRun(fallback.collect())
+          val fallbackPlan = fallback.queryExecution.executedPlan
+          assert(canonicalize(fallbackRows) === canonicalize(sparkRows))
+          val expectedReason =
+            "Expression support is disabled. Set " +
+              s"${CometConf.getExprEnabledConfigKey("MapSort")}=true to enable it."
+          assert(
+            new ExtendedExplainInfo().getFallbackReasons(fallbackPlan).contains(expectedReason),
+            s"expected MapSort-specific fallback reason `$expectedReason` in:\n$fallbackPlan")
+        }
+      }
+    }
+  }
+
+  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",
+      CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true",
+      "spark.sql.legacy.disableMapKeyNormalization" -> "true") {
+      withTable("t_map_sort_fp_key") {
+        sql("CREATE TABLE t_map_sort_fp_key (id INT, m MAP<DOUBLE, INT>) USING parquet")
+        sql("""INSERT INTO t_map_sort_fp_key VALUES
+              |(1, map(CAST('0.0' AS DOUBLE), 10, CAST('-0.0' AS DOUBLE), 20,
+              |        CAST('NaN' AS DOUBLE), 30)),
+              |(2, map(CAST('-0.0' AS DOUBLE), 20, CAST('0.0' AS DOUBLE), 10,
+              |        CAST('NaN' AS DOUBLE), 30)),
+              |(3, map(CAST('NaN' AS DOUBLE), 40, 1.0, 50)),
+              |(4, NULL)""".stripMargin)
+
+        val storedKeys = sql(
+          "SELECT k FROM t_map_sort_fp_key " +
+            "LATERAL VIEW explode(map_keys(m)) e AS k WHERE id = 1").collect().map(_.getDouble(0))
+        val storedBits = storedKeys.map(java.lang.Double.doubleToRawLongBits)
+        assert(
+          storedBits.take(2).sameElements(Array(0L, Long.MinValue)),
+          s"expected stored +0.0 then -0.0 raw bits, got ${storedBits.toSeq}")
+        assert(storedKeys.exists(java.lang.Double.isNaN), "strict fixture must retain a NaN key")
+
+        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")
+        assert(
+          !new ExtendedExplainInfo().getNativeExpressions(cometPlan).contains("mapsort"),
+          s"strict floating-point mapsort must not use native map_sort:\n$cometPlan")
+
+        withSQLConf(CometConf.getExprEnabledConfigKey("MapSort") -> "false") {
+          val fallback = sql("SELECT m, count(*) FROM t_map_sort_fp_key GROUP BY m")
+          assertMapSortInPlan(fallback)
+          assertCodegenDidNotRun {
+            checkSparkAnswerAndFallbackReason(
+              fallback,
+              "Expression support is disabled. Set " +
+                s"${CometConf.getExprEnabledConfigKey("MapSort")}=true to enable it.")
+          }
+        }
+      }
+    }
+  }
+
   test("map_from_entries - binary type routes through codegen dispatcher") {
     val table = "t2"
     withTable(table) {