Skip to content

Commit b0091c6

Browse files
committed
perf: strength reduce hash partition modulo
1 parent 65f337d commit b0091c6

1 file changed

Lines changed: 172 additions & 11 deletions

File tree

  • datafusion/physical-plan/src/repartition

‎datafusion/physical-plan/src/repartition/mod.rs‎

Lines changed: 172 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -425,6 +425,7 @@ enum BatchPartitionerState {
425425
Hash {
426426
exprs: Vec<Arc<dyn PhysicalExpr>>,
427427
num_partitions: usize,
428+
partition_reducer: StrengthReducedU64,
428429
hash_buffer: Vec<u64>,
429430
indices: Vec<Vec<u32>>,
430431
},
@@ -438,6 +439,70 @@ enum BatchPartitionerState {
438439
/// executions and runs.
439440
pub const REPARTITION_RANDOM_STATE: SeededRandomState = SeededRandomState::with_seed(0);
440441

442+
#[derive(Debug, Clone, Copy)]
443+
enum StrengthReducedU64 {
444+
PowerOfTwo { mask: u64 },
445+
Reciprocal { divisor: u64, reciprocal: u128 },
446+
}
447+
448+
impl StrengthReducedU64 {
449+
fn new(divisor: u64) -> Self {
450+
debug_assert!(divisor > 0);
451+
452+
if divisor.is_power_of_two() {
453+
Self::PowerOfTwo { mask: divisor - 1 }
454+
} else {
455+
Self::Reciprocal {
456+
divisor,
457+
// ceil(2^128 / divisor), computed without representing 2^128
458+
reciprocal: u128::MAX / u128::from(divisor) + 1,
459+
}
460+
}
461+
}
462+
463+
fn partition_indices(self, hash_buffer: &[u64], indices: &mut [Vec<u32>]) {
464+
match self {
465+
Self::PowerOfTwo { mask } => {
466+
for (index, hash) in hash_buffer.iter().enumerate() {
467+
indices[(*hash & mask) as usize].push(index as u32);
468+
}
469+
}
470+
Self::Reciprocal {
471+
divisor,
472+
reciprocal,
473+
} => {
474+
for (index, hash) in hash_buffer.iter().enumerate() {
475+
let quotient = Self::quotient(*hash, reciprocal);
476+
let partition = *hash - quotient * divisor;
477+
indices[partition as usize].push(index as u32);
478+
}
479+
}
480+
}
481+
}
482+
483+
#[cfg(test)]
484+
fn remainder(self, value: u64) -> u64 {
485+
match self {
486+
Self::PowerOfTwo { mask } => value & mask,
487+
Self::Reciprocal {
488+
divisor,
489+
reciprocal,
490+
} => value - Self::quotient(value, reciprocal) * divisor,
491+
}
492+
}
493+
494+
#[inline]
495+
fn quotient(value: u64, reciprocal: u128) -> u64 {
496+
let reciprocal_low = reciprocal as u64;
497+
let reciprocal_high = (reciprocal >> 64) as u64;
498+
let low_product = u128::from(value) * u128::from(reciprocal_low);
499+
let high_product = u128::from(value) * u128::from(reciprocal_high);
500+
let carry = ((high_product & u128::from(u64::MAX)) + (low_product >> 64)) >> 64;
501+
502+
((high_product >> 64) + carry) as u64
503+
}
504+
}
505+
441506
impl BatchPartitioner {
442507
/// Create a new [`BatchPartitioner`] for hash-based repartitioning.
443508
///
@@ -446,22 +511,27 @@ impl BatchPartitioner {
446511
/// - `num_partitions`: Total number of output partitions.
447512
/// - `timer`: Metric used to record time spent during repartitioning.
448513
///
449-
/// # Notes
450-
/// This constructor cannot fail and performs no validation.
514+
/// # Errors
515+
/// Returns an error if `num_partitions` is zero.
451516
pub fn new_hash_partitioner(
452517
exprs: Vec<Arc<dyn PhysicalExpr>>,
453518
num_partitions: usize,
454519
timer: metrics::Time,
455-
) -> Self {
456-
Self {
520+
) -> Result<Self> {
521+
if num_partitions == 0 {
522+
return internal_err!("Hash repartition requires at least one partition");
523+
}
524+
525+
Ok(Self {
457526
state: BatchPartitionerState::Hash {
458527
exprs,
459528
num_partitions,
529+
partition_reducer: StrengthReducedU64::new(num_partitions as u64),
460530
hash_buffer: vec![],
461531
indices: vec![vec![]; num_partitions],
462532
},
463533
timer,
464-
}
534+
})
465535
}
466536

467537
/// Create a new [`BatchPartitioner`] for round-robin repartitioning.
@@ -510,7 +580,7 @@ impl BatchPartitioner {
510580
) -> Result<Self> {
511581
match partitioning {
512582
Partitioning::Hash(exprs, num_partitions) => {
513-
Ok(Self::new_hash_partitioner(exprs, num_partitions, timer))
583+
Self::new_hash_partitioner(exprs, num_partitions, timer)
514584
}
515585
Partitioning::RoundRobinBatch(num_partitions) => {
516586
Ok(Self::new_round_robin_partitioner(
@@ -575,7 +645,8 @@ impl BatchPartitioner {
575645
}
576646
BatchPartitionerState::Hash {
577647
exprs,
578-
num_partitions: partitions,
648+
num_partitions: _,
649+
partition_reducer,
579650
hash_buffer,
580651
indices,
581652
} => {
@@ -596,9 +667,7 @@ impl BatchPartitioner {
596667

597668
indices.iter_mut().for_each(|v| v.clear());
598669

599-
for (index, hash) in hash_buffer.iter().enumerate() {
600-
indices[(*hash % *partitions as u64) as usize].push(index as u32);
601-
}
670+
partition_reducer.partition_indices(hash_buffer, indices);
602671

603672
// Finished building index-arrays for output partitions
604673
timer.done();
@@ -1359,7 +1428,7 @@ impl RepartitionExec {
13591428
exprs.clone(),
13601429
*num_partitions,
13611430
metrics.repartition_time.clone(),
1362-
)
1431+
)?
13631432
}
13641433
Partitioning::RoundRobinBatch(num_partitions) => {
13651434
BatchPartitioner::new_round_robin_partitioner(
@@ -1785,6 +1854,98 @@ mod tests {
17851854
use datafusion_execution::runtime_env::RuntimeEnvBuilder;
17861855
use insta::assert_snapshot;
17871856

1857+
#[test]
1858+
fn strength_reduced_u64_remainder_matches_modulo() {
1859+
let divisors = [
1860+
1,
1861+
2,
1862+
3,
1863+
4,
1864+
5,
1865+
7,
1866+
8,
1867+
10,
1868+
16,
1869+
31,
1870+
32,
1871+
63,
1872+
64,
1873+
65,
1874+
97,
1875+
u64::from(u32::MAX),
1876+
u64::from(u32::MAX) + 1,
1877+
1_u64 << 32,
1878+
(1_u64 << 63) - 1,
1879+
1_u64 << 63,
1880+
u64::MAX - 1,
1881+
u64::MAX,
1882+
];
1883+
let values = [
1884+
0,
1885+
1,
1886+
2,
1887+
3,
1888+
4,
1889+
5,
1890+
31,
1891+
32,
1892+
33,
1893+
63,
1894+
64,
1895+
65,
1896+
u64::from(u32::MAX) - 1,
1897+
u64::from(u32::MAX),
1898+
u64::from(u32::MAX) + 1,
1899+
(1_u64 << 32) - 1,
1900+
1_u64 << 32,
1901+
(1_u64 << 32) + 1,
1902+
(1_u64 << 63) - 1,
1903+
1_u64 << 63,
1904+
(1_u64 << 63) + 1,
1905+
u64::MAX - 1,
1906+
u64::MAX,
1907+
];
1908+
1909+
for divisor in divisors {
1910+
let reducer = StrengthReducedU64::new(divisor);
1911+
for value in values {
1912+
assert_eq!(
1913+
reducer.remainder(value),
1914+
value % divisor,
1915+
"value={value} divisor={divisor}"
1916+
);
1917+
}
1918+
1919+
let mut value = 0x1234_5678_9abc_def0 ^ divisor;
1920+
for _ in 0..10_000 {
1921+
value = value
1922+
.wrapping_mul(6_364_136_223_846_793_005)
1923+
.wrapping_add(1_442_695_040_888_963_407);
1924+
assert_eq!(
1925+
reducer.remainder(value),
1926+
value % divisor,
1927+
"value={value} divisor={divisor}"
1928+
);
1929+
}
1930+
}
1931+
}
1932+
1933+
#[test]
1934+
fn hash_partitioner_requires_nonzero_partitions() {
1935+
let metrics = ExecutionPlanMetricsSet::new();
1936+
let timer = MetricBuilder::new(&metrics).subset_time("test", 0);
1937+
1938+
let err = BatchPartitioner::new_hash_partitioner(vec![], 0, timer)
1939+
.err()
1940+
.expect("zero hash partitions should fail")
1941+
.to_string();
1942+
1943+
assert!(
1944+
err.contains("Hash repartition requires at least one partition"),
1945+
"actual: {err}"
1946+
);
1947+
}
1948+
17881949
#[tokio::test]
17891950
async fn one_to_many_round_robin() -> Result<()> {
17901951
// define input partitions

0 commit comments

Comments
 (0)