diff --git a/native/spark-expr/benches/common/matched_maps.rs b/native/spark-expr/benches/common/matched_maps.rs new file mode 100644 index 00000000000..0be1e0c04e1 --- /dev/null +++ b/native/spark-expr/benches/common/matched_maps.rs @@ -0,0 +1,258 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Matched map inputs from CometShuffleBenchmark (PR #5788). + +use arrow::array::{ + Array, ArrayRef, Int32Array, Int32Builder, MapArray, MapBuilder, StringBuilder, StructArray, +}; +use arrow::datatypes::Field; +use criterion::{BenchmarkId, Criterion, Throughput}; +use datafusion::physical_plan::ColumnarValue; +use datafusion_comet_spark_expr::{murmur3::create_murmur3_hashes, spark_map_sort}; +use std::{hint::black_box, sync::Arc}; + +pub const ROWS: usize = 8192; + +#[derive(Clone, Copy)] +// Each benchmark executable registers only its own subset of stages. +#[allow(dead_code)] +pub enum Stage { + HashOnly, + NormalizeOnly, + NormalizeHash, +} + +impl Stage { + fn name(self) -> &'static str { + match self { + Self::HashOnly => "hash_only", + Self::NormalizeOnly => "normalize_only", + Self::NormalizeHash => "normalize_hash", + } + } +} + +#[derive(Clone, Copy)] +enum Shape { + Map, + StructMapInt, +} + +impl Shape { + fn name(self) -> &'static str { + match self { + Self::Map => "map", + Self::StructMapInt => "struct_map_int", + } + } +} + +// Scala Long arithmetic wraps; pmod is applied to the signed result of mix64. +fn c1(row: usize) -> i32 { + let mut z = (row as u64).wrapping_add(0x9e3779b97f4a7c15); + z = (z ^ (z >> 30)).wrapping_mul(0xbf58476d1ce4e5b9); + z = (z ^ (z >> 27)).wrapping_mul(0x94d049bb133111eb); + ((z ^ (z >> 31)) as i64).rem_euclid(1_000_000) as i32 +} + +fn maps(max_entries: usize, reversed: bool) -> ArrayRef { + let mut builder = MapBuilder::new( + Some(crate::common::map_field_names()), + StringBuilder::new(), + Int32Builder::new(), + ); + for row in 0..ROWS { + let value = c1(row); + let count = if max_entries <= 1 { + max_entries + } else { + 2 + value as usize % (max_entries - 1) + }; + for idx in 0..count { + // Singleton is MAP(CAST(c1 AS STRING), c1); larger maps use sequence(1, count). + let x = if max_entries == 1 { + 0 + } else if reversed { + count - idx + } else { + idx + 1 + }; + let entry = value + x as i32; + builder.keys().append_value(entry.to_string()); + builder.values().append_value(entry); + } + builder.append(true).unwrap(); + } + Arc::new(builder.finish()) +} + +fn normalize(args: &[ColumnarValue]) -> ArrayRef { + match spark_map_sort(args).unwrap() { + ColumnarValue::Array(array) => array, + _ => panic!("expected array"), + } +} + +fn wrap(map: ArrayRef, ints: &ArrayRef) -> ArrayRef { + Arc::new(StructArray::new( + vec![ + Arc::new(Field::new("m", map.data_type().clone(), true)), + Arc::new(Field::new("i", ints.data_type().clone(), true)), + ] + .into(), + vec![map, Arc::clone(ints)], + None, + )) +} + +fn hashes(array: &ArrayRef) -> Vec { + let mut hashes = vec![42; ROWS]; + create_murmur3_hashes(std::slice::from_ref(array), &mut hashes).unwrap(); + hashes +} + +/// Shared registration keeps every stage's data and correctness checks identical. +/// Input construction and checks are untimed. Hash storage is allocated once; resetting it +/// to seed 42 is timed. Normalization output allocation/drop, and struct reconstruction in +/// normalize_hash, are timed. normalize_only measures mapsort itself (no struct wrapper). +pub fn bench_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!("matched_maps/{}", stage.name())); + group.throughput(Throughput::Elements(ROWS as u64)); + // Zero entries means ROWS non-null empty maps, not an empty batch. + for max_entries in [0, 1, 10, 50] { + let forward = maps(max_entries, false); + let reversed = maps(max_entries, true); + let normalized = normalize(&[ColumnarValue::Array(Arc::clone(&forward))]); + let normalized_reverse = normalize(&[ColumnarValue::Array(Arc::clone(&reversed))]); + assert_eq!(normalized.to_data(), normalized_reverse.to_data()); + // Independently check lexical ordering and key/value alignment; numeric order does + // not in general imply string order for unpadded decimal keys. + let expected = normalized.as_any().downcast_ref::().unwrap(); + let keys = expected + .keys() + .as_any() + .downcast_ref::() + .unwrap(); + let values = expected + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(expected.len(), ROWS); + for (row, offsets) in expected.value_offsets().windows(2).enumerate() { + let base = c1(row); + let count = if max_entries <= 1 { + max_entries + } else { + 2 + base as usize % (max_entries - 1) + }; + assert_eq!((offsets[1] - offsets[0]) as usize, count); + let mut actual_values: Vec<_> = (offsets[0] as usize..offsets[1] as usize) + .map(|i| values.value(i)) + .collect(); + actual_values.sort_unstable(); + let start = if max_entries == 1 { base } else { base + 1 }; + assert_eq!( + actual_values, + (start..start + count as i32).collect::>() + ); + for i in offsets[0] as usize..offsets[1] as usize { + assert_eq!(keys.value(i), values.value(i).to_string()); + if i > offsets[0] as usize { + assert!(keys.value(i - 1) < keys.value(i)); + } + } + } + // This fixed fixture demonstrates why normalization is needed. Do not require + // every row to differ: Murmur3 collisions are possible for arbitrary inputs. + if max_entries > 1 { + assert_ne!(hashes(&forward), hashes(&reversed)); + assert_ne!( + hashes(&wrap(Arc::clone(&forward), &ints)), + hashes(&wrap(Arc::clone(&reversed), &ints)) + ); + } else { + assert_eq!(forward.to_data(), reversed.to_data()); + } + let mut field_hashes = vec![42; ROWS]; + create_murmur3_hashes( + &[Arc::clone(&normalized), Arc::clone(&ints)], + &mut field_hashes, + ) + .unwrap(); + assert_eq!(field_hashes, hashes(&wrap(Arc::clone(&normalized), &ints))); + for (order, raw) in [("forward", forward), ("reversed", reversed)] { + let args = [ColumnarValue::Array(raw)]; + for shape in [Shape::Map, Shape::StructMapInt] { + // mapsort is identical for either enclosing shape; measure it once. + if matches!((stage, shape), (Stage::NormalizeOnly, Shape::StructMapInt)) { + continue; + } + group.bench_function( + BenchmarkId::new(shape.name(), format!("{max_entries}/{order}")), + |b| match stage { + Stage::NormalizeOnly => b.iter(|| { + black_box(normalize(black_box(&args))); + }), + Stage::HashOnly => { + // Both order labels intentionally hash the same pre-normalized + // bytes. They are repeated-measurement controls, not a test of + // raw hash order independence. + 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]; + // Dispatch outside timing, including whether a struct is rebuilt. + match shape { + Shape::Map => b.iter(|| { + buffer.fill(42); + let map = normalize(black_box(&args)); + create_murmur3_hashes(std::slice::from_ref(&map), &mut buffer) + .unwrap(); + black_box(&buffer); + }), + Shape::StructMapInt => b.iter(|| { + buffer.fill(42); + let key = wrap(normalize(black_box(&args)), &ints); + create_murmur3_hashes(std::slice::from_ref(&key), &mut buffer) + .unwrap(); + black_box(&buffer); + }), + } + } + }, + ); + } + } + } + group.finish(); +} diff --git a/native/spark-expr/benches/common/mod.rs b/native/spark-expr/benches/common/mod.rs index ff484ecfe11..cc0c6f275fe 100644 --- a/native/spark-expr/benches/common/mod.rs +++ b/native/spark-expr/benches/common/mod.rs @@ -15,9 +15,10 @@ // specific language governing permissions and limitations // under the License. -//! Helpers shared by the cast-from-string benchmarks, pulled in with +//! Helpers shared by expression benchmarks, pulled in with //! `#[path = "common/mod.rs"] mod common;`. This lives in a subdirectory so that Cargo's bench //! auto-discovery, which only looks at `benches/*.rs`, does not treat it as a bench target. +//! This directory also holds standalone modules such as `matched_maps.rs`, included directly. #![allow(dead_code)] use arrow::array::{ @@ -284,3 +285,12 @@ pub fn list_arrays( ), ] } + +/// Field names shared by the hash and map-sort benchmark inputs. +pub fn map_field_names() -> arrow::array::MapFieldNames { + arrow::array::MapFieldNames { + entry: "entries".into(), + key: "key".into(), + value: "value".into(), + } +} diff --git a/native/spark-expr/benches/hash.rs b/native/spark-expr/benches/hash.rs index 20ed6eff56d..4d8d4026cfd 100644 --- a/native/spark-expr/benches/hash.rs +++ b/native/spark-expr/benches/hash.rs @@ -13,7 +13,7 @@ // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations -// under the License.use arrow::array::{ArrayRef, BooleanBuilder, Int32Builder, RecordBatch, StringBuilder}; +// under the License. //! Benchmarks for the Spark-compatible hash kernels, which back the `hash` and `xxhash64` //! expressions and, since #5567, native shuffle hash partitioning. @@ -28,7 +28,7 @@ //! to that macro shows up here. use arrow::array::builder::{Int32Builder, ListBuilder, MapBuilder, StringBuilder, StructBuilder}; -use arrow::array::{ArrayRef, Int32Array, ListArray, MapFieldNames, StringArray, StructArray}; +use arrow::array::{ArrayRef, Int32Array, ListArray, StringArray, StructArray}; use arrow::buffer::OffsetBuffer; use arrow::datatypes::{DataType, Field, Fields}; use criterion::{criterion_group, criterion_main, Criterion}; @@ -36,6 +36,10 @@ use datafusion_comet_spark_expr::murmur3::create_murmur3_hashes; use std::hint::black_box; use std::sync::Arc; +mod common; +#[path = "common/matched_maps.rs"] +mod matched_maps; + const NUM_ROWS: usize = 8192; fn struct_fields() -> Fields { @@ -132,11 +136,7 @@ fn skewed_list_of_struct(num_rows: usize, long_len: usize) -> ArrayRef { /// `map`: keys and values are hashed entry by entry. fn maps(num_rows: usize, entries: usize) -> ArrayRef { let mut mb = MapBuilder::new( - Some(MapFieldNames { - entry: "entries".into(), - key: "key".into(), - value: "value".into(), - }), + Some(common::map_field_names()), StringBuilder::new(), Int32Builder::new(), ); @@ -202,11 +202,7 @@ fn list_of_list(num_rows: usize, outer: usize, inner: usize) -> ArrayRef { /// cover, so the value array is hashed recursively instead. fn map_of_struct(num_rows: usize, entries: usize) -> ArrayRef { let mut mb = MapBuilder::new( - Some(MapFieldNames { - entry: "entries".into(), - key: "key".into(), - value: "value".into(), - }), + Some(common::map_field_names()), StringBuilder::new(), struct_builder(), ); @@ -268,5 +264,17 @@ fn bench(c: &mut Criterion) { group.finish(); } -criterion_group!(benches, bench); +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); + c.bench_function("matched_maps/hash_buffer_seed_reset", |b| { + let mut hashes = vec![42u32; matched_maps::ROWS]; + b.iter(|| { + hashes.fill(42); + black_box(&hashes); + }); + }); +} + +criterion_group!(benches, bench, bench_matched_maps); criterion_main!(benches); diff --git a/native/spark-expr/benches/map_sort.rs b/native/spark-expr/benches/map_sort.rs index f2af0b7401d..f32b9eb5bf2 100644 --- a/native/spark-expr/benches/map_sort.rs +++ b/native/spark-expr/benches/map_sort.rs @@ -18,28 +18,24 @@ //! Benchmarks for spark_map_sort. use arrow::array::builder::{Int32Builder, MapBuilder, StringBuilder}; -use arrow::array::{ArrayRef, MapArray, MapFieldNames}; +use arrow::array::{ArrayRef, MapArray}; use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; use datafusion::physical_plan::ColumnarValue; use datafusion_comet_spark_expr::spark_map_sort; use std::hint::black_box; use std::sync::Arc; -const BATCH_SIZE: usize = 8192; +mod common; +#[path = "common/matched_maps.rs"] +mod matched_maps; -fn map_field_names() -> MapFieldNames { - MapFieldNames { - entry: "entries".into(), - key: "key".into(), - value: "value".into(), - } -} +const BATCH_SIZE: usize = 8192; /// Build a MapArray with `BATCH_SIZE` rows where each map has `entries_per_map` entries. /// Keys are integers in reverse order so every map needs a real sort. fn build_int_key_map(entries_per_map: usize) -> MapArray { let mut builder = MapBuilder::new( - Some(map_field_names()), + Some(common::map_field_names()), Int32Builder::new(), Int32Builder::new(), ); @@ -59,7 +55,7 @@ fn build_int_key_map(entries_per_map: usize) -> MapArray { /// Same shape as `build_int_key_map` but with string keys. fn build_string_key_map(entries_per_map: usize) -> MapArray { let mut builder = MapBuilder::new( - Some(map_field_names()), + Some(common::map_field_names()), StringBuilder::new(), Int32Builder::new(), ); @@ -103,5 +99,9 @@ fn bench_map_sort(c: &mut Criterion) { group.finish(); } -criterion_group!(benches, bench_map_sort); +fn bench_matched_maps(c: &mut Criterion) { + matched_maps::bench_maps(c, matched_maps::Stage::NormalizeOnly); +} + +criterion_group!(benches, bench_map_sort, bench_matched_maps); criterion_main!(benches);