@@ -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.
439440pub 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+
441506impl 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