refactor: address review feedback on MapSort

- Use single global sort_to_indices+take instead of per-map take+concat
- Add early-out fast paths (empty array, all-null, is_sorted=true)
- Fall back to Spark for floating-point map keys when strictFloatingPoint=true
- Clean up Arc::clone calls and replace .unwrap() on downcasts with .expect
- Document MapSort behavior in map expressions compatibility guide

Author: andygrove

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

@@ -19,5 +19,16 @@ under the License.
 
 # Map Expressions
 
+## 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`.
+
 <!--BEGIN:EXPR_COMPAT[map]-->
 <!--END:EXPR_COMPAT-->

native/spark-expr/src/map_funcs/map_sort.rs

@@ -15,8 +15,8 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use arrow::array::{Array, ArrayRef, MapArray, StructArray};
-use arrow::compute::{concat, sort_to_indices, take, SortOptions};
+use arrow::array::{Array, ArrayRef, MapArray, StructArray, UInt32Array};
+use arrow::compute::{sort_to_indices, take, SortOptions};
 use arrow::datatypes::DataType;
 use datafusion::common::{exec_err, DataFusionError};
 use datafusion::physical_plan::ColumnarValue;
@@ -31,59 +31,62 @@ pub fn spark_map_sort(args: &[ColumnarValue]) -> Result<ColumnarValue, DataFusio
     }
 
     let arr_arg: ArrayRef = match &args[0] {
-        ColumnarValue::Array(array) => Arc::<dyn Array>::clone(array),
+        ColumnarValue::Array(array) => Arc::clone(array),
         ColumnarValue::Scalar(scalar) => scalar.to_array_of_size(1)?,
     };
 
     let (maps_arg, map_field, is_sorted) = match arr_arg.data_type() {
         DataType::Map(map_field, is_sorted) => {
-            let maps_arg = arr_arg.as_any().downcast_ref::<MapArray>().unwrap();
+            let maps_arg = arr_arg
+                .as_any()
+                .downcast_ref::<MapArray>()
+                .expect("invariant: array data type is Map but downcast to MapArray failed");
             (maps_arg, map_field, is_sorted)
         }
         _ => return exec_err!("spark_map_sort expects Map type as argument"),
     };
 
+    // Fast paths: nothing to sort, all maps null, or input already declared sorted.
+    if maps_arg.is_empty() || maps_arg.null_count() == maps_arg.len() || *is_sorted {
+        return Ok(ColumnarValue::Array(arr_arg));
+    }
+
     let maps_arg_entries = maps_arg.entries();
     let maps_arg_offsets = maps_arg.offsets();
 
-    let mut sorted_map_entries_vec: Vec<ArrayRef> = Vec::with_capacity(maps_arg.len());
+    let sort_options = SortOptions {
+        descending: false,
+        nulls_first: true,
+    };
+
+    // Build one global permutation over the full entries struct, respecting per-map boundaries,
+    // then issue a single `take`. This avoids per-map struct copies and a final `concat`.
+    let mut global_indices: Vec<u32> = Vec::with_capacity(maps_arg_entries.len());
 
     for idx in 0..maps_arg.len() {
         let map_start = maps_arg_offsets[idx] as usize;
         let map_end = maps_arg_offsets[idx + 1] as usize;
-        let map_len = map_end - map_start;
-
-        let map_entries = maps_arg_entries.slice(map_start, map_len);
-
-        if map_len == 0 {
-            sorted_map_entries_vec.push(Arc::new(map_entries));
+        if map_end == map_start {
             continue;
         }
 
-        let map_keys = map_entries.column(0);
-        let sort_options = SortOptions {
-            descending: false,
-            nulls_first: true,
-        };
-        let sorted_indices = sort_to_indices(&map_keys, Some(sort_options), None)?;
-
-        let sorted_map_entries = take(&map_entries, &sorted_indices, None)?;
-        sorted_map_entries_vec.push(sorted_map_entries);
+        let map_keys = maps_arg_entries
+            .column(0)
+            .slice(map_start, map_end - map_start);
+        let local_indices = sort_to_indices(&map_keys, Some(sort_options), None)?;
+        global_indices.extend(local_indices.values().iter().map(|i| map_start as u32 + *i));
     }
 
-    let sorted_map_entries_arr: Vec<&dyn Array> = sorted_map_entries_vec
-        .iter()
-        .map(|arr| arr.as_ref())
-        .collect();
-    let combined_sorted_map_entries = concat(&sorted_map_entries_arr)?;
-    let sorted_map_struct = combined_sorted_map_entries
+    let indices = UInt32Array::from(global_indices);
+    let sorted_entries = take(maps_arg_entries, &indices, None)?;
+    let sorted_map_struct = sorted_entries
         .as_any()
         .downcast_ref::<StructArray>()
-        .unwrap();
+        .expect("invariant: take on StructArray must return StructArray");
 
     // Preserve the original is_sorted flag to keep schema consistent
     let sorted_map_arr = Arc::new(MapArray::try_new(
-        Arc::<arrow::datatypes::Field>::clone(map_field),
+        Arc::clone(map_field),
         maps_arg.offsets().clone(),
         sorted_map_struct.clone(),
         maps_arg.nulls().cloned(),

spark/src/main/spark-4.0/org/apache/comet/shims/CometExprShim.scala

@@ -26,9 +26,10 @@ import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.internal.types.StringTypeWithCollation
 import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, DataTypes, MapType, StringType}
 
+import org.apache.comet.CometConf
 import org.apache.comet.CometSparkSessionExtensions.withInfo
 import org.apache.comet.expressions.{CometCast, CometEvalMode}
-import org.apache.comet.serde.{CommonStringExprs, Compatible, ExprOuterClass, Incompatible}
+import org.apache.comet.serde.{CommonStringExprs, Compatible, ExprOuterClass, Incompatible, SupportLevel}
 import org.apache.comet.serde.ExprOuterClass.{BinaryOutputStyle, Expr}
 import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, optExprWithInfo, scalarFunctionExprToProto, scalarFunctionExprToProtoWithReturnType, supportedScalarSortElementType}
 
@@ -149,6 +150,14 @@ trait CometExprShim extends CommonStringExprs {
         if (!supportedScalarSortElementType(keyType)) {
           withInfo(ms, s"MapSort on map with key type $keyType is not supported")
           None
+        } else if (CometConf.COMET_EXEC_STRICT_FLOATING_POINT.get() &&
+          SupportLevel.containsFloatingPoint(keyType)) {
+          withInfo(
+            ms,
+            "MapSort on floating-point key is not 100% compatible with Spark, and Comet is " +
+              s"running with ${CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key}=true. " +
+              s"${CometConf.COMPAT_GUIDE}")
+          None
         } else {
           val childExpr = exprToProtoInternal(ms.child, inputs, binding)
           val mapSortExpr = scalarFunctionExprToProtoWithReturnType(