feat: add MapSort expression support for Spark 4.0

[andygrove]

Which issue does this PR close?

Closes #1941
Closes #3171

Rationale for this change

Spark 4.0 introduces MapSort, used for normalizing map values when they appear in shuffle hash partitioning keys, in try_element_at, and in other contexts where map ordering must be deterministic. Without native support, queries that touch maps in any of these positions fall back to Spark, which forces the entire enclosing operator off Comet (e.g. an entire shuffle exchange).

What changes are included in this PR?

• New native scalar function map_sort in native/spark-expr/src/map_funcs/map_sort.rs that sorts map entries by key in ascending order, registered via comet_scalar_funcs.rs.
• Wire MapSort into the Spark 4.0 CometExprShim so the expression is converted to the new scalar function during serde.
• The columnar shuffle on map array element test in CometColumnarShuffleSuite now expects shuffle fallback on Spark 4.0+: the new shuffle-key normalization wraps mapsort inside transform(arr, x -> mapsort(x)), and Comet does not currently support ArrayTransform with a lambda body. Answer correctness is still verified via checkSparkAnswer.

How are these changes tested?

• New unit tests in native/spark-expr/src/map_funcs/map_sort.rs cover sorting on each supported key type, null handling, and empty maps.
• Existing CometColumnarShuffleSuite tests for map shuffle keys all pass under the Spark 4.0 profile (41/41).

[mbutrovich]

Nice, glad we're cleaning up this ugliness.

[mbutrovich]

Thanks @andygrove! Categorized suggestions below:

Performance

One global take instead of per-map sort_to_indices + take + concat?

spark_map_sort loops over every map in the batch, calling sort_to_indices and take on each slice, then a final concat over n struct slices. For n maps that's roughly 2n + 1 kernel calls and 2n + 1 allocations (n small index arrays, n full per-map struct copies from take, one full concat copy).

The version below keeps the n sort_to_indices calls but drops the rest: one Vec<u32> grow (pre-sized), one take over the full entries, no concat. That's n + 1 kernel calls and roughly n + 2 allocations, and crucially it eliminates all per-map struct copies plus the final concat copy, which are the expensive ones. Peak transient memory drops from roughly 3x entry volume to roughly 1x.

Compare to DataFusion's sort_batch (datafusion/physical-plan/src/sorts/sort.rs:813), which is what ExternalSorter uses to sort an in-memory batch:

let indices = lexsort_to_indices(&sort_columns, fetch)?;
let columns = take_arrays(batch.columns(), &indices, None)?;

One global permutation, one take. No per-row work, no trailing concat.

map_sort can't call sort_to_indices once on the full key column. It needs to sort within each map boundary. But the same shape applies: build one global index buffer respecting group boundaries, then a single take on the full entries:

let total = maps_arg_entries.len();
let mut global: Vec<u32> = Vec::with_capacity(total);

for idx in 0..maps_arg.len() {
    let start = maps_arg_offsets[idx] as usize;
    let end   = maps_arg_offsets[idx + 1] as usize;
    if end == start { continue; }

    let keys  = maps_arg_entries.column(0).slice(start, end - start);
    let local = sort_to_indices(&keys, Some(sort_options), None)?;
    global.extend(local.values().iter().map(|i| start as u32 + *i));
}

let indices = UInt32Array::from(global);
let sorted  = take(&maps_arg_entries, &indices, None)?;

Offsets and nulls stay untouched. Output is one struct array, matching sort_batch's pattern of global-permutation-plus-single-take.

Early outs?

Worth short-circuiting these up front?

if maps_arg.is_empty() { return Ok(ColumnarValue::Array(arr_arg)); }
if maps_arg.null_count() == maps_arg.len() { return Ok(ColumnarValue::Array(arr_arg)); }
if *is_sorted { return Ok(ColumnarValue::Array(arr_arg)); }

Microbenchmark numbers?

