linkrs_query/executor/streaming/spill/
run_format.rs1use std::path::PathBuf;
2
3use linkrs_core::columnar::MaterializedBatch;
4use linkrs_core::error::QueryError;
5use linkrs_core::Value;
6
7pub(super) const RUN_MAGIC: [u8; 4] = [0x4C, 0x52, 0x53, 0x43];
11
12pub(super) const RUN_VERSION: u32 = 3;
14
15pub(super) const RUN_HEADER_SIZE: u32 = 48;
17
18pub(super) const SECTION_ROWS: usize = 1024;
21
22pub(super) const COMPRESSION_MIN_SIZE: usize = 256;
24
25pub(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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46#[repr(u8)]
47pub enum RunCompression {
48 None = 0,
49 Zstd = 1,
50}
51
52#[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 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 let checksum = fnv1a_64(&buf[0..40]);
97 buf[40..48].copy_from_slice(&checksum.to_le_bytes());
98 buf
99 }
100
101 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
139pub(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
155pub(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#[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
214pub 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}