linkrs_query/executor/streaming/chunk/
pool.rs

1//! Row-buffer pool: reuse row `Vec<Value>` allocations across chunks.
2//!
3//! `DataChunk::clone` remains an explicit deep copy. New code prefers the
4//! move-first discipline (`expand_visible_rows`, `take_indices`, `slice`).
5//! The pool has two uses:
6//! - Multiplicity expansion (`expand_rows_reusing_buffers` in `selection.rs`)
7//!   clones repeated rows into recycled buffers instead of fresh allocations
8//!   (one pool lock per expansion; the common `multiplicity == 1` path never
9//!   touches the pool).
10//! - Residual one-by-one rebuilds: callers take drained rows via
11//!   [`DataChunk::take_rows_for_reuse`], return them with `release_rows`, and
12//!   reacquire them with `acquire_rows` / `acquire_row` instead of allocating
13//!   fresh vectors.
14//!
15//! The pool is process-global, lock-guarded, and bounded: at most
16//! [`MAX_POOLED_ROWS`] row buffers are retained, and oversized row buffers
17//! (capacity above [`MAX_POOLED_ROW_CAPACITY`]) are dropped instead of
18//! retained so a single wide row cannot pin memory.
19
20use linkrs_core::Value;
21use parking_lot::Mutex;
22use std::sync::LazyLock;
23
24pub const MAX_POOLED_ROWS: usize = 64;
25pub const MAX_POOLED_ROW_CAPACITY: usize = 1024;
26
27static ROW_POOL: LazyLock<Mutex<Vec<Vec<Value>>>> = LazyLock::new(|| Mutex::new(Vec::new()));
28
29/// Process-global row-buffer pool.
30pub struct RowBufferPool;
31
32impl RowBufferPool {
33    /// Take one reusable row buffer, or an empty vector when the pool is dry.
34    pub fn acquire_row() -> Vec<Value> {
35        ROW_POOL.lock().pop().unwrap_or_default()
36    }
37
38    /// Return one row buffer after `clear`. Oversized or surplus buffers are dropped.
39    pub fn release_row(row: &mut Vec<Value>) {
40        row.clear();
41        if row.capacity() > MAX_POOLED_ROW_CAPACITY {
42            return;
43        }
44        let mut pool = ROW_POOL.lock();
45        if pool.len() < MAX_POOLED_ROWS {
46            pool.push(std::mem::take(row));
47        }
48    }
49
50    /// Take up to `n` reusable row buffers.
51    pub fn acquire_rows(n: usize) -> Vec<Vec<Value>> {
52        let mut pool = ROW_POOL.lock();
53        let take = n.min(pool.len());
54        let mut out = Vec::with_capacity(n);
55        for _ in 0..take {
56            if let Some(row) = pool.pop() {
57                out.push(row);
58            }
59        }
60        out
61    }
62
63    /// Return drained rows to the pool, up to the capacity bound.
64    pub fn release_rows(rows: &mut Vec<Vec<Value>>) {
65        let mut pool = ROW_POOL.lock();
66        for row in rows.iter_mut() {
67            row.clear();
68            if row.capacity() > MAX_POOLED_ROW_CAPACITY {
69                continue;
70            }
71            if pool.len() >= MAX_POOLED_ROWS {
72                break;
73            }
74            pool.push(std::mem::take(row));
75        }
76        rows.clear();
77    }
78
79    /// Current pooled buffer count (observability / tests).
80    pub fn pool_len() -> usize {
81        ROW_POOL.lock().len()
82    }
83
84    /// Drain the pool (tests only).
85    #[cfg(test)]
86    pub fn clear_for_test() {
87        ROW_POOL.lock().clear();
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94
95    #[test]
96    fn acquire_release_reuses_capacity() {
97        RowBufferPool::clear_for_test();
98        let mut row = vec![Value::Int(1), Value::Int(2)];
99        RowBufferPool::release_row(&mut row);
100        assert!(row.is_empty());
101        assert_eq!(RowBufferPool::pool_len(), 1);
102        let reused = RowBufferPool::acquire_row();
103        assert_eq!(RowBufferPool::pool_len(), 0);
104        assert!(reused.capacity() >= 2);
105    }
106
107    #[test]
108    fn pool_drops_beyond_capacity() {
109        RowBufferPool::clear_for_test();
110        let mut rows: Vec<Vec<Value>> = (0..(MAX_POOLED_ROWS + 16))
111            .map(|_| Vec::<Value>::new())
112            .collect();
113        RowBufferPool::release_rows(&mut rows);
114        assert!(rows.is_empty());
115        assert_eq!(RowBufferPool::pool_len(), MAX_POOLED_ROWS);
116        RowBufferPool::clear_for_test();
117    }
118
119    #[test]
120    fn oversized_row_is_dropped() {
121        RowBufferPool::clear_for_test();
122        let mut row = Vec::with_capacity(MAX_POOLED_ROW_CAPACITY + 1);
123        row.push(Value::Int(1));
124        RowBufferPool::release_row(&mut row);
125        assert_eq!(RowBufferPool::pool_len(), 0);
126    }
127
128    #[test]
129    fn acquire_rows_partial_fill() {
130        RowBufferPool::clear_for_test();
131        let got = RowBufferPool::acquire_rows(4);
132        assert!(got.is_empty());
133        let mut rows = vec![vec![Value::Int(1)], vec![Value::Int(2)]];
134        RowBufferPool::release_rows(&mut rows);
135        let got = RowBufferPool::acquire_rows(4);
136        assert_eq!(got.len(), 2);
137        RowBufferPool::clear_for_test();
138    }
139}
linkrs_query/executor/streaming/chunk/pool.rs

linkrs_query/executor/streaming/chunk/
pool.rs
