diff --git a/docs/source/contributor-guide/expression-audits/map_funcs.md b/docs/source/contributor-guide/expression-audits/map_funcs.md index a2facc5a58..fc9c5584c9 100644 --- a/docs/source/contributor-guide/expression-audits/map_funcs.md +++ b/docs/source/contributor-guide/expression-audits/map_funcs.md @@ -66,6 +66,7 @@ ## map_sort +- Performance (tuned locally 2026-09-13; [PR #5901](https://github.com/apache/datafusion-comet/pull/5901), related to [#5900](https://github.com/apache/datafusion-comet/issues/5900)): reuse per-batch prefix-tuple sorting scratch for multi-entry `Utf8`/`Int32` maps and bulk-fill all-empty offsets, preserving the singleton path from [#5887](https://github.com/apache/datafusion-comet/pull/5887). Against upstream including #5887, matched 2–10-entry forward normalization was about 3x faster; 2–50-entry maps improved 28–39% in the full run and 33–39% in independent paired confirmation. Benchmarks: `native/spark-expr/benches/map_sort.rs`, `hash.rs`, and `common/matched_maps.rs`. - 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 diff --git a/native/spark-expr/benches/common/matched_maps.rs b/native/spark-expr/benches/common/matched_maps.rs index 2bdd448ddc..13d92798d3 100644 --- a/native/spark-expr/benches/common/matched_maps.rs +++ b/native/spark-expr/benches/common/matched_maps.rs @@ -379,3 +379,120 @@ pub fn bench_regression_maps(c: &mut Criterion, stage: Stage) { } group.finish(); } + +// Int-value variants retain the shared regression fixture conventions. +pub fn bench_multi_entry_int_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!("multi_entry_int_maps/{}", stage.name())); + group.throughput(Throughput::Elements(ROWS as u64)); + for (name, null_every, mixed, long) in [ + ("wide_no_null", 0, false, false), + ("wide_sparse_null", 100, false, false), + ("wide_dense_null", 2, false, false), + ("wide_long_unicode", 0, false, true), + ("wide_long_unicode_dense_null", 2, false, true), + ("mixed_dense_null", 2, true, false), + ] { + let mut builder = MapBuilder::new( + Some(crate::common::map_field_names()), + StringBuilder::new(), + Int32Builder::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 { + 2 + c1(row) as usize % 49 + }; + 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((row * 100 + entry) as i32); + } + } + 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 157324e337..9e5d8720d2 100644 --- a/native/spark-expr/benches/hash.rs +++ b/native/spark-expr/benches/hash.rs @@ -116,7 +116,9 @@ 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_multi_entry_int_maps(c, matched_maps::Stage::HashOnly); matched_maps::bench_regression_maps(c, matched_maps::Stage::NormalizeHash); + matched_maps::bench_multi_entry_int_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 8c5b4c9ba1..20f4f5c7ce 100644 --- a/native/spark-expr/benches/map_sort.rs +++ b/native/spark-expr/benches/map_sort.rs @@ -145,6 +145,7 @@ fn bench_unsupported_singleton(c: &mut Criterion) { 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); + matched_maps::bench_multi_entry_int_maps(c, matched_maps::Stage::NormalizeOnly); } criterion_group!( diff --git a/native/spark-expr/src/map_funcs/map_sort.rs b/native/spark-expr/src/map_funcs/map_sort.rs index 58de949dfd..b07ae8bd37 100644 --- a/native/spark-expr/src/map_funcs/map_sort.rs +++ b/native/spark-expr/src/map_funcs/map_sort.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{Array, ArrayRef, MapArray, StructArray, UInt32Array}; +use arrow::array::{Array, ArrayRef, MapArray, StringArray, StructArray, UInt32Array}; use arrow::buffer::OffsetBuffer; use arrow::compute::{sort_to_indices, take, SortOptions}; use arrow::datatypes::DataType; @@ -124,12 +124,39 @@ fn map_sort_indices( let mut rebased_offsets: Vec = Vec::with_capacity(offsets.len()); rebased_offsets.push(0); + if offsets[offsets.len() - 1] == offsets[0] { + // Empty visible slices still need take/rebasing in the caller. + rebased_offsets.resize(offsets.len(), 0); + return Ok((global_indices, rebased_offsets)); + } + + // Restrict allocation reuse to the measured map shape. + // Other types retain Arrow's dispatch and validation. + let string_keys = if entries.column(1).data_type() == &DataType::Int32 { + entries + .column(0) + .as_any() + .downcast_ref::() + .filter(|keys| keys.null_count() == 0) + } else { + None + }; + let mut scratch = Vec::new(); + 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 if let Some(keys) = string_keys { + append_string_map_indices( + keys, + map_start, + map_end, + &mut scratch, + &mut global_indices, + ); } 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)?; @@ -142,6 +169,56 @@ fn map_sort_indices( Ok((global_indices, rebased_offsets)) } +// Keep Arrow's (index, four-byte prefix, length) sort representation and comparator. +// In particular, sorting u32 indices instead can change the permutation of equal keys +// because Rust's unstable sort specializes by element layout. Reusing this buffer avoids +// a key-array slice and Arrow's per-row index/tuple/output allocations without changing +// that behavior. Capacity grows only to the largest row, not the whole entries array. +fn append_string_map_indices( + keys: &StringArray, + start: usize, + end: usize, + scratch: &mut Vec<(u32, u32, u64)>, + global_indices: &mut Vec, +) { + scratch.clear(); + scratch.extend((start..end).map(|index| { + // SAFETY: MapArray offsets bound each row within its entries/key array. + let bytes = unsafe { keys.value_unchecked(index) }.as_bytes(); + let prefix = if bytes.len() >= 4 { + // SAFETY: At least four initialized bytes are available; alignment is not required. + u32::from_be(unsafe { std::ptr::read_unaligned(bytes.as_ptr().cast::()) }) + } else if bytes.is_empty() { + 0 + } else { + let mut prefix = 0u32; + for &byte in bytes { + prefix = (prefix << 8) | u32::from(byte); + } + prefix << (8 * (4 - bytes.len())) + }; + (index as u32, prefix, bytes.len() as u64) + })); + scratch.sort_unstable_by(|a, b| { + let order = a.1.cmp(&b.1); + if !order.is_eq() { + return order; + } + if a.2 < 4 || b.2 < 4 { + let order = a.2.cmp(&b.2); + if !order.is_eq() { + return order; + } + } + // SAFETY: Both indices were generated from this map's valid entry range above. + unsafe { + keys.value_unchecked(a.0 as usize) + .cmp(keys.value_unchecked(b.0 as usize)) + } + }); + global_indices.extend(scratch.iter().map(|entry| entry.0)); +} + #[cfg(test)] mod tests { use super::*; @@ -1030,4 +1107,177 @@ mod tests { } } } + + // A non-default schema makes metadata/field-name preservation observable. + fn map_with_int_values( + 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_sorted_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_multi_entry_unicode_duplicates_and_sliced_null_maps() { + use arrow::buffer::NullBuffer; + let prefix = "資料é".repeat(128); + let keys: ArrayRef = Arc::new(StringArray::from(vec![ + "z".into(), + "y".into(), + "one".into(), + "中".into(), + "a".into(), + "é".into(), + "a".into(), + "\0".into(), + "".into(), + format!("{prefix}B"), + format!("{prefix}A"), + format!("{prefix}A"), + "😀".into(), + ])); + let map = map_with_int_values( + keys, + vec![0, 2, 2, 3, 9, 9, 13], + Some(NullBuffer::from(vec![true, true, true, false, false, true])), + false, + ); + assert_sorted_permutation( + map.slice(1, 5), + vec![2, 8, 7, 4, 6, 5, 3, 10, 11, 9, 12], + vec![0, 0, 1, 7, 7, 11], + ); + } + + #[test] + fn test_multi_entry_permutation_matches_arrow_for_equal_keys() { + // Arrow's unstable sort does not promise stable duplicate ordering. Check the exact + // current kernel permutation, including larger rows where insertion sort no longer applies. + let mut seed = 42u64; + for count in [2, 3, 6, 16, 20, 21, 26, 32, 50, 64, 128, 257] { + for distinct in [1, 3, 17, 1000] { + let mut text = vec!["unused".to_owned()]; + for _ in 0..count { + seed ^= seed << 13; + seed ^= seed >> 7; + seed ^= seed << 17; + text.push(format!("{}", seed % distinct)); + } + let keys: ArrayRef = Arc::new(StringArray::from(text)); + let local = sort_to_indices( + &keys.slice(1, count), + Some(SortOptions { + descending: false, + nulls_first: true, + }), + None, + ) + .unwrap(); + let expected = local.values().iter().map(|i| i + 1).collect(); + let map = map_with_int_values(keys, vec![0, 1, count as i32 + 1], None, false); + assert_sorted_permutation(map.slice(1, 1), expected, vec![0, count as i32]); + } + } + } + + #[test] + fn test_multi_entry_unsupported_keys_preserve_error_under_null_map() { + use arrow::buffer::NullBuffer; + use arrow::datatypes::Field; + let keys: ArrayRef = Arc::new(StructArray::new( + vec![Arc::new(Field::new("k", DataType::Int32, false))].into(), + vec![Arc::new(Int32Array::from(vec![2, 1]))], + None, + )); + let expected = DataFusionError::from( + sort_to_indices( + keys.as_ref(), + Some(SortOptions { + descending: false, + nulls_first: true, + }), + None, + ) + .unwrap_err(), + ) + .to_string(); + let map = map_with_int_values( + keys, + vec![0, 0, 2], + Some(NullBuffer::from(vec![true, false])), + false, + ); + assert_eq!( + spark_map_sort(&[ColumnarValue::Array(Arc::new(map))]) + .unwrap_err() + .to_string(), + expected + ); + } + + #[test] + fn test_empty_visible_maps_rebase_and_preserve_nulls() { + use arrow::buffer::NullBuffer; + for keys in [ + Arc::new(StringArray::from(vec!["z", "a"])) as ArrayRef, + Arc::new(Int32Array::from(vec![2, 1])), + ] { + let map = map_with_int_values( + keys, + vec![0, 2, 2, 2], + Some(NullBuffer::from(vec![true, true, false])), + false, + ); + assert_sorted_permutation(map.slice(1, 2), vec![], vec![0, 0, 0]); + } + } }