Skip to main content

linkrs_query/executor/streaming/spill/
run_format.rs

1use std::path::PathBuf;
2
3use linkrs_core::columnar::MaterializedBatch;
4use linkrs_core::error::QueryError;
5use linkrs_core::Value;
6
7// ── Run file format constants ───────────────────────────────────────────────
8
9/// Magic bytes at the start of every spill run file: `LRSC` = Linkrs Run
10pub(super) const RUN_MAGIC: [u8; 4] = [0x4C, 0x52, 0x53, 0x43];
11
12/// Current run file format version (columnar v2).
13pub(super) const RUN_VERSION: u32 = 3;
14
15/// Size of the run file header in bytes.
16pub(super) const RUN_HEADER_SIZE: u32 = 48;
17
18/// Rows staged per columnar section. Bounds writer staging memory and keeps
19/// each section independently checksummed and decompressible.
20pub(super) const SECTION_ROWS: usize = 1024;
21
22/// Minimum section payload size (bytes) below which compression is skipped.
23pub(super) const COMPRESSION_MIN_SIZE: usize = 256;
24
25// ── Simple FNV-1a 64-bit checksum ───────────────────────────────────────────
26
27pub(super) const FNV1A_64_INIT: u64 = 0xcbf29ce484222325;
28const FNV1A_64_PRIME: u64 = 0x100000001b3;
29
30fn fnv1a_64(data: &[u8]) -> u64 {
31    fnv1a_64_update(FNV1A_64_INIT, data)
32}
33
34pub(super) fn fnv1a_64_update(mut hash: u64, data: &[u8]) -> u64 {
35    for &byte in data {
36        hash ^= byte as u64;
37        hash = hash.wrapping_mul(FNV1A_64_PRIME);
38    }
39    hash
40}
41
42// ── RunHeader ────────────────────────────────────────────────────────────────
43
44/// Compression marker stored per section frame.
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46#[repr(u8)]
47pub enum RunCompression {
48    None = 0,
49    Zstd = 1,
50}
51
52/// Header for a columnar spill run file.
53///
54/// Layout (48 bytes total):
55/// ```text
56/// [0..4)   magic: b"LRSC"
57/// [4..8)   version: u32 LE (= 3)
58/// [8..16)  schema_fingerprint: u64 LE
59/// [16..24) row_count: u64 LE
60/// [24..28) num_columns: u32 LE
61/// [28..32) section_count: u32 LE
62/// [32..36) flags: u32 LE (reserved, zero)
63/// [36..40) reserved: u32 LE (zero)
64/// [40..48) header_checksum: u64 LE (FNV-1a of bytes [0..40))
65/// ```
66/// Body: `section_count` frames back to back. Each frame:
67/// `[payload_len u64 LE][flags u8][payload][checksum u64 LE]`
68/// (`flags` bit 0 = zstd; checksum = FNV-1a over flags byte + stored payload
69/// bytes). Each (decompressed) payload:
70/// `[num_rows u32 LE][num_cols u32 LE]` then per column
71/// `[col_len u64 LE][postcard(Vec<Value>) bytes]` — columns contiguous.
72/// Values stay postcard self-describing; column types are intentionally not
73/// stored (fingerprint + column count validate replay compatibility).
74#[derive(Debug, Clone, Copy)]
75pub struct RunHeader {
76    pub version: u32,
77    pub schema_fingerprint: u64,
78    pub row_count: u64,
79    pub num_columns: u32,
80    pub section_count: u32,
81    pub flags: u32,
82}
83
84impl RunHeader {
85    /// Encode header into a 48-byte buffer.
86    pub(super) fn encode(&self) -> [u8; RUN_HEADER_SIZE as usize] {
87        let mut buf = [0u8; RUN_HEADER_SIZE as usize];
88        buf[0..4].copy_from_slice(&RUN_MAGIC);
89        buf[4..8].copy_from_slice(&self.version.to_le_bytes());
90        buf[8..16].copy_from_slice(&self.schema_fingerprint.to_le_bytes());
91        buf[16..24].copy_from_slice(&self.row_count.to_le_bytes());
92        buf[24..28].copy_from_slice(&self.num_columns.to_le_bytes());
93        buf[28..32].copy_from_slice(&self.section_count.to_le_bytes());
94        buf[32..36].copy_from_slice(&self.flags.to_le_bytes());
95        // bytes 36..40 stay zero (reserved)
96        let checksum = fnv1a_64(&buf[0..40]);
97        buf[40..48].copy_from_slice(&checksum.to_le_bytes());
98        buf
99    }
100
101    /// Decode header from a 48-byte buffer, validating magic, version, and
102    /// header checksum. Legacy row-major files are rejected.
103    pub(super) fn decode(buf: &[u8]) -> Result<Self, QueryError> {
104        if buf.len() < RUN_HEADER_SIZE as usize {
105            return Err(QueryError::execution(
106                "spill run: truncated header".to_string(),
107            ));
108        }
109        if buf[0..4] != RUN_MAGIC {
110            return Err(QueryError::execution(
111                "spill run: invalid magic bytes".to_string(),
112            ));
113        }
114        let version = u32::from_le_bytes(buf[4..8].try_into().unwrap());
115        if version != RUN_VERSION {
116            return Err(QueryError::execution(format!(
117                "spill run: unsupported version {}, expected {}",
118                version, RUN_VERSION
119            )));
120        }
121        let expected = fnv1a_64(&buf[0..40]);
122        let actual = u64::from_le_bytes(buf[40..48].try_into().unwrap());
123        if expected != actual {
124            return Err(QueryError::execution(
125                "spill run: header checksum mismatch".to_string(),
126            ));
127        }
128        Ok(Self {
129            version,
130            schema_fingerprint: u64::from_le_bytes(buf[8..16].try_into().unwrap()),
131            row_count: u64::from_le_bytes(buf[16..24].try_into().unwrap()),
132            num_columns: u32::from_le_bytes(buf[24..28].try_into().unwrap()),
133            section_count: u32::from_le_bytes(buf[28..32].try_into().unwrap()),
134            flags: u32::from_le_bytes(buf[32..36].try_into().unwrap()),
135        })
136    }
137}
138
139/// Encode one staging batch as a columnar section payload (uncompressed).
140pub(super) fn encode_section(batch: &MaterializedBatch) -> Result<Vec<u8>, QueryError> {
141    let num_rows = batch.num_rows();
142    let num_cols = batch.num_columns();
143    let mut payload = Vec::new();
144    payload.extend_from_slice(&(num_rows as u32).to_le_bytes());
145    payload.extend_from_slice(&(num_cols as u32).to_le_bytes());
146    for i in 0..num_cols {
147        let encoded = postcard::to_allocvec(batch.column_slice(i))
148            .map_err(|e| QueryError::execution(format!("run serialize column: {}", e)))?;
149        payload.extend_from_slice(&(encoded.len() as u64).to_le_bytes());
150        payload.extend_from_slice(&encoded);
151    }
152    Ok(payload)
153}
154
155/// Decode one section payload into a batch, validating column count.
156pub(super) fn decode_section(
157    payload: &[u8],
158    expected_cols: Option<u32>,
159) -> Result<MaterializedBatch, QueryError> {
160    let fail = |msg: String| QueryError::execution(format!("spill run: {msg}"));
161    if payload.len() < 8 {
162        return Err(fail("truncated section payload".to_string()));
163    }
164    let num_rows = u32::from_le_bytes(payload[0..4].try_into().unwrap()) as usize;
165    let num_cols = u32::from_le_bytes(payload[4..8].try_into().unwrap());
166    if num_rows > 10_000_000 || num_cols > 100_000 {
167        return Err(fail("section dimensions out of range".to_string()));
168    }
169    if let Some(expected) = expected_cols {
170        if num_cols != expected {
171            return Err(fail(format!(
172                "section column count {} != run column count {}",
173                num_cols, expected
174            )));
175        }
176    }
177    let mut offset = 8usize;
178    let mut batch = MaterializedBatch::new(0, 0);
179    for _ in 0..num_cols {
180        if payload.len() < offset + 8 {
181            return Err(fail("truncated section column".to_string()));
182        }
183        let col_len = u64::from_le_bytes(payload[offset..offset + 8].try_into().unwrap()) as usize;
184        offset += 8;
185        if payload.len() < offset + col_len {
186            return Err(fail("truncated section column data".to_string()));
187        }
188        let values: Vec<Value> = postcard::from_bytes(&payload[offset..offset + col_len])
189            .map_err(|e| QueryError::execution(format!("spill run: deserialize column: {}", e)))?;
190        offset += col_len;
191        if values.len() != num_rows {
192            return Err(fail(format!(
193                "section column length {} != section row count {}",
194                values.len(),
195                num_rows
196            )));
197        }
198        batch.append_column(values);
199    }
200    Ok(batch)
201}
202
203// ── Metadata for a spill file with enhanced run format ───────────────────────
204
205/// Metadata for a single spill run file (sorted spill).
206#[derive(Debug, Clone)]
207pub struct SpilledRun {
208    pub path: PathBuf,
209    pub row_count: u64,
210    pub byte_size: u64,
211    pub schema_fingerprint: u64,
212}
213
214/// Schema fingerprint for a run file: FNV-style hash over column names.
215///
216/// Single implementation shared by every run writer (sort runs, set-operator
217/// spill); readers validate it against the recorded fingerprint on open.
218pub fn schema_fingerprint(col_names: &[String]) -> u64 {
219    use std::hash::Hasher;
220    let mut hasher = std::collections::hash_map::DefaultHasher::new();
221    for name in col_names {
222        hasher.write(name.as_bytes());
223        hasher.write_u8(0);
224    }
225    hasher.finish()
226}