linkrs_query/executor/streaming/spill/
run_io.rs

1use std::io::{BufReader, BufWriter, Read, Write};
2use std::path::{Path, PathBuf};
3
4use linkrs_core::columnar::MaterializedBatch;
5use linkrs_core::error::QueryError;
6use linkrs_core::Value;
7
8use super::run_format::{
9    decode_section, encode_section, fnv1a_64_update, RunCompression, RunHeader, SpilledRun,
10    COMPRESSION_MIN_SIZE, FNV1A_64_INIT, RUN_HEADER_SIZE, RUN_VERSION, SECTION_ROWS,
11};
12
13// ── RunWriter ────────────────────────────────────────────────────────────────
14
15/// Writes a columnar run to disk with a versioned header, schema fingerprint,
16/// and per-section checksums.
17///
18/// Rows are staged in a columnar batch (bounded by [`SECTION_ROWS`]) and
19/// flushed as column-major sections, so writer memory stays flat regardless
20/// of run size. Each section is optionally zstd-compressed and individually
21/// checksummed; the header is patched in at finalize time.
22pub struct RunWriter {
23    pub(crate) writer: BufWriter<std::fs::File>,
24    pub(crate) path: PathBuf,
25    pub(crate) schema_fingerprint: u64,
26    pub(crate) row_count: u64,
27    pub(crate) section_count: u32,
28    pub(crate) body_bytes: u64,
29    pub(crate) staging: MaterializedBatch,
30}
31
32impl RunWriter {
33    /// Create a new run writer with header placeholder already written.
34    /// Callers should use `SpillManager::create_run_writer()` instead.
35    pub(crate) fn new(
36        writer: BufWriter<std::fs::File>,
37        path: PathBuf,
38        schema_fingerprint: u64,
39    ) -> Self {
40        Self {
41            writer,
42            path,
43            schema_fingerprint,
44            row_count: 0,
45            section_count: 0,
46            body_bytes: 0,
47            staging: MaterializedBatch::new(0, 0),
48        }
49    }
50
51    /// Write a single row to the run file (staged, flushed per section).
52    pub fn write_row(&mut self, row: &[Value]) -> Result<(), QueryError> {
53        if self.staging.num_columns() > 0 && row.len() != self.staging.num_columns() {
54            return Err(QueryError::execution(format!(
55                "spill run: row arity {} != run column count {}",
56                row.len(),
57                self.staging.num_columns()
58            )));
59        }
60        self.staging.append_row(row.to_vec());
61        if self.staging.num_rows() >= SECTION_ROWS {
62            self.flush_staging()?;
63        }
64        Ok(())
65    }
66
67    /// Write a batch of rows.
68    pub fn write_rows(&mut self, rows: &[Vec<Value>]) -> Result<(), QueryError> {
69        for row in rows {
70            self.write_row(row)?;
71        }
72        Ok(())
73    }
74
75    /// Write a columnar batch to the run file (staged per section).
76    pub fn write_batch(&mut self, batch: &MaterializedBatch) -> Result<(), QueryError> {
77        for row in batch.rows() {
78            self.write_row(&row)?;
79        }
80        Ok(())
81    }
82
83    /// Encode the staging batch as one columnar frame and append it.
84    fn flush_staging(&mut self) -> Result<(), QueryError> {
85        if self.staging.is_empty() {
86            return Ok(());
87        }
88        let payload = encode_section(&self.staging)?;
89        let (compression, stored) = if payload.len() >= COMPRESSION_MIN_SIZE {
90            match zstd::encode_all(payload.as_slice(), 3) {
91                Ok(compressed) if compressed.len() < payload.len() => {
92                    (RunCompression::Zstd, compressed)
93                }
94                _ => (RunCompression::None, payload),
95            }
96        } else {
97            (RunCompression::None, payload)
98        };
99        let flags = compression as u8;
100        let checksum = fnv1a_64_update(FNV1A_64_INIT, &[flags]);
101        let checksum = fnv1a_64_update(checksum, &stored);
102        self.writer
103            .write_all(&(stored.len() as u64).to_le_bytes())
104            .map_err(|e| QueryError::execution(format!("spill run: write section len: {}", e)))?;
105        self.writer
106            .write_all(&[flags])
107            .map_err(|e| QueryError::execution(format!("spill run: write section flags: {}", e)))?;
108        self.writer
109            .write_all(&stored)
110            .map_err(|e| QueryError::execution(format!("spill run: write section: {}", e)))?;
111        self.writer
112            .write_all(&checksum.to_le_bytes())
113            .map_err(|e| {
114                QueryError::execution(format!("spill run: write section checksum: {}", e))
115            })?;
116        self.body_bytes += 8 + 1 + stored.len() as u64 + 8;
117        self.row_count += self.staging.num_rows() as u64;
118        self.section_count += 1;
119        self.staging.clear();
120        Ok(())
121    }
122
123    /// Finalize the run: flush staging, patch the header, flush.
124    /// Returns `SpilledRun` metadata.
125    ///
126    /// Low-level primitive: production spill paths should route through
127    /// [`SpillManager::finalize_run`] instead so disk quota is enforced and
128    /// manager byte counters stay exact. Direct use remains for tests and the
129    /// no-manager fallback.
130    pub fn finalize(mut self) -> Result<SpilledRun, QueryError> {
131        use std::io::Seek;
132
133        self.flush_staging()?;
134        self.writer
135            .flush()
136            .map_err(|e| QueryError::execution(format!("spill run: flush body: {}", e)))?;
137
138        let header = RunHeader {
139            version: RUN_VERSION,
140            schema_fingerprint: self.schema_fingerprint,
141            row_count: self.row_count,
142            num_columns: self.staging.num_columns() as u32,
143            section_count: self.section_count,
144            flags: 0,
145        };
146
147        // Seek back to position 0 to overwrite header placeholder with actual header.
148        self.writer
149            .seek(std::io::SeekFrom::Start(0))
150            .map_err(|e| QueryError::execution(format!("spill run: seek: {}", e)))?;
151        let header_bytes = header.encode();
152        self.writer
153            .write_all(&header_bytes)
154            .map_err(|e| QueryError::execution(format!("spill run: write header: {}", e)))?;
155        self.writer
156            .flush()
157            .map_err(|e| QueryError::execution(format!("spill run: flush header: {}", e)))?;
158
159        let file_size = self.body_bytes + RUN_HEADER_SIZE as u64;
160        Ok(SpilledRun {
161            path: self.path.clone(),
162            row_count: self.row_count,
163            byte_size: file_size,
164            schema_fingerprint: self.schema_fingerprint,
165        })
166    }
167
168    pub fn path(&self) -> &Path {
169        &self.path
170    }
171
172    pub fn row_count(&self) -> u64 {
173        self.row_count
174    }
175}
176
177impl std::fmt::Debug for RunWriter {
178    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
179        f.debug_struct("RunWriter")
180            .field("path", &self.path)
181            .field("row_count", &self.row_count)
182            .field("body_bytes", &self.body_bytes)
183            .finish()
184    }
185}
186
187// ── RunReader ────────────────────────────────────────────────────────────────
188
189/// Reads a columnar run file back, validating the header (magic, version,
190/// checksum) on open and each section frame (checksum) on decode.
191///
192/// Sections stream one at a time (bounded by [`SECTION_ROWS`]), so reader
193/// memory stays flat regardless of run size.
194pub struct RunReader {
195    reader: BufReader<std::fs::File>,
196    path: PathBuf,
197    header: RunHeader,
198    remaining: u64,
199    section_rows: Vec<Vec<Value>>,
200    section_pos: usize,
201}
202
203impl RunReader {
204    /// Open and validate a run file.
205    pub fn open(run: &SpilledRun) -> Result<Self, QueryError> {
206        let mut f = std::fs::File::open(&run.path)
207            .map_err(|e| QueryError::execution(format!("spill run: open file: {}", e)))?;
208
209        // Read header
210        let mut header_buf = [0u8; RUN_HEADER_SIZE as usize];
211        f.read_exact(&mut header_buf)
212            .map_err(|e| QueryError::execution(format!("spill run: read header: {}", e)))?;
213        let header = RunHeader::decode(&header_buf)?;
214
215        // Validate schema fingerprint if provided
216        if run.schema_fingerprint != 0 && header.schema_fingerprint != run.schema_fingerprint {
217            return Err(QueryError::execution(format!(
218                "spill run: schema fingerprint mismatch: expected {}, got {}",
219                run.schema_fingerprint, header.schema_fingerprint
220            )));
221        }
222
223        // Validate row count if provided
224        if run.row_count != 0 && header.row_count != run.row_count {
225            return Err(QueryError::execution(format!(
226                "spill run: row count mismatch: expected {}, got {}",
227                run.row_count, header.row_count
228            )));
229        }
230
231        Ok(Self {
232            reader: BufReader::new(f),
233            path: run.path.clone(),
234            remaining: header.row_count,
235            header,
236            section_rows: Vec::new(),
237            section_pos: 0,
238        })
239    }
240
241    /// Open a run file with default validation (allows any schema fingerprint).
242    pub fn open_path(path: &Path) -> Result<Self, QueryError> {
243        let dummy = SpilledRun {
244            path: path.to_path_buf(),
245            row_count: 0,
246            byte_size: 0,
247            schema_fingerprint: 0,
248        };
249        Self::open(&dummy)
250    }
251
252    /// Decode the next section frame into the row buffer.
253    fn load_next_section(&mut self) -> Result<(), QueryError> {
254        let mut len_buf = [0u8; 8];
255        self.reader
256            .read_exact(&mut len_buf)
257            .map_err(|e| QueryError::execution(format!("spill run: read section len: {}", e)))?;
258        let payload_len = u64::from_le_bytes(len_buf) as usize;
259        if payload_len > 512 * 1024 * 1024 {
260            return Err(QueryError::execution(
261                "spill run: section payload length out of range".to_string(),
262            ));
263        }
264        let mut flags_buf = [0u8; 1];
265        self.reader
266            .read_exact(&mut flags_buf)
267            .map_err(|e| QueryError::execution(format!("spill run: read section flags: {}", e)))?;
268        let mut stored = vec![0u8; payload_len];
269        self.reader.read_exact(&mut stored).map_err(|e| {
270            QueryError::execution(format!("spill run: read section payload: {}", e))
271        })?;
272        let mut checksum_buf = [0u8; 8];
273        self.reader.read_exact(&mut checksum_buf).map_err(|e| {
274            QueryError::execution(format!("spill run: read section checksum: {}", e))
275        })?;
276        let expected = fnv1a_64_update(FNV1A_64_INIT, &flags_buf);
277        let expected = fnv1a_64_update(expected, &stored);
278        if expected != u64::from_le_bytes(checksum_buf) {
279            return Err(QueryError::execution(
280                "spill run: section checksum mismatch".to_string(),
281            ));
282        }
283        let payload: Vec<u8> = match flags_buf[0] {
284            0 => stored,
285            1 => zstd::decode_all(stored.as_slice())
286                .map_err(|e| QueryError::execution(format!("spill run: decompress: {}", e)))?,
287            other => {
288                return Err(QueryError::execution(format!(
289                    "spill run: unknown section compression {}",
290                    other
291                )));
292            }
293        };
294        let expected_cols = if self.header.num_columns > 0 {
295            Some(self.header.num_columns)
296        } else {
297            None
298        };
299        let mut batch = decode_section(&payload, expected_cols)?;
300        batch.set_fingerprint(self.header.schema_fingerprint);
301        self.section_rows = batch.into_rows();
302        self.section_pos = 0;
303        Ok(())
304    }
305
306    /// Read the next row from the run.
307    pub fn read_row(&mut self) -> Result<Option<Vec<Value>>, QueryError> {
308        if self.remaining == 0 {
309            return Ok(None);
310        }
311        if self.section_pos >= self.section_rows.len() {
312            self.load_next_section()?;
313        }
314        if self.section_pos >= self.section_rows.len() {
315            return Err(QueryError::execution(
316                "spill run: section row underflow".to_string(),
317            ));
318        }
319        let row = std::mem::take(&mut self.section_rows[self.section_pos]);
320        self.section_pos += 1;
321        self.remaining -= 1;
322        Ok(Some(row))
323    }
324
325    /// Read all remaining rows from the run.
326    pub fn read_all(&mut self) -> Result<Vec<Vec<Value>>, QueryError> {
327        let mut rows = Vec::with_capacity(self.remaining as usize);
328        while let Some(row) = self.read_row()? {
329            rows.push(row);
330        }
331        Ok(rows)
332    }
333
334    /// Read up to `n` rows into a columnar batch.
335    ///
336    /// Returns `None` when the run is exhausted. Bounds peak memory to one
337    /// output slice instead of the full run.
338    pub fn read_batch(&mut self, n: usize) -> Result<Option<MaterializedBatch>, QueryError> {
339        if self.remaining == 0 {
340            return Ok(None);
341        }
342        let mut batch = MaterializedBatch::new(0, self.header.schema_fingerprint);
343        let mut count = 0usize;
344        while count < n {
345            match self.read_row()? {
346                Some(row) => {
347                    batch.append_row(row);
348                    count += 1;
349                }
350                None => break,
351            }
352        }
353        if count == 0 {
354            Ok(None)
355        } else {
356            Ok(Some(batch))
357        }
358    }
359
360    pub fn header(&self) -> &RunHeader {
361        &self.header
362    }
363
364    pub fn remaining(&self) -> u64 {
365        self.remaining
366    }
367
368    pub fn path(&self) -> &Path {
369        &self.path
370    }
371}
372
373impl std::fmt::Debug for RunReader {
374    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
375        f.debug_struct("RunReader")
376            .field("path", &self.path)
377            .field("remaining", &self.remaining)
378            .field("header", &self.header)
379            .finish()
380    }
381}
linkrs_query/executor/streaming/spill/run_io.rs

linkrs_query/executor/streaming/spill/
run_io.rs
