You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/04/30 11:39:32 UTC
(datafusion) 01/01: perf: strength reduce hash partition modulo (up to 1.16x faster) (#21900)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-21900-13642865b526b32fbef8d6c29c1f5db4abe8c6b6
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 3deeadd00207addd026dc16bbdccdd198bfd24f7
Author: Daniël Heres <da...@gmail.com>
AuthorDate: Thu Apr 30 13:39:09 2026 +0200
perf: strength reduce hash partition modulo (up to 1.16x faster) (#21900)
## Which issue does this PR close?
- Closes #21843.
## Rationale for this change
Performance improvement on large hash-repartitions.
TPC-H at bigger scale factor shows biggest benefit:
<details>
```
--------------------
Benchmark tpch_sf10.json
--------------------
┏━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Query ┃ HEAD ┃ perf-strength-reduce-hash-partition ┃ Change ┃
┡━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ QQuery 1 │ 327.64 / 329.66 ±1.44 / 331.84 ms │ 327.82 / 330.00 ±1.68 / 331.48 ms │ no change │
│ QQuery 2 │ 131.35 / 138.32 ±3.91 / 143.24 ms │ 125.94 / 126.49 ±0.61 / 127.42 ms │ +1.09x faster │
│ QQuery 3 │ 286.47 / 300.24 ±8.19 / 308.99 ms │ 273.76 / 276.01 ±1.80 / 277.98 ms │ +1.09x faster │
│ QQuery 4 │ 158.30 / 160.81 ±2.21 / 163.54 ms │ 137.56 / 138.79 ±0.98 / 139.80 ms │ +1.16x faster │
│ QQuery 5 │ 428.90 / 437.68 ±4.45 / 440.83 ms │ 390.52 / 396.51 ±4.30 / 403.67 ms │ +1.10x faster │
│ QQuery 6 │ 131.88 / 132.83 ±1.17 / 134.81 ms │ 133.06 / 134.70 ±1.22 / 135.83 ms │ no change │
│ QQuery 7 │ 541.09 / 545.88 ±4.21 / 552.67 ms │ 508.51 / 531.75 ±16.82 / 548.21 ms │ no change │
│ QQuery 8 │ 467.86 / 476.44 ±6.95 / 483.87 ms │ 427.19 / 439.03 ±9.88 / 453.56 ms │ +1.09x faster │
│ QQuery 9 │ 649.16 / 660.07 ±10.12 / 676.72 ms │ 605.25 / 611.87 ±5.72 / 620.70 ms │ +1.08x faster │
│ QQuery 10 │ 327.64 / 339.90 ±6.92 / 348.85 ms │ 321.66 / 330.67 ±4.76 / 334.89 ms │ no change │
│ QQuery 11 │ 104.93 / 107.54 ±1.71 / 110.18 ms │ 92.80 / 101.35 ±12.27 / 125.63 ms │ +1.06x faster │
│ QQuery 12 │ 198.96 / 202.37 ±2.45 / 206.21 ms │ 195.07 / 197.77 ±4.26 / 206.26 ms │ no change │
│ QQuery 13 │ 300.44 / 312.37 ±6.87 / 321.90 ms │ 291.85 / 308.47 ±10.23 / 317.55 ms │ no change │
│ QQuery 14 │ 188.06 / 193.71 ±4.81 / 200.69 ms │ 182.89 / 186.72 ±3.67 / 192.75 ms │ no change │
│ QQuery 15 │ 334.88 / 339.95 ±5.81 / 350.78 ms │ 330.79 / 336.21 ±4.31 / 342.71 ms │ no change │
│ QQuery 16 │ 78.38 / 81.25 ±2.51 / 84.55 ms │ 74.35 / 76.61 ±2.80 / 81.94 ms │ +1.06x faster │
│ QQuery 17 │ 744.08 / 761.70 ±12.84 / 781.69 ms │ 703.40 / 724.66 ±23.35 / 770.05 ms │ no change │
│ QQuery 18 │ 760.17 / 782.23 ±12.12 / 796.85 ms │ 725.45 / 744.71 ±15.59 / 765.59 ms │ no change │
│ QQuery 19 │ 267.90 / 280.99 ±14.61 / 306.80 ms │ 275.58 / 298.23 ±27.69 / 351.75 ms │ 1.06x slower │
│ QQuery 20 │ 311.46 / 323.12 ±10.13 / 341.26 ms │ 312.13 / 319.42 ±4.39 / 324.46 ms │ no change │
│ QQuery 21 │ 816.40 / 837.33 ±19.78 / 870.18 ms │ 766.23 / 778.58 ±8.98 / 792.31 ms │ +1.08x faster │
│ QQuery 22 │ 81.46 / 84.94 ±2.58 / 88.20 ms │ 75.31 / 77.73 ±1.39 / 79.55 ms │ +1.09x faster │
└───────────┴────────────────────────────────────┴─────────────────────────────────────┴───────────────┘
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━┓
┃ Benchmark Summary ┃ ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━┩
│ Total Time (HEAD) │ 7829.34ms │
│ Total Time (perf-strength-reduce-hash-partition) │ 7466.29ms │
│ Average Time (HEAD) │ 355.88ms │
│ Average Time (perf-strength-reduce-hash-partition) │ 339.38ms │
│ Queries Faster │ 10 │
│ Queries Slower │ 1 │
│ Queries with No Change │ 11 │
│ Queries with Failure │ 0 │
└────────────────────────────────────────────────────┴───────────┘
```
</details>
## What changes are included in this PR?
Use strength-reduce to speed up hash % partition
## Are these changes tested?
Existing tests
## Are there any user-facing changes?
A small change to `new_hash_partitioner` to return a `Result` instead of
panic during runtime
---
datafusion/physical-plan/src/repartition/mod.rs | 204 ++++++++++++++++++++++--
1 file changed, 187 insertions(+), 17 deletions(-)
diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs
index f5d5d6dbd0..b4af6e2c09 100644
--- a/datafusion/physical-plan/src/repartition/mod.rs
+++ b/datafusion/physical-plan/src/repartition/mod.rs
@@ -424,7 +424,7 @@ pub struct BatchPartitioner {
enum BatchPartitionerState {
Hash {
exprs: Vec<Arc<dyn PhysicalExpr>>,
- num_partitions: usize,
+ partition_reducer: StrengthReducedU64,
hash_buffer: Vec<u64>,
indices: Vec<Vec<u32>>,
},
@@ -438,6 +438,78 @@ enum BatchPartitionerState {
/// executions and runs.
pub const REPARTITION_RANDOM_STATE: SeededRandomState = SeededRandomState::with_seed(0);
+/// Computes `value % divisor` without division in the hot loop when `divisor`
+/// is fixed for many values.
+///
+/// Hash repartitioning computes a remainder for every row. Integer division is
+/// relatively expensive, so this precomputes the strength-reduced form of the
+/// divisor: powers of two use a bit mask, and other divisors use a reciprocal
+/// multiply to recover the quotient and therefore the remainder. This is the
+/// same invariant-divisor optimization compilers use for `%` by a constant.
+#[derive(Debug, Clone, Copy)]
+enum StrengthReducedU64 {
+ PowerOfTwo { mask: u64 },
+ Reciprocal { divisor: u64, reciprocal: u128 },
+}
+
+impl StrengthReducedU64 {
+ fn new(divisor: u64) -> Self {
+ debug_assert!(divisor > 0);
+
+ if divisor.is_power_of_two() {
+ Self::PowerOfTwo { mask: divisor - 1 }
+ } else {
+ Self::Reciprocal {
+ divisor,
+ // ceil(2^128 / divisor), computed without representing 2^128
+ reciprocal: u128::MAX / u128::from(divisor) + 1,
+ }
+ }
+ }
+
+ fn partition_indices(self, hash_buffer: &[u64], indices: &mut [Vec<u32>]) {
+ match self {
+ Self::PowerOfTwo { mask } => {
+ for (index, hash) in hash_buffer.iter().enumerate() {
+ indices[(*hash & mask) as usize].push(index as u32);
+ }
+ }
+ Self::Reciprocal {
+ divisor,
+ reciprocal,
+ } => {
+ for (index, hash) in hash_buffer.iter().enumerate() {
+ let quotient = Self::quotient(*hash, reciprocal);
+ let partition = *hash - quotient * divisor;
+ indices[partition as usize].push(index as u32);
+ }
+ }
+ }
+ }
+
+ #[cfg(test)]
+ fn remainder(self, value: u64) -> u64 {
+ match self {
+ Self::PowerOfTwo { mask } => value & mask,
+ Self::Reciprocal {
+ divisor,
+ reciprocal,
+ } => value - Self::quotient(value, reciprocal) * divisor,
+ }
+ }
+
+ #[inline]
+ fn quotient(value: u64, reciprocal: u128) -> u64 {
+ let reciprocal_low = reciprocal as u64;
+ let reciprocal_high = (reciprocal >> 64) as u64;
+ let low_product = u128::from(value) * u128::from(reciprocal_low);
+ let high_product = u128::from(value) * u128::from(reciprocal_high);
+ let carry = ((high_product & u128::from(u64::MAX)) + (low_product >> 64)) >> 64;
+
+ ((high_product >> 64) + carry) as u64
+ }
+}
+
impl BatchPartitioner {
/// Create a new [`BatchPartitioner`] for hash-based repartitioning.
///
@@ -446,22 +518,29 @@ impl BatchPartitioner {
/// - `num_partitions`: Total number of output partitions.
/// - `timer`: Metric used to record time spent during repartitioning.
///
- /// # Notes
- /// This constructor cannot fail and performs no validation.
+ /// The partition count is fixed for the lifetime of the partitioner, so this
+ /// precomputes a strength-reduced reducer for `hash % num_partitions`.
+ ///
+ /// # Errors
+ /// Returns an error if `num_partitions` is zero.
pub fn new_hash_partitioner(
exprs: Vec<Arc<dyn PhysicalExpr>>,
num_partitions: usize,
timer: metrics::Time,
- ) -> Self {
- Self {
+ ) -> Result<Self> {
+ if num_partitions == 0 {
+ return internal_err!("Hash repartition requires at least one partition");
+ }
+
+ Ok(Self {
state: BatchPartitionerState::Hash {
exprs,
- num_partitions,
+ partition_reducer: StrengthReducedU64::new(num_partitions as u64),
hash_buffer: vec![],
indices: vec![vec![]; num_partitions],
},
timer,
- }
+ })
}
/// Create a new [`BatchPartitioner`] for round-robin repartitioning.
@@ -501,7 +580,8 @@ impl BatchPartitioner {
/// - `num_input_partitions`: Total number of input partitions.
///
/// # Errors
- /// Returns an error if the provided partitioning scheme is not supported.
+ /// Returns an error if the provided partitioning scheme is not supported,
+ /// or if hash partitioning is requested with zero output partitions.
pub fn try_new(
partitioning: Partitioning,
timer: metrics::Time,
@@ -510,7 +590,7 @@ impl BatchPartitioner {
) -> Result<Self> {
match partitioning {
Partitioning::Hash(exprs, num_partitions) => {
- Ok(Self::new_hash_partitioner(exprs, num_partitions, timer))
+ Self::new_hash_partitioner(exprs, num_partitions, timer)
}
Partitioning::RoundRobinBatch(num_partitions) => {
Ok(Self::new_round_robin_partitioner(
@@ -575,7 +655,7 @@ impl BatchPartitioner {
}
BatchPartitionerState::Hash {
exprs,
- num_partitions: partitions,
+ partition_reducer,
hash_buffer,
indices,
} => {
@@ -596,9 +676,7 @@ impl BatchPartitioner {
indices.iter_mut().for_each(|v| v.clear());
- for (index, hash) in hash_buffer.iter().enumerate() {
- indices[(*hash % *partitions as u64) as usize].push(index as u32);
- }
+ partition_reducer.partition_indices(hash_buffer, indices);
// Finished building index-arrays for output partitions
timer.done();
@@ -653,9 +731,9 @@ impl BatchPartitioner {
// return the number of output partitions
fn num_partitions(&self) -> usize {
- match self.state {
- BatchPartitionerState::RoundRobin { num_partitions, .. } => num_partitions,
- BatchPartitionerState::Hash { num_partitions, .. } => num_partitions,
+ match &self.state {
+ BatchPartitionerState::RoundRobin { num_partitions, .. } => *num_partitions,
+ BatchPartitionerState::Hash { indices, .. } => indices.len(),
}
}
}
@@ -1359,7 +1437,7 @@ impl RepartitionExec {
exprs.clone(),
*num_partitions,
metrics.repartition_time.clone(),
- )
+ )?
}
Partitioning::RoundRobinBatch(num_partitions) => {
BatchPartitioner::new_round_robin_partitioner(
@@ -1783,6 +1861,98 @@ mod tests {
use datafusion_execution::runtime_env::RuntimeEnvBuilder;
use insta::assert_snapshot;
+ #[test]
+ fn strength_reduced_u64_remainder_matches_modulo() {
+ let divisors = [
+ 1,
+ 2,
+ 3,
+ 4,
+ 5,
+ 7,
+ 8,
+ 10,
+ 16,
+ 31,
+ 32,
+ 63,
+ 64,
+ 65,
+ 97,
+ u64::from(u32::MAX),
+ u64::from(u32::MAX) + 1,
+ 1_u64 << 32,
+ (1_u64 << 63) - 1,
+ 1_u64 << 63,
+ u64::MAX - 1,
+ u64::MAX,
+ ];
+ let values = [
+ 0,
+ 1,
+ 2,
+ 3,
+ 4,
+ 5,
+ 31,
+ 32,
+ 33,
+ 63,
+ 64,
+ 65,
+ u64::from(u32::MAX) - 1,
+ u64::from(u32::MAX),
+ u64::from(u32::MAX) + 1,
+ (1_u64 << 32) - 1,
+ 1_u64 << 32,
+ (1_u64 << 32) + 1,
+ (1_u64 << 63) - 1,
+ 1_u64 << 63,
+ (1_u64 << 63) + 1,
+ u64::MAX - 1,
+ u64::MAX,
+ ];
+
+ for divisor in divisors {
+ let reducer = StrengthReducedU64::new(divisor);
+ for value in values {
+ assert_eq!(
+ reducer.remainder(value),
+ value % divisor,
+ "value={value} divisor={divisor}"
+ );
+ }
+
+ let mut value = 0x1234_5678_9abc_def0 ^ divisor;
+ for _ in 0..10_000 {
+ value = value
+ .wrapping_mul(6_364_136_223_846_793_005)
+ .wrapping_add(1_442_695_040_888_963_407);
+ assert_eq!(
+ reducer.remainder(value),
+ value % divisor,
+ "value={value} divisor={divisor}"
+ );
+ }
+ }
+ }
+
+ #[test]
+ fn hash_partitioner_requires_nonzero_partitions() {
+ let metrics = ExecutionPlanMetricsSet::new();
+ let timer = MetricBuilder::new(&metrics).subset_time("test", 0);
+
+ let err = BatchPartitioner::new_hash_partitioner(vec![], 0, timer)
+ .err()
+ .expect("zero hash partitions should fail")
+ .to_string();
+
+ assert!(
+ err.contains("Hash repartition requires at least one partition"),
+ "actual: {err}"
+ );
+ }
+
#[tokio::test]
async fn one_to_many_round_robin() -> Result<()> {
// define input partitions
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org