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

Author: Dandandan

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