laminar_connectors/kafka/
partitioner.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 {
164    use super::*;
165
166    #[test]
167    fn test_murmur2_known_values() {
168        // Empty key: seed goes through final mixing.
169        let h = murmur2(b"");
170        assert_ne!(h, 0);
171
172        // Known values consistent with Kafka's Java implementation.
173        let h = murmur2(b"key1");
174        assert_ne!(h, 0);
175    }
176
177    #[test]
178    fn test_murmur2_deterministic() {
179        let h1 = murmur2(b"test-key");
180        let h2 = murmur2(b"test-key");
181        assert_eq!(h1, h2);
182
183        let h3 = murmur2(b"different-key");
184        assert_ne!(h1, h3);
185    }
186
187    #[test]
188    fn test_key_hash_partitioner_with_key() {
189        let mut p = KeyHashPartitioner::new();
190        let partition = p.partition(Some(b"order-123"), 6);
191        assert!(partition.is_some());
192        let part = partition.unwrap();
193        assert!((0..6).contains(&part));
194
195        // Same key → same partition
196        let partition2 = p.partition(Some(b"order-123"), 6);
197        assert_eq!(partition, partition2);
198    }
199
200    #[test]
201    fn test_key_hash_partitioner_no_key() {
202        let mut p = KeyHashPartitioner::new();
203        assert_eq!(p.partition(None, 6), None);
204    }
205
206    #[test]
207    fn test_round_robin_partitioner() {
208        let mut p = RoundRobinPartitioner::new();
209        assert_eq!(p.partition(None, 3), Some(0));
210        assert_eq!(p.partition(None, 3), Some(1));
211        assert_eq!(p.partition(None, 3), Some(2));
212        assert_eq!(p.partition(None, 3), Some(0)); // wraps
213    }
214
215    #[test]
216    fn test_round_robin_ignores_key() {
217        let mut p = RoundRobinPartitioner::new();
218        assert_eq!(p.partition(Some(b"key"), 3), Some(0));
219        assert_eq!(p.partition(Some(b"key"), 3), Some(1));
220    }
221
222    #[test]
223    fn test_sticky_partitioner() {
224        let mut p = StickyPartitioner::new(3);
225
226        // First 3 records go to partition 0
227        assert_eq!(p.partition(None, 4), Some(0));
228        assert_eq!(p.partition(None, 4), Some(0));
229        assert_eq!(p.partition(None, 4), Some(0));
230
231        // 4th record rotates to partition 1
232        assert_eq!(p.partition(None, 4), Some(1));
233        assert_eq!(p.partition(None, 4), Some(1));
234        assert_eq!(p.partition(None, 4), Some(1));
235
236        // 7th record rotates to partition 2
237        assert_eq!(p.partition(None, 4), Some(2));
238    }
239}
laminar_connectors/kafka/partitioner.rs

laminar_connectors/kafka/
partitioner.rs
