diff --git a/native/spark-expr/src/map_funcs/map_sort.rs b/native/spark-expr/src/map_funcs/map_sort.rs index 16cc51946bb..22e8b840f50 100644 --- a/native/spark-expr/src/map_funcs/map_sort.rs +++ b/native/spark-expr/src/map_funcs/map_sort.rs @@ -16,6 +16,7 @@ // under the License. use arrow::array::{Array, ArrayRef, MapArray, StructArray, UInt32Array}; +use arrow::buffer::OffsetBuffer; use arrow::compute::{sort_to_indices, take, SortOptions}; use arrow::datatypes::DataType; use datafusion::common::{exec_err, DataFusionError}; @@ -61,20 +62,26 @@ pub fn spark_map_sort(args: &[ColumnarValue]) -> Result = Vec::with_capacity(maps_arg_entries.len()); + let mut rebased_offsets: Vec = Vec::with_capacity(maps_arg.len() + 1); + rebased_offsets.push(0); 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; - if map_end == map_start { - continue; + if map_end > map_start { + 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 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)); + rebased_offsets.push(global_indices.len() as i32); } let indices = UInt32Array::from(global_indices); @@ -87,7 +94,7 @@ pub fn spark_map_sort(args: &[ColumnarValue]) -> Result a, + _ => panic!("expected an array"), + }; + let sorted = sorted.as_any().downcast_ref::().unwrap(); + assert_eq!(sorted.len(), 1); + let keys = sorted + .keys() + .as_any() + .downcast_ref::() + .unwrap(); + let values = sorted + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(keys.iter().collect::>(), vec![Some("c"), Some("d")]); + assert_eq!(values.iter().collect::>(), vec![Some(3), Some(4)]); + } + macro_rules! build_map { ( $key_builder:expr, diff --git a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala index f3fba58c605..f4a559b872b 100644 --- a/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometMapExpressionSuite.scala @@ -460,4 +460,21 @@ class CometMapExpressionSuite extends CometTestBase { } } + test("mapsort on a sliced map does not overrun the sorted entries") { + // A native OFFSET slices the batch and Arrow keeps a sliced MapArray's original entry offsets, + // so `mapsort` receives a map whose first entry offset is nonzero. `spark_map_sort` takes only + // the visible entries, so reusing the input offsets overran them and failed the query with + // "Max offset of N exceeds length of entries M". + assume(isSpark40Plus, "Spark 4.0 inserts MapSort for group-by and repartition on map keys") + withParquetTable( + (0 until 20).map(i => (i, Map(s"b${i % 5}" -> i, s"a${i % 5}" -> (i + 1)))), + "tbl") { + // GROUP BY on a map: InsertMapSortInGroupingExpressions adds the mapsort. Repartition on a + // map reaches the same code through InsertMapSortInRepartitionExpressions, but native shuffle + // rejects map partitioning keys today, so this is the reachable path. + checkSparkAnswer( + "SELECT _2, count(*) FROM (SELECT * FROM tbl ORDER BY _1 LIMIT 15 OFFSET 5) GROUP BY _2") + } + } + }