diff --git a/native/Cargo.lock b/native/Cargo.lock index 135a774a319..b40bae55722 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -2087,6 +2087,7 @@ dependencies = [ "tokio", "twox-hash", "uuid", + "xxhash-rust", ] [[package]] @@ -7468,6 +7469,12 @@ version = "0.13.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "66fee0b777b0f5ac1c69bb06d361268faafa61cd4682ae064a171c16c433e9e4" +[[package]] +name = "xxhash-rust" +version = "0.8.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aee1b19627c7c60102ab80d3a9cbe18de90bfe03bfa6c3715447681f0e8c8af6" + [[package]] name = "yoke" version = "0.8.3" diff --git a/native/shuffle/src/partitioners/multi_partition.rs b/native/shuffle/src/partitioners/multi_partition.rs index ccfe14790c2..94657395a80 100644 --- a/native/shuffle/src/partitioners/multi_partition.rs +++ b/native/shuffle/src/partitioners/multi_partition.rs @@ -27,6 +27,7 @@ use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion::execution::runtime_env::RuntimeEnv; use datafusion_comet_common::tracing::{with_trace, with_trace_async}; use datafusion_comet_spark_expr::murmur3::create_murmur3_hashes; +use datafusion_comet_spark_expr::xxh3_64::create_xxh3_64_hashes; use itertools::Itertools; use std::fmt; use std::fmt::{Debug, Formatter}; @@ -36,8 +37,13 @@ use tokio::time::Instant; /// Reusable scratch buffers for computing row-to-partition assignments. #[derive(Default)] struct ScratchSpace { - /// Hashes for each row in the current batch. + /// Murmur3 hashes for each row in the current batch. Allocated only for + /// the `Hash` partitioning arm (Spark-compatible seed 42, 32-bit hash). hashes_buf: Vec, + /// XXH3-64 hashes for each row in the current batch. Allocated only for + /// the `RoundRobin` partitioning arm, which is Comet-internal and does + /// not need Spark hash parity. + xxh3_hashes_buf: Vec, /// Partition ids for each row in the current batch. partition_ids: Vec, /// The row indices of the rows in each partition. This array is conceptually divided into @@ -225,11 +231,16 @@ impl MultiPartitionShuffleRepartitioner { // The initial values are not used. let scratch = ScratchSpace { hashes_buf: match partitioning { - // Allocate hashes_buf for hash and round robin partitioning. - // Round robin hashes all columns to achieve even, deterministic distribution. - CometPartitioning::Hash(_, _) | CometPartitioning::RoundRobin(_, _) => { - vec![0; batch_size] - } + // Spark-compatible Murmur3, seed 42, matches HashPartitioning + // downstream co-partitioning contracts. + CometPartitioning::Hash(_, _) => vec![0; batch_size], + _ => vec![], + }, + xxh3_hashes_buf: match partitioning { + // RoundRobin is Comet-internal (see comment in the RoundRobin + // arm of `partitioning_batch`); XXH3-64 gives faster + // per-row hashing on strings and binary than Murmur3. + CometPartitioning::RoundRobin(_, _) => vec![0; batch_size], _ => vec![], }, partition_ids: vec![0; batch_size], @@ -384,13 +395,14 @@ impl MultiPartitionShuffleRepartitioner { self.scratch = scratch; } CometPartitioning::RoundRobin(num_output_partitions, max_hash_columns) => { - // Comet implements "round robin" as hash partitioning on columns. - // This achieves the same goal as Spark's round robin (even distribution - // without semantic grouping) while being deterministic for fault tolerance. - // - // Note: This produces different partition assignments than Spark's round robin, - // which sorts by UnsafeRow binary representation before assigning partitions. - // However, both approaches provide even distribution and determinism. + // Comet implements "round robin" as hash partitioning on columns + // using XXH3-64. This achieves even distribution and determinism + // for fault tolerance. The hash function is a Comet-internal + // choice: RoundRobinPartitioning has no ShuffleSpec, so no + // downstream Spark stage co-partitions against this layout, and + // Spark's own round-robin sorts by UnsafeRow binary + // representation instead, so the partition assignments already + // differ from Spark either way. let mut scratch = std::mem::take(&mut self.scratch); let (partition_starts, partition_row_indices): (&Vec, &Vec) = { let mut timer = self.metrics.repart_time.timer(); @@ -409,18 +421,19 @@ impl MultiPartitionShuffleRepartitioner { .map(|i| Arc::clone(input.column(i))) .collect(); - // Use identical seed as Spark hash partitioning. - let hashes_buf = &mut scratch.hashes_buf[..num_rows]; - hashes_buf.fill(42_u32); + let hashes_buf = &mut scratch.xxh3_hashes_buf[..num_rows]; + hashes_buf.fill(42_u64); // Compute hash for selected columns - create_murmur3_hashes(&columns_to_hash, hashes_buf)?; + create_xxh3_64_hashes(&columns_to_hash, hashes_buf)?; - // Assign partition IDs based on hash (same as hash partitioning) + // Assign partition IDs based on hash. XXH3-64's low 32 bits + // are fully avalanched, so truncating to u32 for pmod + // preserves distribution without widening pmod itself. let partition_ids = &mut scratch.partition_ids[..num_rows]; hashes_buf.iter().enumerate().for_each(|(idx, hash)| { partition_ids[idx] = - comet_partitioning::pmod(*hash, *num_output_partitions) as u32; + comet_partitioning::pmod(*hash as u32, *num_output_partitions) as u32; }); // We now have partition ids for every input row, map that to partition starts diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 6faa9fec4ec..e1ad4631dad 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -42,6 +42,7 @@ datafusion-comet-jni-bridge = { workspace = true } jni = "0.22.4" futures = { workspace = true } twox-hash = "2.1.2" +xxhash-rust = { version = "0.8", features = ["xxh3"] } rand = { workspace = true } base64 = "0.23.0" # shortest round-trip float formatting, used to match Spark's Double.toString-based diff --git a/native/spark-expr/src/hash_funcs/mod.rs b/native/spark-expr/src/hash_funcs/mod.rs index c6eba4a46bf..650e81f48e9 100644 --- a/native/spark-expr/src/hash_funcs/mod.rs +++ b/native/spark-expr/src/hash_funcs/mod.rs @@ -18,6 +18,7 @@ pub mod murmur3; pub(super) mod utils; mod xxhash64; +pub mod xxh3_64; pub use murmur3::spark_murmur3_hash; pub(crate) use xxhash64::create_xxhash64_hashes; diff --git a/native/spark-expr/src/hash_funcs/xxh3_64.rs b/native/spark-expr/src/hash_funcs/xxh3_64.rs new file mode 100644 index 00000000000..9e3947df4ca --- /dev/null +++ b/native/spark-expr/src/hash_funcs/xxh3_64.rs @@ -0,0 +1,87 @@ +// 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. + +use arrow::array::{Array, ArrayRef}; +use arrow::compute::take; +use datafusion::{ + arrow::{ + array::*, + datatypes::{ArrowDictionaryKeyType, ArrowNativeType}, + }, + error::{DataFusionError, Result}, +}; +use std::sync::Arc; +use xxhash_rust::xxh3::xxh3_64_with_seed; + +use crate::create_hashes_internal; + +/// One-shot XXH3-64 with a chained seed. Used by the shuffle partitioner's +/// `RoundRobin` arm to compute a per-row hash across columns. Not tied to any +/// Spark expression (Spark has no XXH3-64 hash). +#[inline] +fn xxh3_64_oneshot>(data: T, seed: u64) -> u64 { + xxh3_64_with_seed(data.as_ref(), seed) +} + +fn create_xxh3_64_hashes_dictionary( + array: &ArrayRef, + hashes_buffer: &mut [u64], + first_col: bool, +) -> Result<()> { + let dict_array = array.as_any().downcast_ref::>().unwrap(); + if !first_col { + let unpacked = take(dict_array.values().as_ref(), dict_array.keys(), None)?; + create_xxh3_64_hashes(&[unpacked], hashes_buffer)?; + } else { + // Hash each dictionary value once and reuse per key, so a large + // dictionary element (e.g. a long string) is not rehashed per row. + let dict_values = Arc::clone(dict_array.values()); + let mut dict_hashes = vec![42u64; dict_values.len()]; + create_xxh3_64_hashes(&[dict_values], &mut dict_hashes)?; + + for (hash, key) in hashes_buffer.iter_mut().zip(dict_array.keys().iter()) { + if let Some(key) = key { + let idx = key.to_usize().ok_or_else(|| { + DataFusionError::Internal(format!( + "Can not convert key value {:?} to usize in dictionary of type {:?}", + key, + dict_array.data_type() + )) + })?; + *hash = dict_hashes[idx] + } // no update for Null, consistent with other hashes + } + } + Ok(()) +} + +/// Compute a per-row XXH3-64 hash across `arrays`, chaining each column's +/// hash through the previous row hash. The number of rows is +/// `hashes_buffer.len()`. Callers must seed `hashes_buffer` before calling. +pub fn create_xxh3_64_hashes<'a>( + arrays: &[ArrayRef], + hashes_buffer: &'a mut [u64], +) -> Result<&'a mut [u64]> { + create_hashes_internal!( + arrays, + hashes_buffer, + xxh3_64_oneshot, + create_xxh3_64_hashes_dictionary, + create_xxh3_64_hashes + ); + Ok(hashes_buffer) +}