fraiseql_wire/stream/
chunking.rs

1//! Chunking logic for batching rows
2
3use bytes::Bytes;
4
5/// Row chunk (batch of raw JSON bytes)
6pub struct RowChunk {
7    rows: Vec<Bytes>,
8}
9
10impl RowChunk {
11    /// Create new chunk
12    pub const fn new() -> Self {
13        Self { rows: Vec::new() }
14    }
15
16    /// Create with capacity
17    pub fn with_capacity(capacity: usize) -> Self {
18        Self {
19            rows: Vec::with_capacity(capacity),
20        }
21    }
22
23    /// Add row to chunk
24    pub fn push(&mut self, row: Bytes) {
25        self.rows.push(row);
26    }
27
28    /// Check if chunk is empty
29    pub const fn is_empty(&self) -> bool {
30        self.rows.is_empty()
31    }
32
33    /// Get chunk size
34    pub const fn len(&self) -> usize {
35        self.rows.len()
36    }
37
38    /// Consume chunk and return rows
39    pub fn into_rows(self) -> Vec<Bytes> {
40        self.rows
41    }
42}
43
44impl Default for RowChunk {
45    fn default() -> Self {
46        Self::new()
47    }
48}
49
50/// Chunking strategy
51pub struct ChunkingStrategy {
52    chunk_size: usize,
53}
54
55impl ChunkingStrategy {
56    /// Create new strategy with given chunk size
57    pub const fn new(chunk_size: usize) -> Self {
58        Self { chunk_size }
59    }
60
61    /// Check if chunk is full
62    pub const fn is_full(&self, chunk: &RowChunk) -> bool {
63        chunk.len() >= self.chunk_size
64    }
65
66    /// Create new chunk with appropriate capacity
67    pub fn new_chunk(&self) -> RowChunk {
68        RowChunk::with_capacity(self.chunk_size)
69    }
70}
71
72impl Default for ChunkingStrategy {
73    fn default() -> Self {
74        Self::new(256)
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use super::*;
81
82    #[test]
83    fn test_chunk_operations() {
84        let mut chunk = RowChunk::new();
85        assert!(chunk.is_empty());
86
87        chunk.push(Bytes::from_static(b"{}"));
88        assert_eq!(chunk.len(), 1);
89        assert!(!chunk.is_empty());
90    }
91
92    #[test]
93    fn test_chunking_strategy() {
94        let strategy = ChunkingStrategy::new(2);
95        let mut chunk = strategy.new_chunk();
96
97        assert!(!strategy.is_full(&chunk));
98
99        chunk.push(Bytes::from_static(b"{}"));
100        assert!(!strategy.is_full(&chunk));
101
102        chunk.push(Bytes::from_static(b"{}"));
103        assert!(strategy.is_full(&chunk));
104    }
105}
fraiseql_wire/stream/chunking.rs

fraiseql_wire/stream/
chunking.rs