It would be nice to have before/after numbers in the PR description. Spark 4.0 map-shuffle now stays native, so there's a natural story to measure against the prior fallback path.

Spark compatibility

NaN / -0.0 on float-keyed maps

Spark's key ordering goes through Double.compare (NaN sorts largest, -0.0 == 0.0). Arrow's sort_to_indices uses IEEE total ordering (-0.0 < 0.0, NaN at extremes). Comet already tracks this class of incompatibility via spark.comet.exec.strictFloatingPoint.

CometSortOrder (serde/CometSortOrder.scala:34) and CometSortArray (serde/arrays.scala:150) both consult COMET_EXEC_STRICT_FLOATING_POINT and return Incompatible when the child type contains Float/Double. The MapSort shim in this PR doesn't. MapSort should follow the same pattern:

case ms: MapSort =>
  val keyType = ms.dataType.asInstanceOf[MapType].keyType
  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, s"MapSort on floating-point key not compatible with Spark when " +
      s"${CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key}=true")
    None
  } else {
    // ... existing conversion
  }

Default behavior is unchanged (Comet converts the expression, accepting the float divergence). Users who set strictFloatingPoint=true get the same fallback behavior they already get for SortOrder and SortArray.

supportedRangePartitioningDataType in CometShuffleExchangeExec.scala:344 has the same gap but is out of scope for this PR. I will open an issue to follow up on that.

Style nits

• native/spark-expr/src/map_funcs/map_sort.rs:75,128. Arc::<dyn Array>::clone(array) and Arc::<arrow::datatypes::Field>::clone(map_field) can be Arc::clone(array) / map_field.clone().
• .unwrap() on the downcast_ref at :81 and :121. An .expect("invariant: ...") documents what's being assumed.

Docs

Worth a mention in docs/source/user-guide/compatibility.md that Spark 4.0 map-shuffle / group-by-on-map now stays native (closing #1941, #3171)?

[andygrove]

Thanks @mbutrovich. Pushed 1620e33b7 addressing the feedback:

Performance

• Replaced the per-map take + final concat with one global Vec<u32> index buffer plus a single take over the full entries struct, mirroring sort_batch.
• Added the three early-outs (empty input, all-null, *is_sorted) at the top of spark_map_sort.

Spark compatibility

• MapSort now falls back to Spark for floating-point keys when spark.comet.exec.strictFloatingPoint=true, matching CometSortOrder and CometSortArray.

Style

• Switched to Arc::clone(...) and .expect("invariant: ...") on the two downcasts.

Docs

• Added a MapSort section to docs/source/user-guide/latest/compatibility/expressions/map.md describing the Spark 4.0 behavior and the strict floating-point fallback.

I'll add benchmarks next

[andygrove]

Microbenchmark numbers from native/spark-expr/benches/map_sort.rs (6e832541d). Each case is a batch of 8192 maps; "old" is the per-map take + concat impl, "new" is the global-permutation + single take impl. Release build, M-series Mac.

Case	Entries/map	Old	New	Speedup
int_keys/4	4	3.99 ms	0.77 ms	5.2x
int_keys/16	16	4.44 ms	1.25 ms	3.6x
int_keys/64	64	5.02 ms	1.62 ms	3.1x
string_keys/4	4	4.52 ms	0.97 ms	4.7x
string_keys/16	16	7.16 ms	3.87 ms	1.85x
string_keys/64	64	8.95 ms	4.24 ms	2.1x

The biggest wins are the small-map cases where the per-map struct copies and trailing concat dominated. To run locally:

cd native && cargo bench --bench map_sort

[mbutrovich]

Looks awesome! Thanks for the quick turnaround on the feedback, @andygrove!

[comphead]

is this shim only for 4.0 or need to be in 4.x?

[andygrove]

yeah the PR was started before 4.1 was added, and upmerging broke the tests. added shims for 4.1 and 4.2 now

[comphead]

there is also a folder spark/src/main/spark-4.x to serve all Spark4 subversions