feat: RR partition with xxhash-rust

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<u32>,
+    /// 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<u64>,
     /// Partition ids for each row in the current batch.
     partition_ids: Vec<u32>,
     /// The row indices of the rows in each partition. This array is conceptually divided into
@@ -225,11 +231,16 @@ impl<T: PartitionWriter> MultiPartitionShuffleRepartitioner<T> {
         // 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<T: PartitionWriter> MultiPartitionShuffleRepartitioner<T> {
                 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<u32>, &Vec<u32>) = {
                     let mut timer = self.metrics.repart_time.timer();
@@ -409,18 +421,19 @@ impl<T: PartitionWriter> MultiPartitionShuffleRepartitioner<T> {
                         .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

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

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;

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<T: AsRef<[u8]>>(data: T, seed: u64) -> u64 {
+    xxh3_64_with_seed(data.as_ref(), seed)
+}
+
+fn create_xxh3_64_hashes_dictionary<K: ArrowDictionaryKeyType>(
+    array: &ArrayRef,
+    hashes_buffer: &mut [u64],
+    first_col: bool,
+) -> Result<()> {
+    let dict_array = array.as_any().downcast_ref::<DictionaryArray<K>>().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)
+}