From a8f5be0ec36c1c52d76b64553d2e2c365ba8dd12 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sat, 12 Sep 2026 11:46:24 -0700 Subject: [PATCH 1/2] perf: optimize map_sort singleton normalization --- .../expression-audits/map_funcs.md | 4 + .../spark-expr/benches/common/matched_maps.rs | 123 +++++++ native/spark-expr/benches/hash.rs | 2 + native/spark-expr/benches/map_sort.rs | 53 ++- native/spark-expr/src/map_funcs/map_sort.rs | 315 ++++++++++++++++-- 5 files changed, 468 insertions(+), 29 deletions(-) diff --git a/docs/source/contributor-guide/expression-audits/map_funcs.md b/docs/source/contributor-guide/expression-audits/map_funcs.md index ea13e6ab130..bdaff4bded8 100644 --- a/docs/source/contributor-guide/expression-audits/map_funcs.md +++ b/docs/source/contributor-guide/expression-audits/map_funcs.md @@ -64,6 +64,10 @@ - Spark 4.0.1 (audited 2026-05-27): semantics unchanged. - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. +## map_sort + +- Performance (tuned locally 2026-09-12; PR pending): skip Arrow sort dispatch for eligible flat singleton keys, with a batch check and specialized fallback loop for batches without singletons. In the local DataFusion 55.0.0 development cohort, matched singleton normalization measured 19–22x faster in the full run and 18.4x in an independent forward-order confirmation. Benchmarks: `native/spark-expr/benches/map_sort.rs`, `hash.rs`, and `common/matched_maps.rs`; 92 cases cover normalization, hashing, combined execution, nulls, slices, mixed cardinalities, and long Unicode values. Flagged regressions did not remain stable through independent and reversed-order confirmation. + ## map_values - Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8. diff --git a/native/spark-expr/benches/common/matched_maps.rs b/native/spark-expr/benches/common/matched_maps.rs index 0be1e0c04e1..2bdd448ddcc 100644 --- a/native/spark-expr/benches/common/matched_maps.rs +++ b/native/spark-expr/benches/common/matched_maps.rs @@ -256,3 +256,126 @@ pub fn bench_maps(c: &mut Criterion, stage: Stage) { } group.finish(); } + +/// Additional regression shapes use the same rows and hash seed as the matched fixtures. +/// Null maps retain physical entries, exercising normalization underneath null slots too. +/// The leading row is sliced away to exercise nonzero entry offsets in every case. +pub fn bench_regression_maps(c: &mut Criterion, stage: Stage) { + let ints: ArrayRef = Arc::new(Int32Array::from_iter_values((0..ROWS).map(c1))); + let mut group = c.benchmark_group(format!("regression_maps/{}", stage.name())); + group.throughput(Throughput::Elements(ROWS as u64)); + for (name, null_every, mixed, long) in [ + ("singleton_no_null", 0, false, false), + ("singleton_sparse_null", 100, false, false), + ("singleton_dense_null", 2, false, false), + ("mixed_no_null", 0, true, false), + ("mixed_sparse_null", 100, true, false), + ("mixed_dense_null", 2, true, false), + ("singleton_long_unicode", 0, false, true), + ("mixed_long_unicode_dense_null", 2, true, true), + ] { + let mut builder = MapBuilder::new( + Some(crate::common::map_field_names()), + StringBuilder::new(), + StringBuilder::new(), + ); + let prefix = if long { + "資料é".repeat(128) + } else { + String::new() + }; + for row in 0..=ROWS { + let count = if row == 0 { + 3 // Sliced-away prefix must have entries, even for mixed cardinalities. + } else if mixed { + [0, 1, 1, 2, 10, 50][row % 6] + } else { + 1 + }; + for entry in (0..count).rev() { + builder.keys().append_value(format!("{prefix}{entry:04}")); + if null_every != 0 && (row + entry + 1) % null_every == 0 { + builder.values().append_null(); + } else { + builder + .values() + .append_value(format!("{prefix}{row}:{entry}")); + } + } + builder + .append(null_every == 0 || row % null_every != 0) + .unwrap(); + } + let raw: ArrayRef = Arc::new(builder.finish().slice(1, ROWS)); + let args = [ColumnarValue::Array(Arc::clone(&raw))]; + let normalized = normalize(&args); + // Independent expected permutation checks values, child nulls, and schema as well + // as ordering; explicit checks cover map validity and rebased sliced offsets. + let map = raw.as_any().downcast_ref::().unwrap(); + let keys = map + .keys() + .as_any() + .downcast_ref::() + .unwrap(); + let mut permutation = Vec::new(); + let mut offsets = vec![0i32]; + for pair in map.value_offsets().windows(2) { + let mut indices: Vec = (pair[0] as u32..pair[1] as u32).collect(); + indices.sort_by(|a, b| keys.value(*a as usize).cmp(keys.value(*b as usize))); + permutation.extend(indices); + offsets.push(permutation.len() as i32); + } + let expected_entries = arrow::compute::take( + map.entries(), + &arrow::array::UInt32Array::from(permutation), + None, + ) + .unwrap(); + let actual = normalized.as_any().downcast_ref::().unwrap(); + assert_eq!(actual.entries().to_data(), expected_entries.to_data()); + assert_eq!(actual.value_offsets(), offsets); + assert_eq!(actual.nulls(), map.nulls()); + assert_eq!(actual.data_type(), map.data_type()); + for shape in [Shape::Map, Shape::StructMapInt] { + if matches!((stage, shape), (Stage::NormalizeOnly, Shape::StructMapInt)) { + continue; + } + group.bench_function(BenchmarkId::new(shape.name(), name), |b| match stage { + Stage::NormalizeOnly => b.iter(|| black_box(normalize(black_box(&args)))), + Stage::HashOnly => { + let input = match shape { + Shape::Map => Arc::clone(&normalized), + Shape::StructMapInt => wrap(Arc::clone(&normalized), &ints), + }; + let mut buffer = vec![42; ROWS]; + b.iter(|| { + buffer.fill(42); + create_murmur3_hashes(std::slice::from_ref(black_box(&input)), &mut buffer) + .unwrap(); + black_box(&buffer); + }); + } + Stage::NormalizeHash => { + let mut buffer = vec![42; ROWS]; + match shape { + Shape::Map => b.iter(|| { + buffer.fill(42); + let input = normalize(black_box(&args)); + create_murmur3_hashes(std::slice::from_ref(&input), &mut buffer) + .unwrap(); + black_box(&buffer); + }), + Shape::StructMapInt => b.iter(|| { + buffer.fill(42); + let input = wrap(normalize(black_box(&args)), &ints); + create_murmur3_hashes(std::slice::from_ref(&input), &mut buffer) + .unwrap(); + black_box(&buffer); + }), + } + } + }); + } + } + group.finish(); +} diff --git a/native/spark-expr/benches/hash.rs b/native/spark-expr/benches/hash.rs index 4d8d4026cfd..e88dfaa4e90 100644 --- a/native/spark-expr/benches/hash.rs +++ b/native/spark-expr/benches/hash.rs @@ -267,6 +267,8 @@ fn bench(c: &mut Criterion) { fn bench_matched_maps(c: &mut Criterion) { matched_maps::bench_maps(c, matched_maps::Stage::HashOnly); matched_maps::bench_maps(c, matched_maps::Stage::NormalizeHash); + matched_maps::bench_regression_maps(c, matched_maps::Stage::HashOnly); + matched_maps::bench_regression_maps(c, matched_maps::Stage::NormalizeHash); c.bench_function("matched_maps/hash_buffer_seed_reset", |b| { let mut hashes = vec![42u32; matched_maps::ROWS]; b.iter(|| { diff --git a/native/spark-expr/benches/map_sort.rs b/native/spark-expr/benches/map_sort.rs index f32b9eb5bf2..8c5b4c9ba15 100644 --- a/native/spark-expr/benches/map_sort.rs +++ b/native/spark-expr/benches/map_sort.rs @@ -74,7 +74,7 @@ fn build_string_key_map(entries_per_map: usize) -> MapArray { fn bench_map_sort(c: &mut Criterion) { let mut group = c.benchmark_group("spark_map_sort"); - for entries in [4usize, 16, 64] { + for entries in [0usize, 1, 4, 16, 64] { let int_map: ArrayRef = Arc::new(build_int_key_map(entries)); group.bench_with_input( BenchmarkId::new("int_keys", entries), @@ -99,9 +99,58 @@ fn bench_map_sort(c: &mut Criterion) { group.finish(); } +// Struct keys are rejected by Arrow even when each map has just one entry. +// Keep the fallible path in the baseline suite so type validation cannot disappear. +fn bench_unsupported_singleton(c: &mut Criterion) { + use arrow::array::{Array, Int32Array, StructArray}; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::{DataType, Field}; + + let keys = StructArray::new( + vec![Arc::new(Field::new("k", DataType::Int32, false))].into(), + vec![Arc::new(Int32Array::from_iter_values(0..BATCH_SIZE as i32))], + None, + ); + let entries = StructArray::new( + vec![ + Arc::new(Field::new("key", keys.data_type().clone(), false)), + Arc::new(Field::new("value", DataType::Int32, true)), + ] + .into(), + vec![ + Arc::new(keys), + Arc::new(Int32Array::from(vec![1; BATCH_SIZE])), + ], + None, + ); + let map = MapArray::new( + Arc::new(Field::new("entries", entries.data_type().clone(), false)), + OffsetBuffer::new((0..=BATCH_SIZE as i32).collect::>().into()), + entries, + None, + false, + ); + let args = [ColumnarValue::Array(Arc::new(map))]; + assert!(spark_map_sort(&args) + .unwrap_err() + .to_string() + .contains("Sort not supported")); + // Input creation is untimed; error creation and drop are timed. Only the first + // nonempty row is visited, so this is not full-batch throughput. + c.bench_function("spark_map_sort/unsupported_singleton_struct_key", |b| { + b.iter(|| black_box(spark_map_sort(black_box(&args)).unwrap_err())); + }); +} + fn bench_matched_maps(c: &mut Criterion) { matched_maps::bench_maps(c, matched_maps::Stage::NormalizeOnly); + matched_maps::bench_regression_maps(c, matched_maps::Stage::NormalizeOnly); } -criterion_group!(benches, bench_map_sort, bench_matched_maps); +criterion_group!( + benches, + bench_map_sort, + bench_matched_maps, + bench_unsupported_singleton +); criterion_main!(benches); diff --git a/native/spark-expr/src/map_funcs/map_sort.rs b/native/spark-expr/src/map_funcs/map_sort.rs index 22e8b840f50..58de949dfd6 100644 --- a/native/spark-expr/src/map_funcs/map_sort.rs +++ b/native/spark-expr/src/map_funcs/map_sort.rs @@ -55,34 +55,34 @@ 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); + // Arrow rejects some key types even for a singleton (e.g. Struct), and nested sorts + // can fail while ranking child values. Only skip dispatch for flat types whose sort + // is infallible; all other types must retain Arrow's original validation/error path. + let key_type = maps_arg_entries.column(0).data_type(); + let can_skip_singleton_sort = key_type.is_primitive() + || matches!( + key_type, + DataType::Boolean + | DataType::Utf8 + | DataType::LargeUtf8 + | DataType::Utf8View + | DataType::Binary + | DataType::LargeBinary + | DataType::BinaryView + | DataType::FixedSizeBinary(_) + ); - 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 { - 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); - } + // Keep the original loop for batches without eligible singletons. Specializing the + // loop avoids adding a per-row branch to wide maps. All-empty visible slices need + // no scan, including slices whose entries array still contains an unused prefix. + let has_singletons = can_skip_singleton_sort + && maps_arg_offsets[maps_arg.len()] > maps_arg_offsets[0] + && maps_arg_offsets.windows(2).any(|w| w[1] - w[0] == 1); + let (global_indices, rebased_offsets) = if has_singletons { + map_sort_indices::(maps_arg_entries, maps_arg_offsets)? + } else { + map_sort_indices::(maps_arg_entries, maps_arg_offsets)? + }; let indices = UInt32Array::from(global_indices); let sorted_entries = take(maps_arg_entries, &indices, None)?; @@ -103,6 +103,45 @@ pub fn spark_map_sort(args: &[ColumnarValue]) -> Result( + entries: &StructArray, + offsets: &[i32], +) -> Result<(Vec, Vec), DataFusionError> { + 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`. + // + // `take` produces exactly the entries the visible maps refer to, so the result is indexed from + // zero. A sliced MapArray keeps its original entry offsets (a slice of two two-entry maps that + // drops the first has offsets `[2, 4]`), so the input offsets cannot be reused here -- they + // would overrun the taken entries. Rebuild them from the per-map lengths instead. + let mut global_indices: Vec = Vec::with_capacity(entries.len()); + let mut rebased_offsets: Vec = Vec::with_capacity(offsets.len()); + rebased_offsets.push(0); + + for idx in 0..offsets.len() - 1 { + let map_start = offsets[idx] as usize; + let map_end = offsets[idx + 1] as usize; + if map_end > map_start { + if SKIP_SINGLETON && map_end == map_start + 1 { + global_indices.push(map_start as u32); + } else { + let map_keys = 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); + } + + Ok((global_indices, rebased_offsets)) +} + #[cfg(test)] mod tests { use super::*; @@ -769,4 +808,226 @@ mod tests { .to_string() .contains("spark_map_sort expects Map type as argument")); } + // A non-default schema makes metadata/field-name preservation observable. + fn map_with_keys( + keys: ArrayRef, + offsets: Vec, + nulls: Option, + sorted: bool, + ) -> MapArray { + use arrow::datatypes::Field; + let values: ArrayRef = Arc::new(Int32Array::from_iter((0..keys.len()).map(|i| { + if i % 2 == 0 { + None + } else { + Some(i as i32) + } + }))); + let entries = StructArray::new( + vec![ + Arc::new(Field::new("custom_key", keys.data_type().clone(), false)), + Arc::new(Field::new("custom_value", DataType::Int32, true)), + ] + .into(), + vec![keys, values], + None, + ); + MapArray::new( + Arc::new( + Field::new("custom_entries", entries.data_type().clone(), false).with_metadata( + std::collections::HashMap::from([("source".into(), "test".into())]), + ), + ), + OffsetBuffer::new(offsets.into()), + entries, + nulls, + sorted, + ) + } + + fn assert_map_permutation(map: MapArray, permutation: Vec, offsets: Vec) { + let expected = take(map.entries(), &UInt32Array::from(permutation), None).unwrap(); + let result = spark_map_sort(&[ColumnarValue::Array(Arc::new(map.clone()))]).unwrap(); + let ColumnarValue::Array(result) = result else { + panic!("expected array") + }; + let actual = result.as_any().downcast_ref::().unwrap(); + assert_eq!(actual.entries().to_data(), expected.to_data()); + assert_eq!(actual.value_offsets(), offsets); + assert_eq!(actual.nulls(), map.nulls()); + if let Some(expected_nulls) = map.nulls() { + let actual_nulls = actual.nulls().unwrap(); + assert_eq!(actual_nulls.offset(), expected_nulls.offset()); + assert_eq!( + actual_nulls.buffer().as_ptr(), + expected_nulls.buffer().as_ptr() + ); + } + assert_eq!(actual.data_type(), map.data_type()); + } + + #[test] + fn test_singletons_and_mixed_sliced_maps_preserve_buffers_and_schema() { + use arrow::array::{Float64Array, LargeStringArray, StringViewArray}; + use arrow::buffer::NullBuffer; + let numbers = vec![99, 98, 5, 9, 1, 7, 4, 3, 2]; + let strings: Vec<_> = numbers.iter().map(i32::to_string).collect(); + let mut lists = ListBuilder::new(Int32Builder::new()); + for n in &numbers { + lists.values().append_value(*n); + lists.append(true); + } + let mut string_lists = ListBuilder::new(StringBuilder::new()); + for text in &strings { + string_lists.values().append_value(text); + string_lists.append(true); + } + for keys in [ + Arc::new(Int32Array::from(numbers.clone())) as ArrayRef, + Arc::new(Float64Array::from_iter_values( + numbers.iter().map(|n| *n as f64), + )), + Arc::new(StringArray::from(strings.clone())), + Arc::new(LargeStringArray::from(strings.clone())), + Arc::new(StringViewArray::from(strings)), + Arc::new(lists.finish()), + Arc::new(string_lists.finish()), + ] { + let map = map_with_keys( + Arc::clone(&keys), + vec![0, 2, 2, 3, 5, 6, 6, 7, 9], + Some(NullBuffer::from(vec![ + true, true, true, false, false, false, true, true, + ])), + false, + ); + let sliced = map.slice(1, 7); + assert_eq!(sliced.value_offsets()[0], 2); + assert_map_permutation( + sliced, + vec![2, 4, 3, 5, 6, 8, 7], + vec![0, 0, 1, 3, 4, 4, 5, 7], + ); + // Singleton-only batches include a valid map with a null value. + let singleton = map_with_keys(keys, (0..=9).collect(), None, false).slice(2, 5); + assert_map_permutation(singleton, vec![2, 3, 4, 5, 6], vec![0, 1, 2, 3, 4, 5]); + } + } + + #[test] + fn test_unsupported_singleton_keys_keep_arrow_errors_and_early_returns() { + use arrow::buffer::NullBuffer; + use arrow::datatypes::Field; + let structs: ArrayRef = Arc::new(StructArray::new( + vec![Arc::new(Field::new("x", DataType::Int32, false))].into(), + vec![Arc::new(Int32Array::from(vec![1]))], + None, + )); + let lists: ArrayRef = Arc::new(arrow::array::ListArray::new( + Arc::new(Field::new("item", structs.data_type().clone(), true)), + OffsetBuffer::new(vec![0, 1].into()), + Arc::clone(&structs), + None, + )); + let maps: ArrayRef = Arc::new(map_with_keys( + Arc::new(Int32Array::from(vec![1])), + vec![0, 1], + None, + false, + )); + for keys in [structs, lists, maps] { + let arrow_error = sort_to_indices( + keys.as_ref(), + Some(SortOptions { + descending: false, + nulls_first: true, + }), + None, + ) + .unwrap_err(); + let expected_error = DataFusionError::from(arrow_error).to_string(); + for validity in [None, Some(NullBuffer::from(vec![true, false]))] { + // The first row is empty; the second singleton may be null. Its physical + // entry must still produce the original error when the batch isn't all null. + let map = map_with_keys(Arc::clone(&keys), vec![0, 0, 1], validity, false); + assert_eq!( + spark_map_sort(&[ColumnarValue::Array(Arc::new(map))]) + .unwrap_err() + .to_string(), + expected_error + ); + } + for map in [ + map_with_keys(Arc::clone(&keys), vec![0], None, false), + map_with_keys(Arc::clone(&keys), vec![0, 0], None, false), + map_with_keys( + Arc::clone(&keys), + vec![0, 1], + Some(NullBuffer::from(vec![false])), + false, + ), + map_with_keys(Arc::clone(&keys), vec![0, 1], None, true), + ] { + let result = + spark_map_sort(&[ColumnarValue::Array(Arc::new(map.clone()))]).unwrap(); + let ColumnarValue::Array(result) = result else { + panic!("expected array") + }; + // Empty non-null maps rebase/take, while the other cases return early. + if map.len() == 1 && map.value_length(0) == 0 && map.null_count() == 0 { + assert_eq!(result.data_type(), map.data_type()); + assert_eq!(result.len(), 1); + } else { + assert_eq!(result.to_data(), map.to_data()); + } + } + } + } + + #[test] + fn test_singleton_float_bits_and_binary_keys() { + use arrow::array::{ + BinaryArray, BinaryViewArray, BooleanArray, FixedSizeBinaryArray, Float64Array, + LargeBinaryArray, + }; + let bytes: Vec<&[u8]> = vec![b"z", b"a", b"x", b"q"]; + let floats = [ + f64::from_bits(0x7ff8000000000042), + -0.0, + 0.0, + f64::NEG_INFINITY, + ]; + for keys in [ + Arc::new(Float64Array::from(floats.to_vec())) as ArrayRef, + Arc::new(BooleanArray::from(vec![true, false, false, true])), + Arc::new(BinaryArray::from_vec(bytes.clone())), + Arc::new(LargeBinaryArray::from_vec(bytes.clone())), + Arc::new(BinaryViewArray::from_iter_values(bytes.clone())), + Arc::new(FixedSizeBinaryArray::try_from_iter(bytes.into_iter()).unwrap()), + ] { + let map = map_with_keys(keys, vec![0, 1, 2, 3, 4], None, false); + assert_map_permutation(map.clone(), vec![0, 1, 2, 3], vec![0, 1, 2, 3, 4]); + if matches!(map.key_type(), DataType::Float64) { + let ColumnarValue::Array(result) = + spark_map_sort(&[ColumnarValue::Array(Arc::new(map))]).unwrap() + else { + panic!("expected array") + }; + let actual = result.as_any().downcast_ref::().unwrap(); + let actual = actual + .keys() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + actual + .values() + .iter() + .map(|x| x.to_bits()) + .collect::>(), + floats.iter().map(|x| x.to_bits()).collect::>() + ); + } + } + } } From cf22afe92e5d504adf72917a3a12473f9fb3c578 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sat, 12 Sep 2026 11:47:15 -0700 Subject: [PATCH 2/2] docs: link singleton map performance audit to PR --- docs/source/contributor-guide/expression-audits/map_funcs.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/contributor-guide/expression-audits/map_funcs.md b/docs/source/contributor-guide/expression-audits/map_funcs.md index bdaff4bded8..a2facc5a584 100644 --- a/docs/source/contributor-guide/expression-audits/map_funcs.md +++ b/docs/source/contributor-guide/expression-audits/map_funcs.md @@ -66,7 +66,7 @@ ## map_sort -- Performance (tuned locally 2026-09-12; PR pending): skip Arrow sort dispatch for eligible flat singleton keys, with a batch check and specialized fallback loop for batches without singletons. In the local DataFusion 55.0.0 development cohort, matched singleton normalization measured 19–22x faster in the full run and 18.4x in an independent forward-order confirmation. Benchmarks: `native/spark-expr/benches/map_sort.rs`, `hash.rs`, and `common/matched_maps.rs`; 92 cases cover normalization, hashing, combined execution, nulls, slices, mixed cardinalities, and long Unicode values. Flagged regressions did not remain stable through independent and reversed-order confirmation. +- Performance (tuned locally 2026-09-12; [PR #5887](https://github.com/apache/datafusion-comet/pull/5887)): skip Arrow sort dispatch for eligible flat singleton keys, with a batch check and specialized fallback loop for batches without singletons. In the local DataFusion 55.0.0 development cohort, matched singleton normalization measured 19–22x faster in the full run and 18.4x in an independent forward-order confirmation. Benchmarks: `native/spark-expr/benches/map_sort.rs`, `hash.rs`, and `common/matched_maps.rs`; 92 cases cover normalization, hashing, combined execution, nulls, slices, mixed cardinalities, and long Unicode values. Flagged regressions did not remain stable through independent and reversed-order confirmation. ## map_values