laminar_connectors/kafka/partitioner/
mod.rs

1//! Kafka partitioning strategies.
2//!
3//! [`KafkaPartitioner`] determines which Kafka partition each record is
4//! sent to. Three strategies are provided:
5//!
6//! - [`KeyHashPartitioner`]: Murmur2 hash (Kafka-compatible default)
7//! - [`RoundRobinPartitioner`]: Cycle through partitions
8//! - [`StickyPartitioner`]: Batch to one partition until threshold
9
10/// Trait for determining the target Kafka partition for a record.
11///
12/// Implementations may be stateful (e.g., round-robin counter).
13pub trait KafkaPartitioner: Send + Sync {
14    /// Returns the target partition for the given key.
15    ///
16    /// Returns `None` if the partitioner defers to the broker (librdkafka)
17    /// default partitioning.
18    fn partition(&mut self, key: Option<&[u8]>, num_partitions: i32) -> Option<i32>;
19}
20
21/// Key-hash partitioner using Murmur2 (Kafka-compatible).
22///
23/// Produces the same partition assignment as Kafka's `DefaultPartitioner`
24/// for the same key bytes.
25#[derive(Debug, Default)]
26pub struct KeyHashPartitioner;
27
28impl KeyHashPartitioner {
29    /// Uses Kafka-compatible Murmur2 hash on the key bytes.
30    #[must_use]
31    pub fn new() -> Self {
32        Self
33    }
34}
35
36impl KafkaPartitioner for KeyHashPartitioner {
37    fn partition(&mut self, key: Option<&[u8]>, num_partitions: i32) -> Option<i32> {
38        key.map(|k| {
39            let hash = murmur2(k) & 0x7fff_ffff;
40            #[allow(
41                clippy::cast_possible_truncation,
42                clippy::cast_possible_wrap,
43                clippy::cast_sign_loss
44            )]
45            let partition = (hash % num_partitions as u32) as i32;
46            partition
47        })
48    }
49}
50
51/// Round-robin partitioner distributing records evenly across partitions.
52#[derive(Debug)]
53pub struct RoundRobinPartitioner {
54    counter: u64,
55}
56
57impl RoundRobinPartitioner {
58    /// Cycles through partitions sequentially, ignoring keys.
59    #[must_use]
60    pub fn new() -> Self {
61        Self { counter: 0 }
62    }
63}
64
65impl Default for RoundRobinPartitioner {
66    fn default() -> Self {
67        Self::new()
68    }
69}
70
71impl KafkaPartitioner for RoundRobinPartitioner {
72    #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
73    fn partition(&mut self, _key: Option<&[u8]>, num_partitions: i32) -> Option<i32> {
74        let partition = (self.counter % num_partitions as u64) as i32;
75        self.counter += 1;
76        Some(partition)
77    }
78}
79
80/// Sticky partitioner that batches records to the same partition.
81///
82/// Once `batch_threshold` records have been sent to one partition, rotates
83/// to the next. Reduces broker round-trips (Kafka KIP-794).
84#[derive(Debug)]
85pub struct StickyPartitioner {
86    current_partition: i32,
87    records_in_batch: usize,
88    batch_threshold: usize,
89}
90
91impl StickyPartitioner {
92    /// Creates a new sticky partitioner with the given batch threshold.
93    #[must_use]
94    pub fn new(batch_threshold: usize) -> Self {
95        Self {
96            current_partition: 0,
97            records_in_batch: 0,
98            batch_threshold,
99        }
100    }
101}
102
103impl KafkaPartitioner for StickyPartitioner {
104    fn partition(&mut self, _key: Option<&[u8]>, num_partitions: i32) -> Option<i32> {
105        if self.records_in_batch >= self.batch_threshold {
106            self.current_partition = (self.current_partition + 1) % num_partitions;
107            self.records_in_batch = 0;
108        }
109        self.records_in_batch += 1;
110        Some(self.current_partition)
111    }
112}
113
114/// Murmur2 hash function compatible with Kafka's `DefaultPartitioner`.
115///
116/// This is the 32-bit version used by Kafka for key-based partitioning.
117fn murmur2(data: &[u8]) -> u32 {
118    let seed: u32 = 0x9747_b28c;
119    let m: u32 = 0x5bd1_e995;
120    let r: u32 = 24;
121
122    let len = data.len();
123    #[allow(clippy::cast_possible_truncation)] // MurmurHash2 spec: seed XOR with 32-bit length
124    let mut h: u32 = seed ^ (len as u32);
125
126    let chunks = len / 4;
127    for i in 0..chunks {
128        let offset = i * 4;
129        let mut k = u32::from_le_bytes([
130            data[offset],
131            data[offset + 1],
132            data[offset + 2],
133            data[offset + 3],
134        ]);
135        k = k.wrapping_mul(m);
136        k ^= k >> r;
137        k = k.wrapping_mul(m);
138        h = h.wrapping_mul(m);
139        h ^= k;
140    }
141
142    let remainder = len % 4;
143    let tail_start = chunks * 4;
144    if remainder >= 3 {
145        h ^= u32::from(data[tail_start + 2]) << 16;
146    }
147    if remainder >= 2 {
148        h ^= u32::from(data[tail_start + 1]) << 8;
149    }
150    if remainder >= 1 {
151        h ^= u32::from(data[tail_start]);
152        h = h.wrapping_mul(m);
153    }
154
155    h ^= h >> 13;
156    h = h.wrapping_mul(m);
157    h ^= h >> 15;
158
159    h
160}
161
162#[cfg(test)]
163mod tests;
laminar_connectors/kafka/partitioner/mod.rs

laminar_connectors/kafka/partitioner/
mod.rs
