diff --git a/docs/source/contributor-guide/expression-audits/hash_funcs.md b/docs/source/contributor-guide/expression-audits/hash_funcs.md index c3ad4efe7b2..a294069df14 100644 --- a/docs/source/contributor-guide/expression-audits/hash_funcs.md +++ b/docs/source/contributor-guide/expression-audits/hash_funcs.md @@ -34,7 +34,8 @@ - Spark 3.5.8 (audited 2026-05-27): baseline. `Murmur3Hash(children, seed) extends HashExpression[Int]`; produces a Murmur3 hash with a configurable Int seed and `IntegerType` result. Comet routes via `CometMurmur3Hash` to the native `murmur3_hash` UDF. - Spark 4.0.1 (audited 2026-05-27): semantics unchanged; some inner helper refactors only. - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. -- Known limitation: `DecimalType` children with precision > 18 fall back because Spark hashes them through Java `BigDecimal`; `TimeType` (Spark 4.0+) is also unsupported. The same limitations apply to `xxhash64`, `sha1`, `sha2` through the shared `HashUtils`. +- Decimals with precision 19–38, including nested decimals, hash natively using the minimal signed big-endian bytes of the unscaled value, matching Spark's `BigInteger.toByteArray()`. Precision ≤18 retains unscaled-long hashing. `TimeType` remains unsupported through the shared `HashUtils`. +- Performance (tuned 2026-09-28, issue #5994): typed wide-decimal list scans avoid per-element Arrow slices and recursive dispatch. Local Criterion measurements were 5.0–14.2x faster across 8,192-row `decimal(38,0)` List, LargeList and FixedSizeList inputs with 2/32 elements and no/sparse/dense nulls; paired scalar remeasurement found no significant slowdown. Benchmark: `benches/decimal_hash.rs`. ## md5 @@ -66,10 +67,12 @@ ## xxhash64 +- Performance (tuned 2026-09-28, issue #5994): the shared typed wide-decimal list path was 2.1–15.8x faster in local Criterion measurements across 8,192-row `decimal(38,0)` List, LargeList and FixedSizeList inputs with 2/32 elements and no/sparse/dense nulls; paired scalar remeasurement found no significant slowdown. Benchmark: `benches/decimal_hash.rs`. + - Spark 3.4.3 (audited 2026-05-27): identical to 3.5.8. - Spark 3.5.8 (audited 2026-05-27): baseline. `XxHash64(children, seed) extends HashExpression[Long]`; produces an xxHash64 hash with a configurable Long seed and `LongType` result. Comet routes via `CometXxHash64` to the native `xxhash64` UDF. - Spark 4.0.1 (audited 2026-05-27): semantics unchanged. - Spark 4.1.1 (audited 2026-05-27): identical to 4.0.1. -- Upstream (2026-09-15): Comet's native `xxhash64` UDF delegates compatible arguments at Spark's default seed (`42`) to `datafusion-spark::SparkXxhash64`. Differential tests in `native/spark-expr/src/hash_funcs/xxhash64_diff.rs` compare Comet's kernel against `SparkXxhash64` for primitives, both Decimal128 widths, dictionaries, lists, maps, and nested combinations, and record where the two differ on a NaN whose bits are not canonical. The Comet kernel is retained for a non-default seed (`SparkXxhash64` hardcodes 42), `Struct` (upstream does not push a parent null mask into children; see #5753), a `Dictionary` nested in a list/map (upstream restarts those hashes from 42), `Time64`, and `Float32`/`Float64`, including values nested in lists and maps (upstream hashes the raw bits of a NaN, where Spark hashes every NaN as the canonical NaN; see [apache/datafusion#25913](https://github.com/apache/datafusion/issues/25913)). `murmur3` is unchanged; `create_xxhash64_hashes` remains for `approx_count_distinct` and the fallback path. +- Native routing: Comet's native `xxhash64` UDF delegates compatible arguments at Spark's default seed (`42`) to `datafusion-spark::SparkXxhash64`. Differential tests in `native/spark-expr/src/hash_funcs/xxhash64_diff.rs` compare Comet's kernel against `SparkXxhash64` for primitives, narrow decimals, dictionaries, lists, maps, and nested combinations, and record where the two differ on a NaN whose bits are not canonical. Wide-decimal tests verify that the expression stays on Comet's kernel: the pinned DataFusion implementation uses fixed-width little-endian bytes rather than Spark's minimal signed big-endian bytes. The Comet kernel is retained for wide decimals (also within nested types), a non-default seed (`SparkXxhash64` hardcodes 42), `Struct` (upstream does not push a parent null mask into children; see #5753), a `Dictionary` nested in a list/map (upstream restarts those hashes from 42), `Time64`, and `Float32`/`Float64`, including values nested in lists and maps (upstream hashes the raw bits of a NaN, where Spark hashes every NaN as the canonical NaN; see [apache/datafusion#25913](https://github.com/apache/datafusion/issues/25913)). `create_xxhash64_hashes` remains for `approx_count_distinct` and the non-delegated native path. [Spark Expression Support]: ../../user-guide/latest/expressions.md diff --git a/docs/source/contributor-guide/jvm_shuffle.md b/docs/source/contributor-guide/jvm_shuffle.md index 1d8552b0f0f..474764b90e2 100644 --- a/docs/source/contributor-guide/jvm_shuffle.md +++ b/docs/source/contributor-guide/jvm_shuffle.md @@ -52,9 +52,7 @@ JVM shuffle (`CometColumnarExchange`) is used instead of native shuffle (`CometE range key always falls back here. `HashPartitioning` keys must be primitive only by default: setting `spark.comet.shuffle.native.partitioning.hash.nested.enabled` to `true` keeps struct and array keys, and map keys on Spark 4.0 and later, on the native path. The config defaults to - `false`, so a complex hash key falls back to JVM columnar shuffle unless it is enabled. - Decimal hash keys with precision greater than 18, including nested decimal leaves, also use - JVM shuffle when there is more than one partition and the shuffle mode allows it. See + `false`, so a complex hash key falls back to JVM columnar shuffle unless it is enabled. See [Supported partition key types](native_shuffle.md#when-native-shuffle-is-used) for the exact rules. Complex types are fully supported as data columns in both implementations. diff --git a/docs/source/contributor-guide/native_shuffle.md b/docs/source/contributor-guide/native_shuffle.md index 80d57fe8081..4483d40dc41 100644 --- a/docs/source/contributor-guide/native_shuffle.md +++ b/docs/source/contributor-guide/native_shuffle.md @@ -56,9 +56,9 @@ Native shuffle (`CometExchange`) is selected when all of the following condition `spark.comet.convert.shuffleInput.enabled` is set. Then a shuffle that JVM shuffle would take uses native shuffle instead, provided native shuffle supports the partitioning and the columns, and a `CometSparkToColumnarExec` converts the child's rows to Arrow for it. A shuffle - that hashes a string, a value computed from a string, or a decimal wider than 18 digits stays - on JVM shuffle, because native shuffle would not put every row in the partition Spark's - partitioner does (`CometShuffleExchangeExec.convertsInputForNativeShuffle`). + that hashes a string or a value computed from a string stays on JVM shuffle, because native + shuffle would not put every row in the partition Spark's partitioner does + (`CometShuffleExchangeExec.convertsInputForNativeShuffle`). 3. **Supported partitioning type**: Native shuffle supports: - `HashPartitioning` @@ -83,9 +83,6 @@ Native shuffle (`CometExchange`) is selected when all of the following condition Spark's `mapsort` normalization makes physical entry order irrelevant. A collated string at any depth still disqualifies the key. The config defaults to `false` pending measurement of the nested hashing paths, so by default a complex hash key falls back to JVM shuffle. - Decimal leaves with precision greater than 18 also disqualify a hash key when there is more - than one output partition: native decimal hashing differs from Spark's. A one-partition hash - exchange stays native because the writer serializes it as `SinglePartition` without hashing. ## Architecture diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index 2fd8333a79b..9bcafbcc746 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -357,12 +357,12 @@ to Spark ([#2837](https://github.com/apache/datafusion-comet/issues/2837)). Enab | Function | Status | Implementation | Notes | | --- | --- | --- | --- | | `crc32` | ✅ | Native | | -| `hash` | ✅ | Native | | +| `hash` | ✅ | Native | Includes decimals with precision 19–38 and nested decimals | | `md5` | ✅ | Native | | | `sha` | ✅ | Native | | | `sha1` | ✅ | Native | | | `sha2` | ✅ | Native | | -| `xxhash64` | ✅ | Native | | +| `xxhash64` | ✅ | Native | Includes decimals with precision 19–38 and nested decimals | --- diff --git a/docs/source/user-guide/latest/tuning/shuffle.md b/docs/source/user-guide/latest/tuning/shuffle.md index 31aca7bd024..f82d29b57ed 100644 --- a/docs/source/user-guide/latest/tuning/shuffle.md +++ b/docs/source/user-guide/latest/tuning/shuffle.md @@ -53,17 +53,6 @@ scalar types. Hash partitioning keys must be scalar types unless and later) map keys. That setting is disabled by default until the performance of the nested hashing paths has been measured. Columns that are not partitioning keys may contain complex types like maps, structs, and arrays. -Hash partitioning into multiple partitions on decimal keys with precision greater than 18 falls back because native -hashing does not match Spark's partition assignments. Mixing the two partitioners can silently lose join rows. -The restriction also applies to decimal leaves in nested hash keys. Wider decimals remain supported as payload -columns, range partitioning keys, and in single-partition shuffles, including one-partition hash exchanges. - -With `spark.comet.shuffle.mode=auto`, Comet uses Columnar Shuffle when eligible, adding a columnar-to-row-to-columnar -conversion. With `native`, it uses Spark shuffle, which can also restore aggregates such as `collect_list` and -`collect_set` to Spark to preserve buffer compatibility. Celeborn uses Spark/Celeborn shuffle for unsupported native -keys because it has no Comet columnar fallback. Spark-compatible native decimal hashing is tracked in -[#5994](https://github.com/apache/datafusion-comet/issues/5994). - ### Columnar (JVM) Shuffle Comet Columnar shuffle is JVM-based and supports `HashPartitioning`, `RoundRobinPartitioning`, `RangePartitioning`, and @@ -82,9 +71,9 @@ converts the child's rows to Arrow, as the `spark.comet.convert.*` settings do f batches natively. With hash partitioning that is usually faster than columnar shuffle, while with range partitioning the two come out about even. This applies only where native shuffle supports the partitioning and Comet can convert all of the child's columns, which rules out calendar intervals, and arrays and maps other than `array` and -`map`. A shuffle that hashes a decimal wider than 18 digits, or a string or a value computed from a -string, stays on columnar shuffle, so that it partitions rows as Spark does. The converted batches are allocated on -the JVM outside Comet's memory pools. The setting is disabled by default. +`map`. A shuffle that hashes a string or a value computed from a string stays on columnar shuffle, so +that it partitions rows as Spark does. The converted batches are allocated on the JVM outside Comet's memory pools. +The setting is disabled by default. ### Automatic Revert to Spark Shuffle diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 3011f1c5c6f..4314dd0d7e1 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -75,6 +75,10 @@ path = "src/lib.rs" name = "hash" harness = false +[[bench]] +name = "decimal_hash" +harness = false + [[bench]] name = "hash_alloc" harness = false diff --git a/native/spark-expr/benches/decimal_hash.rs b/native/spark-expr/benches/decimal_hash.rs new file mode 100644 index 00000000000..01e7580cc87 --- /dev/null +++ b/native/spark-expr/benches/decimal_hash.rs @@ -0,0 +1,117 @@ +// 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. + +//! Wide-decimal scalar controls and typed list hashing, including sliced inputs, +//! short/long arrays and no/sparse/dense nulls. Both public SQL hash paths are measured. + +use arrow::array::{ + Array, ArrayRef, Decimal128Array, FixedSizeListArray, LargeListArray, ListArray, +}; +use arrow::buffer::{NullBuffer, OffsetBuffer}; +use arrow::datatypes::Field; +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; +use datafusion::common::ScalarValue; +use datafusion::logical_expr::ColumnarValue; +use datafusion_comet_spark_expr::{spark_murmur3_hash, spark_xxhash64}; +use std::{hint::black_box, sync::Arc}; + +fn bench(c: &mut Criterion) { + let rows = 8192; + let mut group = c.benchmark_group("decimal_hash"); + for (null_every, tag) in [(0, "no_null"), (100, "sparse_null"), (2, "dense_null")] { + for width in [2, 32] { + let values: ArrayRef = Arc::new( + (0..(rows + 2) * width) + .map(|i| { + if null_every != 0 && i % null_every == 0 { + None + } else { + // Mix short signed encodings and values occupying all 16 bytes. + let magnitude = if i % 3 == 0 { 10_i128.pow(37) } else { 1 }; + Some((i as i128 % 7 - 3) * magnitude) + } + }) + .collect::() + .with_precision_and_scale(38, 0) + .unwrap(), + ); + let field = Arc::new(Field::new("item", values.data_type().clone(), true)); + let nulls = (null_every != 0) + .then(|| NullBuffer::from_iter((0..rows + 2).map(|i| (i + 1) % null_every != 0))); + let cases: Vec<(&str, ArrayRef)> = vec![ + ("scalar", values.slice(width, rows)), + ( + "list", + (Arc::new(ListArray::new( + Arc::clone(&field), + OffsetBuffer::from_lengths(std::iter::repeat_n(width, rows + 2)), + Arc::clone(&values), + nulls.clone(), + )) as ArrayRef) + .slice(1, rows), + ), + ( + "large_list", + (Arc::new(LargeListArray::new( + Arc::clone(&field), + OffsetBuffer::from_lengths(std::iter::repeat_n(width, rows + 2)), + Arc::clone(&values), + nulls.clone(), + )) as ArrayRef) + .slice(1, rows), + ), + ( + "fixed_list", + (Arc::new(FixedSizeListArray::new(field, width as i32, values, nulls)) + as ArrayRef) + .slice(1, rows), + ), + ]; + for (shape, array) in cases { + // A scalar control needs only one width. + if shape == "scalar" && width != 2 { + continue; + } + for name in ["murmur3", "xxhash64"] { + let args = vec![ + ColumnarValue::Array(Arc::clone(&array)), + ColumnarValue::Scalar(if name == "murmur3" { + ScalarValue::Int32(Some(42)) + } else { + ScalarValue::Int64(Some(42)) + }), + ]; + group.bench_with_input( + BenchmarkId::new(name, format!("{shape}/{width}/{tag}")), + &args, + |b, args| { + if name == "murmur3" { + b.iter(|| black_box(spark_murmur3_hash(black_box(args)).unwrap())); + } else { + b.iter(|| black_box(spark_xxhash64(black_box(args)).unwrap())); + } + }, + ); + } + } + } + } + group.finish(); +} + +criterion_group!(benches, bench); +criterion_main!(benches); diff --git a/native/spark-expr/src/hash_funcs/utils.rs b/native/spark-expr/src/hash_funcs/utils.rs index 55beb6cd378..a44102c5bc4 100644 --- a/native/spark-expr/src/hash_funcs/utils.rs +++ b/native/spark-expr/src/hash_funcs/utils.rs @@ -234,33 +234,25 @@ macro_rules! hash_array_small_decimal { }; } +/// Hash one wide decimal without allocating BigInteger's signed byte representation. +#[macro_export] +macro_rules! hash_decimal { + ($value:expr, $seed:expr, $hash_method:ident) => {{ + let value: i128 = $value; + let bytes = value.to_be_bytes(); + // Match BigInteger.toByteArray(): drop redundant sign bytes, retaining + // one sign bit. Zero and -1 still occupy one byte. + let sign_bits = value.leading_zeros().max(value.leading_ones()); + let start = (sign_bits.saturating_sub(1) / 8) as usize; + $hash_method(&bytes[start..], $seed) + }}; +} + #[macro_export] macro_rules! hash_array_decimal { ($array_type:ident, $column: ident, $hashes: ident, $hash_method: ident) => { - let array = $column - .as_any() - .downcast_ref::<$array_type>() - .unwrap_or_else(|| { - panic!( - "Failed to downcast column to {}. Actual data type: {:?}.", - stringify!($array_type), - $column.data_type() - ) - }); - - if array.null_count() == 0 { - // Fast path: no nulls, use direct indexing - for i in 0..$hashes.len() { - $hashes[i] = $hash_method(array.value(i).to_le_bytes(), $hashes[i]); - } - } else { - // Slow path: check nulls - for i in 0..$hashes.len() { - if !array.is_null(i) { - $hashes[i] = $hash_method(array.value(i).to_le_bytes(), $hashes[i]); - } - } - } + let hash_decimal = |value: &i128, seed| $crate::hash_decimal!(*value, seed, $hash_method); + $crate::hash_array!($array_type, $column, $hashes, hash_decimal); }; } @@ -423,6 +415,12 @@ macro_rules! hash_list_with_primitive_elements { let elem_array = $values.as_any().downcast_ref::().unwrap(); $crate::hash_list_primitive!(offsets: $offsets, $list_array, elem_array, $hashes_buffer, $hash_method, |v: i64| v.to_le_bytes()); } + DataType::Decimal128(precision, _) if *precision > 18 => { + // Downcast once per list column instead of slicing and dispatching each element. + let elem_array = $values.as_any().downcast_ref::().unwrap(); + let hash_decimal = |value: i128, seed| $crate::hash_decimal!(value, seed, $hash_method); + $crate::hash_list_primitive!(offsets: $offsets, $list_array, elem_array, $hashes_buffer, hash_decimal, |v| v); + } _ => { // Fall back to recursive approach for complex element types $crate::hash_list_array!($list_array_type, $fallback_offset_type, $col, $hashes_buffer, $recursive_hash_method); @@ -518,6 +516,11 @@ macro_rules! hash_list_with_primitive_elements { let elem_array = $values.as_any().downcast_ref::().unwrap(); $crate::hash_list_primitive!(fixed_size: $list_size, $list_array, elem_array, $hashes_buffer, $hash_method, |v: i64| v.to_le_bytes()); } + DataType::Decimal128(precision, _) if *precision > 18 => { + let elem_array = $values.as_any().downcast_ref::().unwrap(); + let hash_decimal = |value: i128, seed| $crate::hash_decimal!(value, seed, $hash_method); + $crate::hash_list_primitive!(fixed_size: $list_size, $list_array, elem_array, $hashes_buffer, hash_decimal, |v| v); + } _ => { // Fall back to recursive approach for complex element types if $list_array.null_count() == 0 { @@ -1323,3 +1326,210 @@ pub(crate) mod test_utils { }; } } + +#[cfg(test)] +mod tests { + use super::super::{spark_murmur3_hash, spark_xxhash64}; + use arrow::array::{Array, ArrayRef, BinaryArray, Decimal128Array, Int32Array}; + use datafusion::common::{Result, ScalarValue}; + use datafusion::logical_expr::ColumnarValue; + use num::BigInt; + use std::sync::Arc; + + type HashFunction = fn(&[ColumnarValue]) -> Result; + + #[test] + fn wide_decimal_hashes_match_signed_big_endian_bytes() { + for precision in 19..=38 { + let max = 10_i128.pow(precision as u32) - 1; + let mut values = vec![Some(0), Some(max), Some(-max), None]; + // Both sides of every signed byte-length boundary, including 127/128 + // and -128/-129, as well as values that fit in i64 despite the type. + for bit in (7..127).step_by(8) { + for value in [(1_i128 << bit) - 1, 1_i128 << bit, (1_i128 << bit) + 1] { + if value <= max { + values.extend([Some(value), Some(-value)]); + } + } + } + let bytes: Vec<_> = values + .iter() + .map(|v| v.map(|v| BigInt::from(v).to_signed_bytes_be())) + .collect(); + let binary: ArrayRef = + Arc::new(BinaryArray::from_iter(bytes.iter().map(|v| v.as_deref()))); + for scale in [0, 2, precision as i8] { + let decimal: ArrayRef = Arc::new( + Decimal128Array::from(values.clone()) + .with_precision_and_scale(precision, scale) + .unwrap(), + ); + // Exercise both null-aware and null-free loops, including a slice. + for (offset, len) in [(0, values.len()), (4, values.len() - 4)] { + let decimal = decimal.slice(offset, len); + let binary = binary.slice(offset, len); + let prefix: ArrayRef = Arc::new(Int32Array::from_iter_values(0..len as i32)); + for seed in [0_i64, 42, -1, i32::MIN as i64] { + for chained in [false, true] { + let args = |value: ArrayRef, seed: ScalarValue| { + let mut args = vec![ColumnarValue::Array(value)]; + if chained { + args.insert(0, ColumnarValue::Array(Arc::clone(&prefix))); + args.push(ColumnarValue::Array(Arc::clone(&prefix))); + } + args.push(ColumnarValue::Scalar(seed)); + args + }; + for (hash, seed) in [ + ( + spark_murmur3_hash as HashFunction, + ScalarValue::Int32(Some(seed as i32)), + ), + ( + spark_xxhash64 as HashFunction, + ScalarValue::Int64(Some(seed)), + ), + ] { + let actual = hash(&args(Arc::clone(&decimal), seed.clone())) + .unwrap() + .into_array(len) + .unwrap(); + let expected = hash(&args(Arc::clone(&binary), seed)) + .unwrap() + .into_array(len) + .unwrap(); + assert_eq!( + actual.to_data(), + expected.to_data(), + "precision={precision}, scale={scale}, chained={chained}" + ); + } + } + } + } + } + } + } + + #[test] + fn decimal_lists_match_independent_binary_encoding() { + use arrow::array::{FixedSizeListArray, LargeListArray, ListArray}; + use arrow::buffer::{NullBuffer, OffsetBuffer}; + use arrow::datatypes::Field; + + for precision in [18, 19, 38] { + let max = 10_i128.pow(precision as u32) - 1; + let values = vec![ + 0, -1, 127, 128, -128, -129, max, -max, 255, 256, -256, -257, 32767, 32768, -32768, + -32769, + ]; + for scale in [0, 2] { + for parent_nulls in [false, true] { + for child_nulls in [false, true] { + let nulls = child_nulls + .then(|| NullBuffer::from_iter((0..values.len()).map(|i| i % 3 != 1))); + // Keep nonzero values under null slots, and values under null parents. + let decimal: ArrayRef = Arc::new( + Decimal128Array::new(values.clone().into(), nulls.clone()) + .with_precision_and_scale(precision, scale) + .unwrap(), + ); + let bytes: Vec<_> = values + .iter() + .map(|&v| { + if precision <= 18 { + (v as i64).to_le_bytes().to_vec() + } else { + BigInt::from(v).to_signed_bytes_be() + } + }) + .collect(); + let binary: ArrayRef = Arc::new(BinaryArray::from_iter( + bytes.iter().enumerate().map(|(i, v)| { + if nulls.as_ref().is_some_and(|n| n.is_null(i)) { + None + } else { + Some(v.as_slice()) + } + }), + )); + for kind in ["list", "large_list", "fixed_list"] { + let wrap = |values: ArrayRef| -> ArrayRef { + let field = + Arc::new(Field::new("item", values.data_type().clone(), true)); + let nulls = parent_nulls + .then(|| NullBuffer::from_iter((0..8).map(|i| i % 3 != 1))); + match kind { + "list" => Arc::new(ListArray::new( + field, + OffsetBuffer::new( + vec![0_i32, 2, 2, 5, 8, 11, 12, 14, 16].into(), + ), + values, + nulls, + )), + "large_list" => Arc::new(LargeListArray::new( + field, + OffsetBuffer::new( + vec![0_i64, 2, 2, 5, 8, 11, 12, 14, 16].into(), + ), + values, + nulls, + )), + _ => Arc::new(FixedSizeListArray::new(field, 2, values, nulls)), + } + }; + let decimal = wrap(Arc::clone(&decimal)); + let binary = wrap(Arc::clone(&binary)); + for (offset, len) in [(0, 8), (1, 6), (3, 1), (8, 0)] { + let prefix: ArrayRef = + Arc::new(Int32Array::from_iter_values(0..len as i32)); + for seed in [0_i64, 42, -1, i32::MIN as i64] { + for chained in [false, true] { + let args = |array: &ArrayRef, seed: ScalarValue| { + let mut args = vec![ColumnarValue::Array( + array.slice(offset, len), + )]; + if chained { + args.insert( + 0, + ColumnarValue::Array(Arc::clone(&prefix)), + ); + args.push(ColumnarValue::Array(Arc::clone( + &prefix, + ))); + } + args.push(ColumnarValue::Scalar(seed)); + args + }; + for (hash, seed) in [ + ( + spark_murmur3_hash as HashFunction, + ScalarValue::Int32(Some(seed as i32)), + ), + ( + spark_xxhash64 as HashFunction, + ScalarValue::Int64(Some(seed)), + ), + ] { + let actual = hash(&args(&decimal, seed.clone())) + .unwrap() + .into_array(len) + .unwrap(); + let expected = hash(&args(&binary, seed)) + .unwrap() + .into_array(len) + .unwrap(); + assert_eq!(actual.to_data(), expected.to_data(), + "{kind}, precision={precision}, scale={scale}, parent_nulls={parent_nulls}, child_nulls={child_nulls}, offset={offset}, len={len}, chained={chained}"); + } + } + } + } + } + } + } + } + } + } +} diff --git a/native/spark-expr/src/hash_funcs/xxhash64.rs b/native/spark-expr/src/hash_funcs/xxhash64.rs index 24697114620..69f5ff966cc 100644 --- a/native/spark-expr/src/hash_funcs/xxhash64.rs +++ b/native/spark-expr/src/hash_funcs/xxhash64.rs @@ -47,6 +47,8 @@ const SPARK_DEFAULT_SEED: i64 = 42; /// - `Struct` (and anything containing one): `SparkXxhash64` does not push a parent null /// mask into children, so hidden values of a NULL struct would affect the hash /// - a `Dictionary` nested in a list/map: `SparkXxhash64` restarts those hashes from 42 +/// - wide decimals: `SparkXxhash64` hashes fixed-width little-endian bytes instead +/// of Spark's minimal signed big-endian representation /// - `Time64`, which `SparkXxhash64` does not dispatch /// - `Float32`/`Float64` (and anything containing one): `SparkXxhash64` hashes the raw bits of a /// NaN, where Spark hashes every NaN as the canonical NaN @@ -132,7 +134,7 @@ fn type_compatible_with_spark_xxhash64(dt: &DataType, in_list_or_map: bool) -> b Boolean | Int8 | Int16 | Int32 | Int64 => true, Utf8 | LargeUtf8 | Binary | LargeBinary | FixedSizeBinary(_) => true, Date32 | Date64 | Timestamp(_, _) => true, - Decimal128(_, _) => true, + Decimal128(precision, _) => *precision <= 18, Dictionary(_, value) if !in_list_or_map => { type_compatible_with_spark_xxhash64(value.as_ref(), true) } diff --git a/native/spark-expr/src/hash_funcs/xxhash64_diff.rs b/native/spark-expr/src/hash_funcs/xxhash64_diff.rs index bf16263358d..c427de88609 100644 --- a/native/spark-expr/src/hash_funcs/xxhash64_diff.rs +++ b/native/spark-expr/src/hash_funcs/xxhash64_diff.rs @@ -29,11 +29,11 @@ //! | Utf8, LargeUtf8, Binary, LargeBinary, FixedSizeBinary | yes | //! | Date32, Date64, Timestamp | yes | //! | Decimal128 precision ≤ 18 | yes | -//! | Decimal128 precision > 18 | yes | +//! | Decimal128 precision > 18 | **no** — upstream uses fixed-width little-endian bytes | //! | Dictionary (top-level) | yes | //! | List / LargeList / FixedSizeList of primitives | yes | //! | Map<Utf8, Int32> / Map<Int32, Utf8> / Map<Utf8, Utf8> / Map<Int32, Int32> | yes | -//! | Map<Utf8, Decimal128> | yes | +//! | Map<Utf8, Decimal128> | only for precision ≤ 18 | //! | Struct (non-null parent) | yes, but not routed (see below) | //! | Struct NULL with hidden children | **no** — `SparkXxhash64` hashes hidden children | //! | List<Dictionary> | **no** — upstream restarts from seed 42 | @@ -125,6 +125,18 @@ fn assert_compatible(label: &str, arrays: &[ArrayRef]) { assert_eq!(expr, upstream, "{label}: expression mismatch"); } +// The pinned DataFusion implementation still uses the old decimal encoding. +// Spark parity is checked separately; routing must not reintroduce that encoding. +fn assert_wide_decimal_uses_comet(label: &str, arrays: &[ArrayRef]) { + let kernel = comet_kernel(arrays, SPARK_DEFAULT_SEED).unwrap(); + let upstream = spark_xxhash64_upstream(arrays).unwrap(); + assert_ne!( + kernel, upstream, + "{label}: expected different decimal encodings" + ); + assert_eq!(comet_expr(arrays, 42).unwrap(), kernel, "{label}: routing"); +} + fn col(array: impl Array + 'static) -> Vec { vec![Arc::new(array) as ArrayRef] } @@ -605,7 +617,7 @@ fn decimal128_precision_18_fits_i64() { #[test] fn decimal128_precision_20_does_not_fit_i64() { let too_big = 10_000_000_000_000_000_000i128; // 1e19, beyond i64::MAX - assert_compatible( + assert_wide_decimal_uses_comet( "Decimal128(20,2)", &col(decimal128( 20, @@ -625,7 +637,7 @@ fn decimal128_precision_20_does_not_fit_i64() { #[test] fn decimal128_precision_38() { let wide = 10i128.pow(28); - assert_compatible( + assert_wide_decimal_uses_comet( "Decimal128(38,10)", &col(decimal128( 38, @@ -642,6 +654,27 @@ fn decimal128_precision_38() { ); } +#[test] +fn wide_decimal_list_and_dictionary_use_comet() { + let values: ArrayRef = Arc::new(decimal128(38, 0, vec![Some(128), Some(-129), None])); + let dict: ArrayRef = Arc::new( + DictionaryArray::::try_new( + Int8Array::from(vec![Some(1), Some(0), Some(2), None]), + Arc::clone(&values), + ) + .unwrap(), + ); + let list: ArrayRef = Arc::new(ListArray::new( + Arc::new(Field::new("item", values.data_type().clone(), true)), + OffsetBuffer::new(vec![0, 2, 3].into()), + values, + None, + )); + for array in [dict, list] { + assert_wide_decimal_uses_comet("nested wide decimal", &[array]); + } +} + // --------------------------------------------------------------------------- // Dictionary // --------------------------------------------------------------------------- @@ -944,7 +977,7 @@ fn map_utf8_decimal128_small() { #[test] fn map_utf8_decimal128_large() { let wide = 10_000_000_000_000_000_000i128; - assert_compatible( + assert_wide_decimal_uses_comet( "Map", &[map_utf8_decimal( 20, diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 9bd32816c64..949ce8da5b8 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -281,9 +281,8 @@ object CometConf extends ShimCometConf { "comes from a Spark operator converts that input to Arrow format and uses Comet " + "native shuffle instead. This applies only where native shuffle supports the " + "partitioning and Comet can convert all of the input's columns, which rules out " + - "calendar intervals. A shuffle that hashes a decimal wider than 18 digits, or a " + - "string or a value computed from a string, stays on the JVM columnar shuffle, so " + - "that it partitions rows as Spark does.") + "calendar intervals. A shuffle that hashes a string or a value computed from a string " + + "stays on the JVM columnar shuffle, so that it partitions rows as Spark does.") .booleanConf .createWithDefault(false) diff --git a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala index 1222bdd9476..6df17048240 100644 --- a/spark/src/main/scala/org/apache/comet/serde/aggregates.scala +++ b/spark/src/main/scala/org/apache/comet/serde/aggregates.scala @@ -1310,8 +1310,7 @@ object CometApproxCountDistinct extends CometAggregateExpressionSerde[HyperLogLo // Types that Comet's native `xxhash64` hashes identically to Spark's `XxHash64Function`. // `StringType` here is the default UTF8_BINARY collation; a collated `StringType(collationId)` // falls through to the `case _` because Spark hashes those via the collation sort key. - // `DecimalType` is limited to precision <= 18, since Spark hashes wider decimals through - // `BigDecimal` (two's-complement big-endian bytes), which the native `i128` path does not match. + // Wide decimals remain unsupported here pending HLL register compatibility tests. private def hashableType(dt: DataType): Boolean = dt match { case BooleanType => true case ByteType | ShortType | IntegerType | LongType => true diff --git a/spark/src/main/scala/org/apache/comet/serde/hash.scala b/spark/src/main/scala/org/apache/comet/serde/hash.scala index 760bb3c2829..04cd184fe41 100644 --- a/spark/src/main/scala/org/apache/comet/serde/hash.scala +++ b/spark/src/main/scala/org/apache/comet/serde/hash.scala @@ -20,7 +20,7 @@ package org.apache.comet.serde import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, Murmur3Hash, Sha1, Sha2, XxHash64} -import org.apache.spark.sql.types.{ArrayType, DataType, DecimalType, IntegerType, LongType, MapType, StringType, StructType} +import org.apache.spark.sql.types.{ArrayType, DataType, IntegerType, LongType, MapType, StringType, StructType} import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, isTimeType, scalarFunctionExprToProtoWithReturnType, serializeDataType, supportedDataType} @@ -116,12 +116,10 @@ object CometSha1 extends CometExpressionSerde[Sha1] { private object HashUtils { - private val unsupportedDecimalReason = - "`DecimalType` with precision > 18 is not supported (Spark hashes via Java `BigDecimal`)" private val unsupportedTimeTypeReason = "`TimeType` is not supported" val unsupportedReasons: Seq[String] = - Seq(unsupportedDecimalReason, unsupportedTimeTypeReason, "Unsupported child data type") + Seq(unsupportedTimeTypeReason, "Unsupported child data type") def supportLevelForChildren(expr: Expression): SupportLevel = { expr.children.iterator @@ -134,8 +132,6 @@ private object HashUtils { } private def unsupportedReasonFor(dt: DataType): Option[String] = dt match { - // Keep in sync with CometShuffleExchangeExec's hash-key restriction until #5994 is fixed. - case d: DecimalType if d.precision > 18 => Some(unsupportedDecimalReason) case s: StructType => s.fields.iterator.flatMap(f => unsupportedReasonFor(f.dataType).iterator).toSeq.headOption case a: ArrayType => unsupportedReasonFor(a.elementType) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala index 4113fa920b3..4dcbca75cd9 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala @@ -624,11 +624,8 @@ object CometShuffleExchangeExec _: FloatType | _: DoubleType | _: StringType | _: BinaryType | _: TimestampType | _: TimestampNTZType | _: DateType => true - case d: DecimalType => - // Match the SQL hash restriction in serde/HashUtils until #5994 fixes native encoding. - // Different partition assignments break mixed native/Spark joins. A single partition - // does not hash the key: CometNativeShuffleWriter serializes it as SinglePartition. - d.precision <= 18 || s.outputPartitioning.numPartitions == 1 + case _: DecimalType => + true case dt if isTimeType(dt) => true case StructType(fields) if nestedHashPartitioningEnabled => diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q49/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q49/extended.txt index 1f9d12fc934..6deb75bea15 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q49/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v1_4/q49/extended.txt @@ -1,7 +1,7 @@ CometColumnarToRow +- CometTakeOrderedAndProject +- CometHashAggregate - +- CometColumnarExchange + +- CometExchange +- CometHashAggregate +- CometUnion :- CometProject diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark3_5/q70a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark3_5/q70a/extended.txt index 09918cf843c..185ce8a5247 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark3_5/q70a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark3_5/q70a/extended.txt @@ -5,7 +5,7 @@ CometColumnarToRow +- CometSort +- CometExchange +- CometHashAggregate - +- CometColumnarExchange + +- CometExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark4_1/q14a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark4_1/q14a/extended.txt index 2a3100957f1..9a11b4895a7 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark4_1/q14a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7-spark4_1/q14a/extended.txt @@ -1,7 +1,7 @@ CometColumnarToRow +- CometTakeOrderedAndProject +- CometHashAggregate - +- CometColumnarExchange + +- CometExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q14a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q14a/extended.txt index 8f5b9835de5..a4ef99efa8e 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q14a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q14a/extended.txt @@ -1,7 +1,7 @@ CometColumnarToRow +- CometTakeOrderedAndProject +- CometHashAggregate - +- CometColumnarExchange + +- CometExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q36a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q36a/extended.txt index 04405d9b84e..f11c15706a4 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q36a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q36a/extended.txt @@ -5,7 +5,7 @@ CometColumnarToRow +- CometSort +- CometExchange +- CometHashAggregate - +- CometColumnarExchange + +- CometExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q49/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q49/extended.txt index 1f9d12fc934..6deb75bea15 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q49/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q49/extended.txt @@ -1,7 +1,7 @@ CometColumnarToRow +- CometTakeOrderedAndProject +- CometHashAggregate - +- CometColumnarExchange + +- CometExchange +- CometHashAggregate +- CometUnion :- CometProject diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q5a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q5a/extended.txt index 4b41ce8a173..f42b9795c47 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q5a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q5a/extended.txt @@ -1,7 +1,7 @@ CometColumnarToRow +- CometTakeOrderedAndProject +- CometHashAggregate - +- CometColumnarExchange + +- CometExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q70a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q70a/extended.txt index 6b6725c38f5..594ae116008 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q70a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q70a/extended.txt @@ -5,7 +5,7 @@ CometColumnarToRow +- CometSort +- CometExchange +- CometHashAggregate - +- CometColumnarExchange + +- CometExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q77a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q77a/extended.txt index a3fbdd2c329..daccf86c3f7 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q77a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q77a/extended.txt @@ -1,7 +1,7 @@ CometColumnarToRow +- CometTakeOrderedAndProject +- CometHashAggregate - +- CometColumnarExchange + +- CometExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q80a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q80a/extended.txt index ea972a9b91d..f4db0be1188 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q80a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q80a/extended.txt @@ -1,7 +1,7 @@ CometColumnarToRow +- CometTakeOrderedAndProject +- CometHashAggregate - +- CometColumnarExchange + +- CometExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q86a/extended.txt b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q86a/extended.txt index b2fe83f4b36..6269f5ddc51 100644 --- a/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q86a/extended.txt +++ b/spark/src/test/resources/tpcds-plan-stability/approved-plans-v2_7/q86a/extended.txt @@ -5,7 +5,7 @@ CometColumnarToRow +- CometSort +- CometExchange +- CometHashAggregate - +- CometColumnarExchange + +- CometExchange +- CometHashAggregate +- CometUnion :- CometHashAggregate diff --git a/spark/src/test/scala/org/apache/comet/CometFuzzAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/CometFuzzAggregateSuite.scala index 32b21cb92ad..3674887f80f 100644 --- a/spark/src/test/scala/org/apache/comet/CometFuzzAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometFuzzAggregateSuite.scala @@ -19,10 +19,6 @@ package org.apache.comet -import org.apache.spark.sql.execution.aggregate.HashAggregateExec -import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec -import org.apache.spark.sql.types.DecimalType - import org.apache.comet.DataTypeSupport.isComplexType class CometFuzzAggregateSuite extends CometFuzzTestBase { @@ -35,16 +31,7 @@ class CometFuzzAggregateSuite extends CometFuzzTestBase { val (_, cometPlan) = checkSparkAnswer(sql) assert(1 == collectNativeScans(cometPlan).length) - val hasWideDecimalKey = df.schema(col).dataType match { - case d: DecimalType => d.precision > 18 - case _ => false - } - // Wide decimal hash keys require Spark shuffle when columnar shuffle is disabled. - if (CometConf.COMET_SHUFFLE_MODE.get() == "native" && hasWideDecimalKey) { - checkSparkAnswerAndOperator(sql, classOf[HashAggregateExec], classOf[ShuffleExchangeExec]) - } else { - checkSparkAnswerAndOperator(sql) - } + checkSparkAnswerAndOperator(sql) } } @@ -68,18 +55,7 @@ class CometFuzzAggregateSuite extends CometFuzzTestBase { val (_, cometPlan) = checkSparkAnswer(sql) assert(1 == collectNativeScans(cometPlan).length) - val hasWideDecimalKey = Seq("c1", "c2", "c3", col).exists { key => - df.schema(key).dataType match { - case d: DecimalType => d.precision > 18 - case _ => false - } - } - // Check both GROUP BY and DISTINCT keys, not unrelated decimal payload columns. - if (CometConf.COMET_SHUFFLE_MODE.get() == "native" && hasWideDecimalKey) { - checkSparkAnswerAndOperator(sql, classOf[HashAggregateExec], classOf[ShuffleExchangeExec]) - } else { - checkSparkAnswerAndOperator(sql) - } + checkSparkAnswerAndOperator(sql) } } diff --git a/spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala b/spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala index 91bd098c1f3..f8d8f307934 100644 --- a/spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala @@ -160,17 +160,6 @@ class CometFuzzTestSuite extends CometFuzzTestBase { } test("distribute by single column (complex types)") { - // Inspect only the key schema: any wide decimal leaf requires Spark's hash partition - // assignments, even when native hashing of nested keys is enabled. - def hasWideDecimal(dataType: DataType): Boolean = dataType match { - case decimal: DecimalType => decimal.precision > 18 - case StructType(fields) => fields.exists(field => hasWideDecimal(field.dataType)) - case ArrayType(elementType, _) => hasWideDecimal(elementType) - case MapType(keyType, valueType, _) => - hasWideDecimal(keyType) || hasWideDecimal(valueType) - case _ => false - } - val df = spark.read.parquet(filename) df.createOrReplaceTempView("t1") val columns = df.schema.fields.filter(f => isComplexType(f.dataType)).map(_.name) @@ -191,19 +180,14 @@ class CometFuzzTestSuite extends CometFuzzTestBase { } assert(cometShuffleExchanges.length == expectedNumCometShuffles) - // Enabling nested keys admits supported types, but wide decimal leaves still require - // Spark's hash partition assignments. JVM shuffle supports both kinds of key. + // Enabling nested keys admits native shuffle, including wide decimal leaves. withSQLConf(CometConf.COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_NESTED_ENABLED.key -> "true") { val enabledDf = spark.sql(sql) enabledDf.collect() val enabledPlan = enabledDf.queryExecution.executedPlan.asInstanceOf[AdaptiveSparkPlanExec].executedPlan - val expectedEnabledShuffles = - if (CometConf.COMET_SHUFFLE_MODE.get() == "native" && - hasWideDecimal(df.schema(col).dataType)) 0 - else 1 assert( - collectCometShuffleExchanges(enabledPlan).length == expectedEnabledShuffles, + collectCometShuffleExchanges(enabledPlan).length == 1, s"Unexpected shuffle for ${df.schema(col)}") } } diff --git a/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala index dc9ae920ef6..7e823fb347c 100644 --- a/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometHashExpressionSuite.scala @@ -25,6 +25,7 @@ import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, ParquetGenerator, SchemaGenOptions} +import org.apache.comet.udf.codegen.CometScalaUDFCodegen /** * Test suite for Spark `hash` (murmur3) and `xxhash64` compatibility between Spark and Comet. @@ -34,54 +35,80 @@ import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, ParquetGener */ class CometHashExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper { + private def checkNativeHashes(query: String): Unit = { + withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") { + CometScalaUDFCodegen.resetStats() + checkSparkAnswerAndImpl(sql(query), native = Seq("hash", "xxhash64")) + val stats = CometScalaUDFCodegen.stats() + assert( + stats.compileCount + stats.cacheHitCount == 0, + s"unexpected dispatcher activity: $stats") + } + } + test("hash - decimal (precision <= 18)") { Seq((10, 2), (18, 0), (18, 10)).foreach { case (precision, scale) => withTable("t") { sql(s"CREATE TABLE t(c DECIMAL($precision, $scale)) USING parquet") sql("INSERT INTO t VALUES (1.23), (-1.23), (0.0), (null)") - checkSparkAnswerAndOperator("SELECT c, hash(c), xxhash64(c) FROM t ORDER BY c") + checkNativeHashes("SELECT c, hash(c), xxhash64(c) FROM t ORDER BY c") } } } test("hash - decimal (precision > 18)") { - Seq((20, 2), (38, 10)).foreach { case (precision, scale) => + Seq((19, 0), (20, 2), (38, 0), (38, 10), (38, 38)).foreach { case (precision, scale) => withTable("t") { sql(s"CREATE TABLE t(c DECIMAL($precision, $scale)) USING parquet") - sql("INSERT INTO t VALUES (1.23), (-1.23), (0.0), (null)") - // Large decimals may fall back to Spark, so just check the answer - checkSparkAnswer("SELECT c, hash(c), xxhash64(c) FROM t ORDER BY c") + val max = BigInt(10).pow(precision) - 1 + val unscaled = Seq(BigInt(0), max, -max) ++ + Seq(1, 127, 128, 129, 255, 256, 32767, 32768).flatMap(v => + Seq(BigInt(v), -BigInt(v))) ++ + Seq( + BigInt(Long.MaxValue), + BigInt(Long.MaxValue) + 1, + BigInt(Long.MinValue), + BigInt(Long.MinValue) - 1) + val values = unscaled.map { v => + val decimal = new java.math.BigDecimal(v.bigInteger, scale).toPlainString + s"(CAST('$decimal' AS DECIMAL($precision, $scale)))" + } + sql(s"INSERT INTO t VALUES ${(values :+ "(null)").mkString(",")}") + checkNativeHashes("SELECT c, hash(c), xxhash64(c) FROM t") + checkNativeHashes("SELECT hash(1, c, c), xxhash64(1, c, c) FROM t") } } } - test("hash - array of decimal (precision > 18) falls back to Spark") { + test("hash - array of decimal (precision > 18)") { withTable("t") { - sql("CREATE TABLE t(c ARRAY) USING parquet") - sql("INSERT INTO t VALUES (array(1.23, 2.34)), (null)") - // Should fall back to Spark due to nested high-precision decimal - checkSparkAnswerAndFallbackReason("SELECT c, hash(c), xxhash64(c) FROM t", "precision > 18") + sql("CREATE TABLE t(c ARRAY) USING parquet") + sql("""INSERT INTO t VALUES (array(127, 128, -128, -129)), + (array(99999999999999999999999999999999999999BD, null, -1)), + (array()), (array(null)), (null)""") + checkNativeHashes("SELECT hash(c), xxhash64(c), hash(1, c, c), xxhash64(1, c, c) FROM t") } } - test("hash - struct with decimal (precision > 18) falls back to Spark") { + test("hash - struct with decimal (precision > 18)") { withTable("t") { - sql("CREATE TABLE t(c STRUCT) USING parquet") - sql("INSERT INTO t VALUES (named_struct('a', 1, 'b', 1.23)), (null)") - // Should fall back to Spark due to nested high-precision decimal - checkSparkAnswerAndFallbackReason("SELECT c, hash(c), xxhash64(c) FROM t", "precision > 18") + sql("CREATE TABLE t(c STRUCT) USING parquet") + sql("""INSERT INTO t VALUES (named_struct('a', 1, 'b', 1.28)), + (named_struct('a', 2, 'b', -1.29)), (named_struct('a', null, 'b', null)), + (named_struct('a', 3, 'b', 999999999999999999999999999999999999.99BD)), (null)""") + checkNativeHashes("SELECT hash(c), xxhash64(c), hash(1, c, c), xxhash64(1, c, c) FROM t") + checkNativeHashes("SELECT hash(array(c, c)), xxhash64(array(c, c)) FROM t") } } - test("hash - map with decimal (precision > 18) value falls back to Spark") { + test("hash - map with decimal (precision > 18) keys and values") { withSQLConf("spark.sql.legacy.allowHashOnMapType" -> "true") { withTable("t") { - sql("CREATE TABLE t(c MAP) USING parquet") - sql("INSERT INTO t VALUES (map('a', 1.23)), (null)") - // Should fall back to Spark due to nested high-precision decimal - checkSparkAnswerAndFallbackReason( - "SELECT c, hash(c), xxhash64(c) FROM t", - "precision > 18") + sql("CREATE TABLE t(c MAP) USING parquet") + sql("""INSERT INTO t VALUES (map(128, -129, -128, 127)), + (map(99999999999999999999999999999999999999BD, null)), + (map()), (null)""") + checkNativeHashes("SELECT hash(c), xxhash64(c), hash(1, c, c), xxhash64(1, c, c) FROM t") } } } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala index 8d91ec71765..eb1086a9a2d 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeShuffleSuite.scala @@ -50,7 +50,7 @@ import org.apache.spark.sql.execution.aggregate.ObjectHashAggregateExec import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec} import org.apache.spark.sql.functions.{broadcast, col, count, countDistinct, spark_partition_id, sum} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.types.{ArrayType, DataType, DecimalType, IntegerType, LongType, MapType, StringType, StructField, StructType} +import org.apache.spark.sql.types.{ArrayType, DataType, IntegerType, LongType, MapType, StringType, StructField, StructType} import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} import org.apache.comet.{CometConf, CometExecIterator, CometExplainInfo, CometShuffleBlockIterator, CometShuffleSizeLimitException, Native} @@ -441,14 +441,7 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper val shuffled = df .select($"_1") .repartition(10, col(c)) - val nativeHashSupported = df.schema(c).dataType match { - case d: DecimalType => d.precision <= 18 - case _ => true - } - checkShuffleAnswer( - shuffled, - if (nativeHashSupported) 1 else 0, - checkNativeOperators = nativeHashSupported) + checkShuffleAnswer(shuffled, 1, checkNativeOperators = true) } } } @@ -468,12 +461,10 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper for (mode <- Seq("native", "auto", "jvm")) { withSQLConf(CometConf.COMET_SHUFFLE_MODE.key -> mode) { val shuffled = spark.table("decimal_shuffle").repartition(7, $"k") - val native = precision <= 18 && mode != "jvm" - val sparkShuffle = precision > 18 && mode == "native" - checkCometExchange(shuffled, if (sparkShuffle) 0 else 1, native) + checkCometExchange(shuffled, 1, native = mode != "jvm") assert(shuffled.queryExecution.executedPlan.collect { case _: ShuffleExchangeExec => 1 - }.sum == (if (sparkShuffle) 1 else 0)) + }.sum == 0) // Result equality alone cannot detect a different hash partition assignment. checkSparkAnswer(shuffled.withColumn("partition", spark_partition_id())) } @@ -508,8 +499,7 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper } } - test( - "wide decimal hash shuffle keeps native aggregates unless multiple partitions need fallback") { + test("wide decimal hash shuffle keeps native aggregates") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", SQLConf.USE_OBJECT_HASH_AGG.key -> "true", @@ -528,18 +518,12 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper val df = sql(s"SELECT $key, sort_array($function(_2)) FROM decimal_shuffle GROUP BY $key") val plan = df.queryExecution.executedPlan - val nativeExpected = precision <= 18 || partitions == 1 - assert( - plan.collect { case _: CometHashAggregateExec => 1 }.sum == - (if (nativeExpected) 2 else 0), - plan.treeString) + assert(plan.collect { case _: CometHashAggregateExec => 1 }.sum == 2, plan.treeString) assert( - plan.collect { case _: ObjectHashAggregateExec => 1 }.sum == - (if (nativeExpected) 0 else 2), + plan.collect { case _: ObjectHashAggregateExec => 1 }.sum == 0, plan.treeString) - // Restoring Spark's aggregate buffers must retain the accelerated input scan. assert(plan.collect { case _: CometNativeScanExec => 1 }.sum == 1, plan.treeString) - val exchanges = checkCometExchange(df, if (nativeExpected) 1 else 0, native = true) + val exchanges = checkCometExchange(df, 1, native = true) // GROUP BY retains HashPartitioning even when there is only one partition. assert(exchanges.forall(_.outputPartitioning.isInstanceOf[HashPartitioning])) checkSparkAnswer(df) @@ -575,7 +559,8 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper val plan = df.queryExecution.executedPlan val exchanges = collect(plan) { case e: CometShuffleExchangeExec => e } assert(exchanges.size == 2, plan.treeString) - assert(exchanges.forall(_.shuffleType == CometColumnarShuffle), plan.treeString) + assert(exchanges.count(_.shuffleType == CometNativeShuffle) == 1, plan.treeString) + assert(exchanges.count(_.shuffleType == CometColumnarShuffle) == 1, plan.treeString) assert(collect(plan) { case _: CometNativeScanExec => 1 }.sum == 1, plan.treeString) assert(collect(plan) { case _: FileSourceScanExec => 1 }.sum == 1, plan.treeString) } @@ -614,11 +599,7 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper case e: CometShuffleExchangeExec => Some(e.shuffleType) case _: ShuffleExchangeExec => None } - val expected = - if (precision <= 18) Some(CometNativeShuffle) - else if (mode == "auto") Some(CometColumnarShuffle) - else None - assert(shuffles == Seq(expected), plan.treeString) + assert(shuffles == Seq(Some(CometNativeShuffle)), plan.treeString) } } } @@ -642,7 +623,7 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper |""".stripMargin) for (key <- Seq("s", "a")) { val shuffled = spark.table("decimal_shuffle").repartition(7, col(key)) - checkCometExchange(shuffled, if (precision <= 18) 1 else 0, native = true) + checkCometExchange(shuffled, 1, native = true) checkSparkAnswer(shuffled.withColumn("partition", spark_partition_id())) } } @@ -1002,6 +983,66 @@ class CometNativeShuffleSuite extends CometTestBase with AdaptiveSparkPlanHelper } } + test("native shuffle on wide decimal hash partitioning keys") { + withNestedHashPartitioning { + withTable("wide_decimals") { + sql("CREATE TABLE wide_decimals(id INT, c DECIMAL(38, 0)) USING parquet") + sql("""INSERT INTO wide_decimals VALUES (0, null), (1, 0), (2, 128), (3, -129), + (4, 99999999999999999999999999999999999999BD), + (5, -99999999999999999999999999999999999999BD), (6, 128)""") + Seq(Seq("c"), Seq("a"), Seq("s"), Seq("id", "c")).foreach { keys => + def shuffled = sql("""SELECT id, c, array(c, c) AS a, + named_struct('d', c) AS s FROM wide_decimals""") + .repartition(10, keys.map(col): _*) + checkCometExchange(shuffled, 1, native = true) + // This pins the common hash helper's shuffle caller, including chained and nested + // inputs. The old fixed-width encoding produced different partition assignments. + checkSparkAnswer(shuffled.selectExpr("id", "spark_partition_id()")) + checkSparkAnswer(shuffled.groupBy("c").count()) + } + } + } + } + + test("wide decimal joins keep native and Spark shuffle inputs copartitioned") { + withSQLConf( + CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "true", + CometConf.COMET_CONVERT_FROM_JSON_ENABLED.key -> "false", + SQLConf.USE_V1_SOURCE_LIST.key -> "parquet,json", + SQLConf.SHUFFLE_PARTITIONS.key -> "7", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false") { + withTable("decimal_native", "decimal_spark") { + Seq("decimal_native" -> "parquet", "decimal_spark" -> "json").foreach { + case (table, format) => + sql(s"CREATE TABLE $table(id INT, k DECIMAL(38, 0)) USING $format") + sql(s"""INSERT INTO $table VALUES (1, 1), (2, -1), (3, 128), (4, -129), + (5, 99999999999999999999999999999999999999BD), + (6, -99999999999999999999999999999999999999BD)""") + } + for (mode <- Seq("native", "auto"); adaptive <- Seq(false, true)) { + withSQLConf( + CometConf.COMET_SHUFFLE_MODE.key -> mode, + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) { + val df = sql("""SELECT n.id, s.id FROM decimal_native n + JOIN decimal_spark s ON n.k = s.k""") + val (_, plan) = checkSparkAnswer(df) + val exchanges = collect(plan) { case e: CometShuffleExchangeExec => e } + assert(exchanges.count(_.shuffleType == CometNativeShuffle) == 1, plan.treeString) + if (mode == "auto") { + assert(exchanges.count(_.shuffleType == CometColumnarShuffle) == 1, plan.treeString) + } else { + assert(collect(plan) { case _: ShuffleExchangeExec => 1 }.sum == 1, plan.treeString) + } + assert(collect(plan) { case _: CometNativeScanExec => 1 }.sum == 1, plan.treeString) + assert(collect(plan) { case _: FileSourceScanExec => 1 }.sum == 1, plan.treeString) + } + } + } + } + } + test("native shuffle on struct hash partitioning key") { withNestedHashPartitioning { Seq(10, 201).foreach { numPartitions => diff --git a/spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala index 03181136005..fe9b595302f 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala @@ -163,6 +163,26 @@ class CometShuffleInputConversionSuite extends CometTestBase { Double.NegativeInfinity), DecimalType(18, 2) -> Seq("0.00", "-0.01", "9999999999999999.99", "-9999999999999999.99") .map(new java.math.BigDecimal(_)), + DecimalType(19, 0) -> Seq( + "0", + "-1", + "9223372036854775806", + "9223372036854775807", + "9223372036854775808", + "9999999999999999999", + "-9999999999999999999").map(new java.math.BigDecimal(_)), + DecimalType(38, 0) -> Seq( + "0", + "-1", + "9223372036854775808", + "-9223372036854775809", + "99999999999999999999999999999999999999", + "-99999999999999999999999999999999999999").map(new java.math.BigDecimal(_)), + DecimalType(38, 18) -> Seq( + "0.000000000000000000", + "-0.000000000000000001", + "99999999999999999999.999999999999999999", + "-99999999999999999999.999999999999999999").map(new java.math.BigDecimal(_)), DateType -> Seq("1970-01-01", "1969-12-31", "0001-01-01", "9999-12-31").map(Date.valueOf), TimestampType -> Seq( new Timestamp(0L), @@ -303,17 +323,23 @@ class CometShuffleInputConversionSuite extends CometTestBase { } } - convertTest("a shuffle that hashes a wide decimal stays on the JVM columnar shuffle") { - // Native shuffle hashes a decimal wider than 18 digits differently from Spark's partitioner - // (#5994). The right input, with its array column, stays on the JVM columnar shuffle, - // so the left one has to as well, or matching keys land in different partitions. + convertTest("a converted wide-decimal shuffle joins the JVM columnar shuffle") { withShuffledJoins { - val left = rowsDf(100).select(col("k").cast(DecimalType(38, 0)).as("lk"), col("l")) - val right = arraysDf().select(col("k").cast(DecimalType(38, 0)).as("rk"), col("xs")) - val df = left.join(right, col("lk") === col("rk")) - val (_, plan) = checkSparkAnswer(df) - assert(conversions(plan).isEmpty, plan) - checkCometExchange(df, 2, native = false) + Seq("true", "false").foreach { aqe => + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe, + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false", + SQLConf.SKEW_JOIN_ENABLED.key -> "false") { + val left = rowsDf(100).select(col("k").cast(DecimalType(38, 0)).as("lk"), col("l")) + // The array column keeps the right input on the JVM columnar shuffle. + val right = arraysDf().select(col("k").cast(DecimalType(38, 0)).as("rk"), col("xs")) + val (_, plan) = checkSparkAnswer(left.join(right, col("lk") === col("rk"))) + assert(convertedShuffles(plan).length == 1, s"AQE $aqe:\n$plan") + assert( + cometShuffles(plan).count(_.shuffleType == CometColumnarShuffle) == 1, + s"AQE $aqe:\n$plan") + } + } } } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometTypedDatasetSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometTypedDatasetSuite.scala index c9e1a349e3d..f2222b37f5d 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometTypedDatasetSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometTypedDatasetSuite.scala @@ -24,6 +24,7 @@ import java.util.concurrent.atomic.AtomicLong import org.apache.spark.sql.{CometTestBase, DataFrame, Dataset} import org.apache.spark.sql.catalyst.expressions.aggregate.Partial import org.apache.spark.sql.comet.{CometBroadcastHashJoinExec, CometHashAggregateExec, CometSparkToColumnarExec} +import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.execution.{ColumnarToRowExec, ColumnarToRowTransition, SerializeFromObjectExec, SparkPlan, WholeStageCodegenExec} import org.apache.spark.sql.functions.{broadcast, col, size, sum} import org.apache.spark.sql.internal.SQLConf @@ -210,27 +211,33 @@ class CometTypedDatasetSuite extends CometTestBase { } convertTest("a join on wide decimal keys with an input that is not converted") { - // The left input converts and the right one, with its array column, does not. Native - // shuffle hashes a decimal wider than 18 digits differently from Spark's partitioner (#5994), - // so the shuffle above the conversion has to stay on Comet's columnar shuffle like the right - // input's, or matching keys land in different partitions and the join loses rows. + // The left input converts its output to Arrow for native shuffle; the right input's + // array keeps JVM shuffle. + // Both shuffles must place matching wide decimal keys in the same partitions. withSQLConf( - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SHUFFLE_MODE.key -> "auto", SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false", SQLConf.SHUFFLE_PARTITIONS.key -> "10") { - val left = spark - .range(0, 100, 1, 2) - .map(i => TypedDsDecimal(new java.math.BigDecimal(i.longValue), i.longValue)) - .alias("l") - val right = spark - .range(0, 100, 1, 2) - .map(i => TypedDsDecimalInts(new java.math.BigDecimal(i.longValue), Seq(i.intValue))) - .alias("r") - val df = left.join(right, col("l.k") === col("r.k")).select(col("l.v"), col("r.xs")) - val (_, plan) = checkSparkAnswer(df) - assert(conversions(plan).nonEmpty, plan) - // One columnar shuffle for each input of the join. - checkCometExchange(df, 2, native = false) + for (adaptive <- Seq(false, true)) { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString) { + val left = spark + .range(0, 100, 1, 2) + .map(i => TypedDsDecimal(new java.math.BigDecimal(i.longValue), i.longValue)) + .alias("l") + val right = spark + .range(0, 100, 1, 2) + .map(i => TypedDsDecimalInts(new java.math.BigDecimal(i.longValue), Seq(i.intValue))) + .alias("r") + val df = left.join(right, col("l.k") === col("r.k")).select(col("l.v"), col("r.xs")) + val (_, plan) = checkSparkAnswer(df) + assert(conversions(plan).nonEmpty, plan) + val exchanges = collect(plan) { case e: CometShuffleExchangeExec => e } + assert(exchanges.count(_.shuffleType == CometNativeShuffle) == 1, plan.treeString) + assert(exchanges.count(_.shuffleType == CometColumnarShuffle) == 1, plan.treeString) + } + } } } diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashExpressionBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashExpressionBenchmark.scala index 238d8b662b1..6450b822fe2 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashExpressionBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometHashExpressionBenchmark.scala @@ -19,6 +19,11 @@ package org.apache.spark.sql.benchmark +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.Row + +import org.apache.comet.{CometConf, ExtendedExplainInfo} + case class HashExprConfig( name: String, query: String, @@ -46,6 +51,10 @@ object CometHashExpressionBenchmark extends CometBenchmarkBase { override def runCometBenchmark(mainArgs: Array[String]): Unit = { val values = 1024 * 1024 + if (mainArgs.contains("--wide-decimal")) { + runWideDecimalBenchmarks(values) + return + } runBenchmarkWithTable("Hash expression benchmarks", values) { v => withTempPath { dir => @@ -74,6 +83,60 @@ object CometHashExpressionBenchmark extends CometBenchmarkBase { runMurmur3HashBenchmarks(values) } + // Run only these cases with --wide-decimal. The fallback arm is the path used by + // wide-decimal hash on main before native support; codegen dispatch is tracked in #5835. + private def runWideDecimalBenchmarks(values: Int): Unit = { + runBenchmarkWithTable("Wide decimal hash", values) { rows => + withTempPath { dir => + withTempTable("parquetV1Table") { + prepareTable( + dir, + spark.sql(s""" + SELECT d, array(d, d) AS a FROM ( + SELECT CASE WHEN value % 100 = 0 THEN NULL + WHEN value % 2 = 0 THEN CAST(value AS DECIMAL(38, 0)) + ELSE CAST(value AS DECIMAL(38, 0)) * 1000000000000000000BD + END AS d FROM $tbl) + """)) + for (function <- Seq("hash", "xxhash64"); column <- Seq("d", "a")) { + val exprName = if (function == "hash") "Murmur3Hash" else "XxHash64" + // A checksum consumes every hash result and avoids collecting a million rows. + val query = + s"SELECT sum(CAST($function($column) AS DECIMAL(38, 0))) FROM parquetV1Table" + val cases = Seq( + ("Spark", false, false), + ("Comet with Spark hash fallback", true, false), + ("Comet native", true, true)) + var expected: Seq[Row] = Seq.empty + val benchmark = + new Benchmark(s"$function($column): decimal(38,0)", rows.toLong, output = output) + cases.foreach { case (name, comet, native) => + val conf = Seq( + CometConf.COMET_ENABLED.key -> comet.toString, + CometConf.COMET_EXEC_ENABLED.key -> comet.toString, + CometConf.getExprEnabledConfigKey(exprName) -> native.toString) + withSQLConf(conf: _*) { + val df = spark.sql(query) + val actual = df.collect().toSeq + if (!comet) expected = actual else assert(actual == expected, name) + if (comet) { + val plan = df.queryExecution.executedPlan + val info = new ExtendedExplainInfo() + assert(info.getNativeExpressions(plan).contains(function) == native, name) + assert(!info.getCodegenDispatchExpressions(plan).contains(function), name) + } + } + benchmark.addCase(name) { _ => + withSQLConf(conf: _*) { spark.sql(query).collect() } + } + } + benchmark.run() + } + } + } + } + } + /** * Comprehensive benchmarks for murmur3 hash function across different data types. These * benchmarks cover primitive types, complex types (arrays, structs), and nested structures to diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala index 5be64f696dd..af6178668ea 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala @@ -676,13 +676,15 @@ class CometCelebornShufflePlanningSuite extends CometTestBase { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, CometConf.COMET_SHUFFLE_MODE.key -> "native") { - val nativeRegistrations = manager.nativeRegistrations.get() - val plan = input.repartition(2, col("value")).queryExecution.executedPlan - assert(plan.isInstanceOf[AdaptiveSparkPlanExec] == adaptive) - val exchanges = cometExchanges(plan) - assert(exchanges.nonEmpty, s"expected native exchange:\n$plan") - assert(exchanges.forall(_.shuffleType == CometNativeShuffle)) - assert(manager.nativeRegistrations.get() == nativeRegistrations) + for (key <- Seq(col("value"), col("value").cast("decimal(38, 0)"))) { + val nativeRegistrations = manager.nativeRegistrations.get() + val plan = input.repartition(2, key).queryExecution.executedPlan + assert(plan.isInstanceOf[AdaptiveSparkPlanExec] == adaptive) + val exchanges = cometExchanges(plan) + assert(exchanges.nonEmpty, s"expected native exchange:\n$plan") + assert(exchanges.forall(_.shuffleType == CometNativeShuffle)) + assert(manager.nativeRegistrations.get() == nativeRegistrations) + } } } @@ -703,27 +705,19 @@ class CometCelebornShufflePlanningSuite extends CometTestBase { } } - for (wideDecimal <- Seq(false, true)) { - test( - s"unsupported repartition executes Spark fallback: wideDecimal=$wideDecimal, " + - s"AQE=$adaptive") { - withSQLConf( - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, - CometConf.COMET_SHUFFLE_MODE.key -> "native", - CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_ENABLED.key -> "false") { - val nativeRegistrations = manager.nativeRegistrations.get() - val sparkRegistrations = manager.sparkRegistrations.get() - val query = if (wideDecimal) { - input.repartition(2, col("value").cast("decimal(38, 0)")) - } else { - input.repartition(2) - } - assertSparkExchange(query.queryExecution.executedPlan) - checkAnswer(query, (1L to 32L).map(Row(_))) - assert(cometExchanges(query.queryExecution.executedPlan).isEmpty) - assert(manager.nativeRegistrations.get() == nativeRegistrations) - assert(manager.sparkRegistrations.get() > sparkRegistrations) - } + test(s"unsupported repartition executes Spark fallback: AQE=$adaptive") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> adaptive.toString, + CometConf.COMET_SHUFFLE_MODE.key -> "native", + CometConf.COMET_SHUFFLE_NATIVE_ROUND_ROBIN_PARTITIONING_ENABLED.key -> "false") { + val nativeRegistrations = manager.nativeRegistrations.get() + val sparkRegistrations = manager.sparkRegistrations.get() + val query = input.repartition(2) + assertSparkExchange(query.queryExecution.executedPlan) + checkAnswer(query, (1L to 32L).map(Row(_))) + assert(cometExchanges(query.queryExecution.executedPlan).isEmpty) + assert(manager.nativeRegistrations.get() == nativeRegistrations) + assert(manager.sparkRegistrations.get() > sparkRegistrations) } }