1use crate::error::{MongrelError, Result};
16use crate::memtable::Value;
17use crate::page::Encoding;
18use crate::schema::TypeId;
19use bincode::Options as _;
20use serde::{Deserialize, Serialize};
21use std::io::Read;
22
23const MAX_PAGE_DECOMPRESSED_BYTES: usize = 64 * 1024 * 1024;
24const GENERATED_EMBEDDING_TRAILER_MAGIC: &[u8; 4] = b"GEM1";
25
26pub fn encode_column(ty: TypeId, values: &[Value]) -> Result<Vec<u8>> {
28 let n = values.len();
29 let validity = validity_bitmap(values);
30 let payload = match ty {
31 TypeId::Int64 => fixed_encode(values, 8, |v| match v {
32 Value::Int64(x) => Ok(x.to_be_bytes().to_vec()),
33 Value::Null => Ok(vec![0; 8]),
34 _ => Err(type_mismatch(ty.clone(), v)),
35 })?,
36 TypeId::Float64 => fixed_encode(values, 8, |v| match v {
37 Value::Float64(f) => Ok(f.to_bits().to_be_bytes().to_vec()),
38 Value::Null => Ok(vec![0; 8]),
39 _ => Err(type_mismatch(ty.clone(), v)),
40 })?,
41 TypeId::TimestampNanos | TypeId::Date64 | TypeId::Time64 => {
42 fixed_encode(values, 8, |v| match v {
43 Value::Int64(x) => Ok(x.to_be_bytes().to_vec()),
44 Value::Null => Ok(vec![0; 8]),
45 _ => Err(type_mismatch(ty.clone(), v)),
46 })?
47 }
48 TypeId::Interval => fixed_encode(values, 20, |v| match v {
49 Value::Interval {
50 months,
51 days,
52 nanos,
53 } => {
54 let mut out = Vec::with_capacity(20);
55 out.extend_from_slice(&months.to_be_bytes());
56 out.extend_from_slice(&days.to_be_bytes());
57 out.extend_from_slice(&nanos.to_be_bytes());
58 Ok(out)
59 }
60 Value::Null => Ok(vec![0; 20]),
61 _ => Err(type_mismatch(ty.clone(), v)),
62 })?,
63 TypeId::Bool => fixed_encode(values, 1, |v| match v {
64 Value::Bool(b) => Ok(vec![*b as u8]),
65 Value::Null => Ok(vec![0]),
66 _ => Err(type_mismatch(ty.clone(), v)),
67 })?,
68 TypeId::Int32 | TypeId::UInt32 | TypeId::Date32 => fixed_encode(values, 4, |v| match v {
69 Value::Int64(x) => Ok((*x as i32).to_be_bytes().to_vec()),
70 Value::Null => Ok(vec![0; 4]),
71 _ => Err(type_mismatch(ty.clone(), v)),
72 })?,
73 TypeId::Bytes | TypeId::Json | TypeId::Enum { .. } => bytes_encode(values)?,
74 TypeId::Uuid => fixed_encode(values, 16, |v| match v {
75 Value::Uuid(b) => Ok(b.to_vec()),
76 Value::Null => Ok(vec![0; 16]),
77 _ => Err(type_mismatch(ty.clone(), v)),
78 })?,
79 TypeId::Embedding { dim } => embedding_encode(values, dim)?,
80 TypeId::Decimal128 { .. } => fixed_encode(values, 16, |v| match v {
81 Value::Decimal(d) => Ok(d.to_be_bytes().to_vec()),
82 Value::Null => Ok(vec![0; 16]),
83 _ => Err(type_mismatch(ty.clone(), v)),
84 })?,
85 other => {
86 return Err(MongrelError::Schema(format!(
87 "encoding for type {other:?} not implemented yet"
88 )))
89 }
90 };
91 let mut page = Vec::with_capacity(4 + validity.len() + payload.len());
92 page.extend_from_slice(&(validity.len() as u32).to_be_bytes());
93 page.extend_from_slice(&validity);
94 page.extend_from_slice(&payload);
95 let _ = n;
96 Ok(page)
97}
98
99pub fn decode_column(ty: TypeId, page: &[u8], n: usize, le: bool) -> Result<Vec<Value>> {
104 if page.len() < 4 {
105 return Err(MongrelError::InvalidArgument(
106 "page too short for header".into(),
107 ));
108 }
109 let vlen = u32::from_be_bytes([page[0], page[1], page[2], page[3]]) as usize;
110 if 4 + vlen > page.len() {
111 return Err(MongrelError::InvalidArgument(
112 "page validity out of range".into(),
113 ));
114 }
115 let validity = &page[4..4 + vlen];
116 let payload = &page[4 + vlen..];
117
118 let mut out = Vec::with_capacity(n);
119 let mut cur = 0usize;
120 let i64_at = |b: &[u8]| -> i64 {
121 if le {
122 i64::from_le_bytes(b.try_into().unwrap())
123 } else {
124 i64::from_be_bytes(b.try_into().unwrap())
125 }
126 };
127 let u64_at = |b: &[u8]| -> u64 {
128 if le {
129 u64::from_le_bytes(b.try_into().unwrap())
130 } else {
131 u64::from_be_bytes(b.try_into().unwrap())
132 }
133 };
134 for i in 0..n {
135 let non_null = (validity.get(i / 8).copied().unwrap_or(0) >> (i % 8)) & 1 == 1;
136 if !non_null {
137 out.push(Value::Null);
138 advance_null(&ty, payload, &mut cur)?; continue;
140 }
141 let val = match ty {
142 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date64 | TypeId::Time64 => {
143 let b = take(payload, &mut cur, 8)?;
144 Value::Int64(i64_at(&b))
145 }
146 TypeId::Float64 => {
147 let b = take(payload, &mut cur, 8)?;
148 Value::Float64(f64::from_bits(u64_at(&b)))
149 }
150 TypeId::Bool => {
151 let b = take(payload, &mut cur, 1)?;
152 Value::Bool(b[0] != 0)
153 }
154 TypeId::Int32 | TypeId::UInt32 | TypeId::Date32 => {
155 let b = take(payload, &mut cur, 4)?;
156 let v = if le {
157 i32::from_le_bytes(b.try_into().unwrap())
158 } else {
159 i32::from_be_bytes(b.try_into().unwrap())
160 };
161 Value::Int64(v as i64)
162 }
163 TypeId::Bytes | TypeId::Json | TypeId::Enum { .. } => {
164 let bytes_start = (n + 1) * 8;
165 let lo = read_off(payload, i, le);
166 let hi = read_off(payload, i + 1, le);
167 let data = payload[bytes_start + lo..bytes_start + hi].to_vec();
168 if matches!(ty, TypeId::Json) {
169 Value::Json(data)
170 } else {
171 Value::Bytes(data)
172 }
173 }
174 TypeId::Uuid => {
175 let b = take(payload, &mut cur, 16)?;
176 let mut arr = [0u8; 16];
177 arr.copy_from_slice(&b);
178 Value::Uuid(arr)
179 }
180 TypeId::Embedding { dim } => {
181 let mut acc = Vec::with_capacity(dim as usize);
182 for _ in 0..dim {
183 let b = take(payload, &mut cur, 4)?;
184 acc.push(f32::from_bits(u32::from_be_bytes(b.try_into().unwrap())));
185 }
186 Value::Embedding(acc)
187 }
188 TypeId::Decimal128 { .. } => {
189 let b = take(payload, &mut cur, 16)?;
190 let arr: [u8; 16] = b.try_into().unwrap();
191 Value::Decimal(i128::from_be_bytes(arr))
192 }
193 TypeId::Interval => {
194 let b = take(payload, &mut cur, 20)?;
195 let months = i64::from_be_bytes(b[0..8].try_into().unwrap());
196 let days = i32::from_be_bytes(b[8..12].try_into().unwrap());
197 let nanos = i64::from_be_bytes(b[12..20].try_into().unwrap());
198 Value::Interval {
199 months,
200 days,
201 nanos,
202 }
203 }
204 other => {
205 return Err(MongrelError::Schema(format!(
206 "decoding for type {other:?} not implemented yet"
207 )))
208 }
209 };
210 out.push(val);
211 }
212 if matches!(ty, TypeId::Embedding { .. }) && cur < payload.len() {
213 decode_generated_embedding_trailer(&payload[cur..], &mut out)?;
214 }
215 Ok(out)
216}
217
218fn validity_bitmap(values: &[Value]) -> Vec<u8> {
221 let n = values.len();
222 let mut bits = vec![0u8; n.div_ceil(8)];
223 for (i, v) in values.iter().enumerate() {
224 if !matches!(v, Value::Null) {
225 bits[i / 8] |= 1 << (i % 8);
226 }
227 }
228 bits
229}
230
231fn fixed_encode(
232 values: &[Value],
233 _width: usize,
234 mut enc: impl FnMut(&Value) -> Result<Vec<u8>>,
235) -> Result<Vec<u8>> {
236 let mut out = Vec::new();
237 for v in values {
238 out.extend_from_slice(&enc(v)?);
239 }
240 Ok(out)
241}
242
243fn embedding_encode(values: &[Value], dim: u32) -> Result<Vec<u8>> {
244 let mut out = Vec::with_capacity(values.len() * dim as usize * 4);
245 let mut metadata = Vec::with_capacity(values.len());
246 for v in values {
247 match v.as_embedding() {
248 Some(vector) => {
249 if vector.len() != dim as usize {
250 return Err(MongrelError::Schema(format!(
251 "embedding dimension mismatch: expected {dim}, got {}",
252 vector.len()
253 )));
254 }
255 for x in vector {
256 out.extend_from_slice(&x.to_bits().to_be_bytes());
257 }
258 }
259 None if matches!(v, Value::Null) => {
260 for _ in 0..dim {
261 out.extend_from_slice(&0u32.to_be_bytes());
262 }
263 }
264 None => return Err(type_mismatch(TypeId::Embedding { dim }, v)),
265 }
266 metadata.push(v.generated_embedding_metadata().cloned());
267 }
268 if metadata.iter().any(Option::is_some) {
269 let encoded = bincode::DefaultOptions::new()
270 .with_fixint_encoding()
271 .serialize(&metadata)
272 .map_err(|error| MongrelError::Other(format!("embedding metadata encode: {error}")))?;
273 out.extend_from_slice(GENERATED_EMBEDDING_TRAILER_MAGIC);
274 out.extend_from_slice(
275 &u32::try_from(encoded.len())
276 .map_err(|_| MongrelError::Full("embedding metadata page is too large".into()))?
277 .to_be_bytes(),
278 );
279 out.extend_from_slice(&encoded);
280 }
281 Ok(out)
282}
283
284fn decode_generated_embedding_trailer(trailer: &[u8], values: &mut [Value]) -> Result<()> {
285 if trailer.len() < 8 || &trailer[..4] != GENERATED_EMBEDDING_TRAILER_MAGIC {
286 return Err(MongrelError::InvalidArgument(
287 "invalid generated embedding metadata trailer".into(),
288 ));
289 }
290 let encoded_len = u32::from_be_bytes(trailer[4..8].try_into().unwrap()) as usize;
291 if trailer.len() != 8 + encoded_len {
292 return Err(MongrelError::InvalidArgument(
293 "generated embedding metadata trailer length mismatch".into(),
294 ));
295 }
296 let metadata: Vec<Option<crate::embedding::GeneratedEmbeddingMetadata>> =
297 bincode::DefaultOptions::new()
298 .with_fixint_encoding()
299 .with_limit(encoded_len as u64)
300 .reject_trailing_bytes()
301 .deserialize(&trailer[8..])
302 .map_err(|error| MongrelError::Other(format!("embedding metadata decode: {error}")))?;
303 if metadata.len() != values.len() {
304 return Err(MongrelError::InvalidArgument(
305 "generated embedding metadata row count mismatch".into(),
306 ));
307 }
308 for (value, metadata) in values.iter_mut().zip(metadata) {
309 let Some(metadata) = metadata else {
310 continue;
311 };
312 let Value::Embedding(vector) = value else {
313 return Err(MongrelError::InvalidArgument(
314 "generated embedding metadata targets a null row".into(),
315 ));
316 };
317 *value = Value::GeneratedEmbedding(Box::new(crate::embedding::GeneratedEmbeddingValue {
318 vector: std::mem::take(vector),
319 metadata,
320 }));
321 }
322 Ok(())
323}
324
325fn bytes_encode(values: &[Value]) -> Result<Vec<u8>> {
326 let n = values.len();
327 let mut offsets = Vec::with_capacity((n + 1) * 8);
328 let mut data = Vec::new();
329 let mut off = 0u64;
330 offsets.extend_from_slice(&off.to_be_bytes()); for v in values {
332 if let Value::Bytes(b) | Value::Json(b) = v {
333 data.extend_from_slice(b);
334 off = off
335 .checked_add(b.len() as u64)
336 .ok_or_else(|| MongrelError::InvalidArgument("bytes length overflow".into()))?;
337 }
338 offsets.extend_from_slice(&off.to_be_bytes());
339 }
340 let mut out = offsets;
341 out.extend_from_slice(&data);
342 Ok(out)
343}
344
345fn read_off(payload: &[u8], idx: usize, le: bool) -> usize {
346 let s = idx * 8;
347 if le {
348 u64::from_le_bytes(payload[s..s + 8].try_into().unwrap()) as usize
349 } else {
350 u64::from_be_bytes(payload[s..s + 8].try_into().unwrap()) as usize
351 }
352}
353
354fn take(payload: &[u8], cur: &mut usize, n: usize) -> Result<Vec<u8>> {
355 if *cur + n > payload.len() {
356 return Err(MongrelError::InvalidArgument("payload truncated".into()));
357 }
358 let s = &payload[*cur..*cur + n];
359 *cur += n;
360 Ok(s.to_vec())
361}
362
363fn advance_null(ty: &TypeId, payload: &[u8], cur: &mut usize) -> Result<()> {
365 let w = match ty {
366 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos => 8,
367 TypeId::Int32 | TypeId::UInt32 | TypeId::Date32 => 4,
368 TypeId::Bool => 1,
369 TypeId::Decimal128 { .. } => 16,
370 TypeId::Date64 | TypeId::Time64 => 8,
371 TypeId::Interval => 20,
372 TypeId::Uuid => 16,
373 TypeId::Embedding { dim } => *dim as usize * std::mem::size_of::<f32>(),
374 TypeId::Json | TypeId::Enum { .. } => return Ok(()), TypeId::Bytes => return Ok(()),
378 _ => return Ok(()),
379 };
380 if *cur + w > payload.len() {
381 return Err(MongrelError::InvalidArgument(
382 "payload truncated at null".into(),
383 ));
384 }
385 *cur += w;
386 Ok(())
387}
388
389fn type_mismatch(ty: TypeId, v: &Value) -> MongrelError {
390 MongrelError::Schema(format!("type mismatch: column {ty:?}, value {v:?}"))
391}
392
393const ALGO_PLAIN: u8 = 0; const ALGO_ZSTD_PLAIN: u8 = 1; const ALGO_ZSTD_DICT: u8 = 2; const ALGO_ZSTD_DELTA: u8 = 3; const ALGO_LZ4_PLAIN: u8 = 4; const ALGO_LZ4_DICT: u8 = 5; const ALGO_LZ4_DELTA: u8 = 6; const ALGO_LE_FLAG: u8 = 1 << 3;
414
415#[inline]
416fn algo_with_le(algo: u8, le: bool) -> u8 {
417 if le {
418 algo | ALGO_LE_FLAG
419 } else {
420 algo
421 }
422}
423
424#[derive(Debug, Clone, Copy)]
428pub enum Compress {
429 Plain,
431 Zstd(i32),
433 Lz4,
436}
437
438fn zstd_compress(data: &[u8]) -> Result<Vec<u8>> {
439 zstd_compress_level(data, 3)
440}
441
442fn zstd_compress_level(data: &[u8], level: i32) -> Result<Vec<u8>> {
447 if data.len() > MAX_PAGE_DECOMPRESSED_BYTES {
448 return Err(MongrelError::ResourceLimitExceeded {
449 resource: "column page bytes",
450 requested: data.len(),
451 limit: MAX_PAGE_DECOMPRESSED_BYTES,
452 });
453 }
454 zstd::encode_all(data, level)
455 .map_err(|e| MongrelError::InvalidArgument(format!("zstd compress: {e}")))
456}
457
458fn lz4_compress(data: &[u8]) -> Vec<u8> {
464 lz4_flex::block::compress_prepend_size(data)
465}
466
467fn max_decompressed_bytes(ty: TypeId, n: usize, algo: u8) -> usize {
472 let validity = 4 + n.div_ceil(8);
474 if matches!(algo, ALGO_ZSTD_DELTA | ALGO_LZ4_DELTA) {
475 return validity
476 .saturating_add(n.saturating_mul(8))
477 .min(MAX_PAGE_DECOMPRESSED_BYTES);
478 }
479 if matches!(algo, ALGO_ZSTD_DICT | ALGO_LZ4_DICT) {
480 return MAX_PAGE_DECOMPRESSED_BYTES;
481 }
482 let payload = match ty {
483 TypeId::Bytes | TypeId::Json | TypeId::Enum { .. } => {
486 return MAX_PAGE_DECOMPRESSED_BYTES;
487 }
488 TypeId::Embedding { dim } => (dim as usize).saturating_mul(8).saturating_mul(n),
489 _ => n.saturating_mul(ty.fixed_size().unwrap_or(8)),
490 };
491 validity
492 .saturating_add(payload)
493 .min(MAX_PAGE_DECOMPRESSED_BYTES)
494}
495
496fn lz4_decompress(data: &[u8], max_bytes: usize) -> Result<Vec<u8>> {
497 if data.len() < 4 {
502 return Err(MongrelError::InvalidArgument("lz4 page truncated".into()));
503 }
504 let declared = u32::from_le_bytes([data[0], data[1], data[2], data[3]]) as usize;
505 if declared > max_bytes {
506 return Err(MongrelError::InvalidArgument(format!(
507 "lz4 declared size {declared} exceeds page limit {max_bytes}"
508 )));
509 }
510 lz4_flex::block::decompress(&data[4..], declared)
511 .map_err(|e| MongrelError::InvalidArgument(format!("lz4 decompress: {e}")))
512}
513
514fn zstd_decompress(data: &[u8], max_bytes: usize) -> Result<Vec<u8>> {
515 let decoder = zstd::stream::read::Decoder::new(data)
518 .map_err(|e| MongrelError::InvalidArgument(format!("zstd decompress: {e}")))?;
519 let mut out = Vec::with_capacity(max_bytes.min(1024 * 1024));
520 decoder
521 .take(max_bytes.saturating_add(1) as u64)
522 .read_to_end(&mut out)
523 .map_err(|e| MongrelError::InvalidArgument(format!("zstd decompress: {e}")))?;
524 if out.len() > max_bytes {
525 return Err(MongrelError::InvalidArgument(format!(
526 "zstd output {} exceeds page limit {max_bytes}",
527 out.len()
528 )));
529 }
530 Ok(out)
531}
532
533pub fn encode_page(ty: TypeId, values: &[Value], encoding: Encoding) -> Result<Vec<u8>> {
536 Ok(match encoding {
537 Encoding::Plain => {
538 let mut out = vec![ALGO_PLAIN];
539 out.extend(encode_column(ty, values)?);
540 out
541 }
542 Encoding::Dictionary if matches!(ty, TypeId::Bytes | TypeId::Enum { .. }) => {
543 let dict = dict_encode_bytes(values);
544 let mut out = vec![ALGO_ZSTD_DICT];
545 out.extend(zstd_compress(&dict)?);
546 out
547 }
548 _ => {
549 let mut out = vec![ALGO_ZSTD_PLAIN];
550 let enc = encode_column(ty, values)?;
551 out.extend(zstd_compress(&enc)?);
552 out
553 }
554 })
555}
556
557pub fn decode_page(ty: TypeId, page: &[u8], n: usize) -> Result<Vec<Value>> {
559 if page.is_empty() {
560 return Err(MongrelError::InvalidArgument("empty page".into()));
561 }
562 let algo = page[0];
563 let le = algo & ALGO_LE_FLAG != 0;
568 let base = algo & !ALGO_LE_FLAG;
569 let body = &page[1..];
570 let raw_owned;
572 let raw: &[u8] = match base {
573 ALGO_PLAIN => body,
574 ALGO_ZSTD_PLAIN | ALGO_ZSTD_DICT | ALGO_ZSTD_DELTA => {
575 raw_owned = zstd_decompress(body, max_decompressed_bytes(ty.clone(), n, base))?;
576 &raw_owned
577 }
578 ALGO_LZ4_PLAIN | ALGO_LZ4_DICT | ALGO_LZ4_DELTA => {
579 raw_owned = lz4_decompress(body, max_decompressed_bytes(ty.clone(), n, base))?;
580 &raw_owned
581 }
582 other => {
583 return Err(MongrelError::InvalidArgument(format!(
584 "unknown page algo {other}"
585 )))
586 }
587 };
588 match base {
589 ALGO_PLAIN | ALGO_ZSTD_PLAIN | ALGO_LZ4_PLAIN => decode_column(ty, raw, n, le),
590 ALGO_ZSTD_DICT | ALGO_LZ4_DICT => dict_decode_bytes(raw, n),
591 ALGO_ZSTD_DELTA | ALGO_LZ4_DELTA => decode_int64_delta_values(raw, ty, n, le),
592 _ => unreachable!(),
593 }
594}
595
596fn decode_int64_delta_values(raw: &[u8], ty: TypeId, n: usize, le: bool) -> Result<Vec<Value>> {
599 if !matches!(ty, TypeId::Int64 | TypeId::TimestampNanos) {
600 return Err(MongrelError::InvalidArgument(format!(
601 "delta page not valid for {ty:?}"
602 )));
603 }
604 let (validity, p) = split_validity(raw)?;
605 let deltas = if le {
606 take_i64_le(p, n)?
607 } else {
608 take_i64_be(p, n)?
609 };
610 let data = delta_prefix_sum_i64(&deltas);
611 let mut out = Vec::with_capacity(n);
612 for (i, &v) in data.iter().enumerate() {
613 let non_null = (validity.get(i / 8).copied().unwrap_or(0) >> (i % 8)) & 1 == 1;
614 out.push(if non_null {
615 Value::Int64(v)
616 } else {
617 Value::Null
618 });
619 }
620 Ok(out)
621}
622
623fn dict_encode_bytes(values: &[Value]) -> Vec<u8> {
627 let validity = validity_bitmap(values);
628 let mut table: Vec<Vec<u8>> = Vec::new();
629 let mut index_of: std::collections::HashMap<&[u8], u32> = std::collections::HashMap::new();
630 let mut indices: Vec<u32> = Vec::with_capacity(values.len());
631 for v in values {
632 let idx = match v {
633 Value::Bytes(b) => {
634 if let Some(&i) = index_of.get(b.as_slice()) {
635 i
636 } else {
637 let i = table.len() as u32;
638 index_of.insert(b.as_slice(), i);
639 table.push(b.clone());
640 i
641 }
642 }
643 _ => 0,
644 };
645 indices.push(idx);
646 }
647
648 let mut out = Vec::new();
649 out.extend_from_slice(&(validity.len() as u32).to_be_bytes());
650 out.extend_from_slice(&validity);
651 out.extend_from_slice(&(indices.len() as u32).to_be_bytes());
652 for i in &indices {
653 out.extend_from_slice(&i.to_be_bytes());
654 }
655 out.extend_from_slice(&(table.len() as u32).to_be_bytes());
656 for entry in &table {
657 out.extend_from_slice(&(entry.len() as u32).to_be_bytes());
658 out.extend_from_slice(entry);
659 }
660 out
661}
662
663fn dict_decode_bytes(data: &[u8], n: usize) -> Result<Vec<Value>> {
664 let mut cur = 0usize;
665 let vlen = read_u32_be(data, &mut cur)? as usize;
666 let validity = checked_slice(data, &mut cur, vlen)?;
667 let index_count = read_u32_be(data, &mut cur)? as usize;
668 let mut indices = Vec::with_capacity(index_count.min(n));
669 for _ in 0..index_count {
670 indices.push(read_u32_be(data, &mut cur)?);
671 }
672 let table_count = read_u32_be(data, &mut cur)? as usize;
673 let mut table: Vec<Vec<u8>> = Vec::with_capacity(table_count);
674 for _ in 0..table_count {
675 let len = read_u32_be(data, &mut cur)? as usize;
676 table.push(checked_slice(data, &mut cur, len)?.to_vec());
677 }
678
679 let mut out = Vec::with_capacity(n);
680 for (i, &idx) in indices.iter().enumerate().take(n) {
681 let non_null = (validity.get(i / 8).copied().unwrap_or(0) >> (i % 8)) & 1 == 1;
682 if !non_null {
683 out.push(Value::Null);
684 } else {
685 let entry = table
686 .get(idx as usize)
687 .cloned()
688 .ok_or_else(|| MongrelError::InvalidArgument("dict index out of range".into()))?;
689 out.push(Value::Bytes(entry));
690 }
691 }
692 Ok(out)
693}
694
695fn read_u32_be(data: &[u8], cur: &mut usize) -> Result<u32> {
696 if *cur + 4 > data.len() {
697 return Err(MongrelError::InvalidArgument(
698 "dict payload truncated".into(),
699 ));
700 }
701 let v = u32::from_be_bytes([data[*cur], data[*cur + 1], data[*cur + 2], data[*cur + 3]]);
702 *cur += 4;
703 Ok(v)
704}
705
706fn checked_slice<'a>(data: &'a [u8], cur: &mut usize, len: usize) -> Result<&'a [u8]> {
709 if *cur + len > data.len() {
710 return Err(MongrelError::InvalidArgument(
711 "dict payload truncated".into(),
712 ));
713 }
714 let s = &data[*cur..*cur + len];
715 *cur += len;
716 Ok(s)
717}
718
719#[cfg(test)]
720mod compressed_tests {
721 use super::*;
722
723 #[test]
724 fn zstd_plain_round_trip_int64() {
725 let vals: Vec<Value> = (0..1000).map(Value::Int64).collect();
726 let page = encode_page(TypeId::Int64, &vals, Encoding::Zstd).unwrap();
727 assert!(
728 page.len() < vals.len() * 8,
729 "zstd must shrink sequential ints"
730 );
731 let back = decode_page(TypeId::Int64, &page, vals.len()).unwrap();
732 assert_eq!(back, vals);
733 }
734
735 #[test]
736 fn dictionary_round_trip_low_card_bytes() {
737 let palette: &[&[u8]] = &[b"red", b"green", b"blue", b"red"];
738 let vals: Vec<Value> = (0..500)
739 .map(|i| Value::Bytes(palette[i % palette.len()].to_vec()))
740 .collect();
741 let page = encode_page(TypeId::Bytes, &vals, Encoding::Dictionary).unwrap();
742 assert!(
743 page.len() < 100,
744 "4 distinct strings over 500 rows must compress to a tiny page, got {}",
745 page.len()
746 );
747 let back = decode_page(TypeId::Bytes, &page, vals.len()).unwrap();
748 assert_eq!(back, vals);
749 }
750
751 #[test]
752 fn variable_page_larger_than_16_kib_round_trips() {
753 for (ty, value) in [
754 (TypeId::Bytes, Value::Bytes(vec![b'x'; 16_385])),
755 (TypeId::Json, Value::Json(vec![b' '; 16_385])),
756 ] {
757 let values = vec![value];
758 let page = encode_page(ty.clone(), &values, Encoding::Zstd).unwrap();
759 assert_eq!(decode_page(ty, &page, 1).unwrap(), values);
760 }
761 }
762
763 #[test]
764 fn plain_page_still_round_trips() {
765 let vals = vec![Value::Int64(1), Value::Null, Value::Int64(9)];
766 let page = encode_page(TypeId::Int64, &vals, Encoding::Plain).unwrap();
767 assert_eq!(page[0], ALGO_PLAIN);
768 assert_eq!(decode_page(TypeId::Int64, &page, 3).unwrap(), vals);
769 }
770}
771
772#[cfg(test)]
773mod tests {
774 use super::*;
775
776 #[test]
777 fn round_trips_int64_with_nulls() {
778 let vals = vec![
779 Value::Int64(1),
780 Value::Null,
781 Value::Int64(-5),
782 Value::Int64(1 << 40),
783 ];
784 let page = encode_column(TypeId::Int64, &vals).unwrap();
785 let back = decode_column(TypeId::Int64, &page, vals.len(), false).unwrap();
786 assert_eq!(back, vals);
787 }
788
789 #[test]
790 fn round_trips_bytes() {
791 let vals = vec![
792 Value::Bytes(b"hello".to_vec()),
793 Value::Null,
794 Value::Bytes(b"".to_vec()),
795 Value::Bytes(b"wide \x00 byte".to_vec()),
796 ];
797 let page = encode_column(TypeId::Bytes, &vals).unwrap();
798 let back = decode_column(TypeId::Bytes, &page, vals.len(), false).unwrap();
799 assert_eq!(back, vals);
800 }
801
802 #[test]
803 fn round_trips_embedding() {
804 let vals = vec![
805 Value::Embedding(vec![1.0, -2.5, 3.0]),
806 Value::Null,
807 Value::Embedding(vec![4.0, 5.0, 6.0]),
808 ];
809 let page = encode_column(TypeId::Embedding { dim: 3 }, &vals).unwrap();
810 let back = decode_column(TypeId::Embedding { dim: 3 }, &page, vals.len(), false).unwrap();
811 assert_eq!(back, vals);
812 }
813
814 #[test]
815 fn round_trips_bool() {
816 let vals = vec![
817 Value::Bool(true),
818 Value::Bool(false),
819 Value::Null,
820 Value::Bool(true),
821 ];
822 let page = encode_column(TypeId::Bool, &vals).unwrap();
823 assert_eq!(
824 decode_column(TypeId::Bool, &page, vals.len(), false).unwrap(),
825 vals
826 );
827 }
828}
829
830#[derive(Debug, Clone, Serialize, Deserialize)]
839pub enum NativeColumn {
840 Int64 {
841 data: Vec<i64>,
842 validity: Vec<u8>,
843 },
844 Float64 {
845 data: Vec<f64>,
846 validity: Vec<u8>,
847 },
848 Bool {
850 data: Vec<u8>,
851 validity: Vec<u8>,
852 },
853 Bytes {
855 offsets: Vec<u32>,
856 values: Vec<u8>,
857 validity: Vec<u8>,
858 },
859}
860
861impl NativeColumn {
862 pub fn len(&self) -> usize {
863 match self {
864 NativeColumn::Int64 { data, .. } => data.len(),
865 NativeColumn::Float64 { data, .. } => data.len(),
866 NativeColumn::Bool { data, .. } => data.len(),
867 NativeColumn::Bytes { offsets, .. } => offsets.len().saturating_sub(1),
868 }
869 }
870
871 pub fn is_empty(&self) -> bool {
872 self.len() == 0
873 }
874
875 pub fn validity(&self) -> &[u8] {
878 match self {
879 NativeColumn::Int64 { validity, .. }
880 | NativeColumn::Float64 { validity, .. }
881 | NativeColumn::Bool { validity, .. }
882 | NativeColumn::Bytes { validity, .. } => validity,
883 }
884 }
885
886 pub fn validate(&self) -> bool {
890 match self {
891 NativeColumn::Int64 { data, validity } => {
892 validity.len() == data.len().div_ceil(8) || validity.is_empty()
893 }
894 NativeColumn::Float64 { data, validity } => {
895 validity.len() == data.len().div_ceil(8) || validity.is_empty()
896 }
897 NativeColumn::Bool { data, validity } => {
898 validity.len() == data.len().div_ceil(8) || validity.is_empty()
899 }
900 NativeColumn::Bytes {
901 offsets,
902 values,
903 validity,
904 } => {
905 let n = offsets.len().saturating_sub(1);
906 (validity.len() == n.div_ceil(8) || validity.is_empty())
907 && offsets
908 .last()
909 .map(|&last| (last as usize) <= values.len())
910 .unwrap_or(true)
911 }
912 }
913 }
914
915 pub fn null_count(&self, n: usize) -> usize {
918 if n == 0 {
919 return 0;
920 }
921 let validity = match self {
922 NativeColumn::Int64 { validity, .. }
923 | NativeColumn::Float64 { validity, .. }
924 | NativeColumn::Bool { validity, .. }
925 | NativeColumn::Bytes { validity, .. } => validity,
926 };
927 if validity.is_empty() {
928 return 0;
929 }
930 let full = n / 8;
933 let mut set = validity[..full.min(validity.len())]
934 .iter()
935 .map(|b| b.count_ones() as usize)
936 .sum::<usize>();
937 let tail_bits = n % 8;
938 if tail_bits > 0 {
939 if let Some(&b) = validity.get(full) {
940 set += (b & ((1u8 << tail_bits) - 1)).count_ones() as usize;
941 }
942 }
943 n - set
944 }
945
946 pub fn approx_bytes(&self) -> u64 {
948 match self {
949 NativeColumn::Int64 { data, validity } => {
950 (data.len() as u64) * 8 + validity.len() as u64
951 }
952 NativeColumn::Float64 { data, validity } => {
953 (data.len() as u64) * 8 + validity.len() as u64
954 }
955 NativeColumn::Bool { data, validity } => data.len() as u64 + validity.len() as u64,
956 NativeColumn::Bytes {
957 offsets,
958 values,
959 validity,
960 } => values.len() as u64 + (offsets.len() as u64) * 4 + validity.len() as u64,
961 }
962 }
963
964 pub fn int64_sequence(start: i64, n: usize) -> Self {
966 NativeColumn::Int64 {
967 data: (0..n).map(|i| start + i as i64).collect(),
968 validity: full_validity(n),
969 }
970 }
971
972 pub fn int64_constant(value: i64, n: usize) -> Self {
974 NativeColumn::Int64 {
975 data: vec![value; n],
976 validity: full_validity(n),
977 }
978 }
979
980 pub fn bool_constant(value: bool, n: usize) -> Self {
982 NativeColumn::Bool {
983 data: vec![if value { 1 } else { 0 }; n],
984 validity: full_validity(n),
985 }
986 }
987
988 pub fn gather(&self, indices: &[usize]) -> NativeColumn {
991 let bit = |v: &[u8], i: usize| (v.get(i / 8).copied().unwrap_or(0) >> (i % 8)) & 1 == 1;
992 match self {
993 NativeColumn::Int64 { data, validity } => NativeColumn::Int64 {
994 data: indices.iter().map(|&i| data[i]).collect(),
995 validity: validity_bitmap_from(indices.iter().map(|&i| bit(validity, i))),
996 },
997 NativeColumn::Float64 {
998 data: fdata,
999 validity: fval,
1000 } => NativeColumn::Float64 {
1001 data: indices.iter().map(|&i| fdata[i]).collect(),
1002 validity: validity_bitmap_from(indices.iter().map(|&i| bit(fval, i))),
1003 },
1004 NativeColumn::Bool {
1005 data: bdata,
1006 validity: bval,
1007 } => NativeColumn::Bool {
1008 data: indices.iter().map(|&i| bdata[i]).collect(),
1009 validity: validity_bitmap_from(indices.iter().map(|&i| bit(bval, i))),
1010 },
1011 NativeColumn::Bytes {
1012 offsets,
1013 values,
1014 validity,
1015 } => {
1016 let mut out_offsets = Vec::with_capacity(indices.len() + 1);
1017 let mut out_values = Vec::new();
1018 out_offsets.push(0);
1019 for &i in indices {
1020 let lo = offsets[i] as usize;
1021 let hi = offsets[i + 1] as usize;
1022 out_values.extend_from_slice(&values[lo..hi]);
1023 out_offsets.push(out_values.len() as u32);
1024 }
1025 NativeColumn::Bytes {
1026 offsets: out_offsets,
1027 values: out_values,
1028 validity: validity_bitmap_from(indices.iter().map(|&i| bit(validity, i))),
1029 }
1030 }
1031 }
1032 }
1033
1034 pub fn value_at(&self, idx: usize) -> Option<Value> {
1039 match self {
1040 NativeColumn::Int64 { data, validity } => {
1041 if !validity_bit(validity, idx) {
1042 return None;
1043 }
1044 data.get(idx).copied().map(Value::Int64)
1045 }
1046 NativeColumn::Float64 { data, validity } => {
1047 if !validity_bit(validity, idx) {
1048 return None;
1049 }
1050 data.get(idx).copied().map(Value::Float64)
1051 }
1052 NativeColumn::Bool { data, validity } => {
1053 if !validity_bit(validity, idx) {
1054 return None;
1055 }
1056 data.get(idx).copied().map(|b| Value::Bool(b != 0))
1057 }
1058 NativeColumn::Bytes {
1059 offsets,
1060 values,
1061 validity,
1062 } => {
1063 if !validity_bit(validity, idx) {
1064 return None;
1065 }
1066 if idx + 1 >= offsets.len() {
1067 return None;
1068 }
1069 let lo = offsets[idx] as usize;
1070 let hi = offsets[idx + 1] as usize;
1071 Some(Value::Bytes(values[lo..hi].to_vec()))
1072 }
1073 }
1074 }
1075
1076 pub fn slice_range(&self, start: usize, end: usize) -> NativeColumn {
1079 let mk_validity = |v: &[u8]| -> Vec<bool> {
1080 (0..(end - start))
1081 .map(|i| validity_bit(v, start + i))
1082 .collect()
1083 };
1084 match self {
1085 NativeColumn::Int64 { data, validity } => NativeColumn::Int64 {
1086 data: data[start..end].to_vec(),
1087 validity: validity_bitmap_from(mk_validity(validity)),
1088 },
1089 NativeColumn::Float64 { data, validity } => NativeColumn::Float64 {
1090 data: data[start..end].to_vec(),
1091 validity: validity_bitmap_from(mk_validity(validity)),
1092 },
1093 NativeColumn::Bool { data, validity } => NativeColumn::Bool {
1094 data: data[start..end].to_vec(),
1095 validity: validity_bitmap_from(mk_validity(validity)),
1096 },
1097 NativeColumn::Bytes {
1098 offsets,
1099 values,
1100 validity,
1101 } => {
1102 let lo = offsets[start] as usize;
1103 let hi = offsets[end] as usize;
1104 let new_offsets: Vec<u32> = offsets[start..=end]
1105 .iter()
1106 .map(|o| *o - offsets[start])
1107 .collect();
1108 NativeColumn::Bytes {
1109 offsets: new_offsets,
1110 values: values[lo..hi].to_vec(),
1111 validity: validity_bitmap_from(mk_validity(validity)),
1112 }
1113 }
1114 }
1115 }
1116
1117 pub fn concat(parts: &[NativeColumn]) -> NativeColumn {
1120 match parts.first() {
1121 Some(NativeColumn::Int64 { .. }) => {
1122 let mut data = Vec::new();
1123 let mut non_null: Vec<bool> = Vec::new();
1124 for p in parts {
1125 if let NativeColumn::Int64 { data: d, validity } = p {
1126 data.extend_from_slice(d);
1127 non_null.extend((0..d.len()).map(|i| validity_bit(validity, i)));
1128 }
1129 }
1130 NativeColumn::Int64 {
1131 data,
1132 validity: validity_bitmap_from(non_null),
1133 }
1134 }
1135 Some(NativeColumn::Float64 { .. }) => {
1136 let mut data = Vec::new();
1137 let mut non_null: Vec<bool> = Vec::new();
1138 for p in parts {
1139 if let NativeColumn::Float64 { data: d, validity } = p {
1140 data.extend_from_slice(d);
1141 non_null.extend((0..d.len()).map(|i| validity_bit(validity, i)));
1142 }
1143 }
1144 NativeColumn::Float64 {
1145 data,
1146 validity: validity_bitmap_from(non_null),
1147 }
1148 }
1149 Some(NativeColumn::Bool { .. }) => {
1150 let mut data = Vec::new();
1151 let mut non_null: Vec<bool> = Vec::new();
1152 for p in parts {
1153 if let NativeColumn::Bool { data: d, validity } = p {
1154 data.extend_from_slice(d);
1155 non_null.extend((0..d.len()).map(|i| validity_bit(validity, i)));
1156 }
1157 }
1158 NativeColumn::Bool {
1159 data,
1160 validity: validity_bitmap_from(non_null),
1161 }
1162 }
1163 Some(NativeColumn::Bytes { .. }) => {
1164 let mut offsets: Vec<u32> = vec![0];
1165 let mut values = Vec::new();
1166 let mut non_null: Vec<bool> = Vec::new();
1167 for p in parts {
1168 if let NativeColumn::Bytes {
1169 offsets: off,
1170 values: val,
1171 validity,
1172 } = p
1173 {
1174 for w in off.windows(2) {
1175 values.extend_from_slice(&val[w[0] as usize..w[1] as usize]);
1176 offsets.push(values.len() as u32);
1177 }
1178 non_null.extend((0..off.len() - 1).map(|i| validity_bit(validity, i)));
1179 }
1180 }
1181 NativeColumn::Bytes {
1182 offsets,
1183 values,
1184 validity: validity_bitmap_from(non_null),
1185 }
1186 }
1187 None => NativeColumn::Bytes {
1188 offsets: vec![0],
1189 values: Vec::new(),
1190 validity: Vec::new(),
1191 },
1192 }
1193 }
1194}
1195
1196fn full_validity(n: usize) -> Vec<u8> {
1197 validity_bitmap_from(std::iter::repeat_n(true, n))
1198}
1199
1200pub fn null_native(ty: TypeId, n: usize) -> NativeColumn {
1203 let validity = vec![0u8; n.div_ceil(8)];
1204 match ty {
1205 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date64 | TypeId::Time64 => {
1206 NativeColumn::Int64 {
1207 data: vec![0; n],
1208 validity,
1209 }
1210 }
1211 TypeId::Float64 => NativeColumn::Float64 {
1212 data: vec![0.0; n],
1213 validity,
1214 },
1215 TypeId::Bool => NativeColumn::Bool {
1216 data: vec![0; n],
1217 validity,
1218 },
1219 _ => NativeColumn::Bytes {
1220 offsets: vec![0u32; n + 1],
1221 values: Vec::new(),
1222 validity,
1223 },
1224 }
1225}
1226
1227#[inline]
1229pub fn validity_bit(validity: &[u8], i: usize) -> bool {
1230 (validity.get(i / 8).copied().unwrap_or(0) >> (i % 8)) & 1 == 1
1231}
1232
1233pub fn all_non_null(validity: &[u8], n: usize) -> bool {
1236 if n == 0 {
1237 return true;
1238 }
1239 let full = n / 8;
1240 if !validity[..full].iter().all(|&b| b == 0xFF) {
1241 return false;
1242 }
1243 if !n.is_multiple_of(8) {
1244 let mask = (1u8 << (n % 8)) - 1;
1245 (validity.get(full).copied().unwrap_or(0) & mask) == mask
1246 } else {
1247 true
1248 }
1249}
1250
1251pub fn native_min_max(ty: TypeId, col: &NativeColumn) -> (Option<Vec<u8>>, Option<Vec<u8>>, u64) {
1256 let _ = ty;
1257 match col {
1258 NativeColumn::Int64 { data, validity } => {
1259 let (mut mn, mut mx, mut nulls) = (None::<i64>, None::<i64>, 0u64);
1260 for (i, v) in data.iter().enumerate() {
1261 if !validity_bit(validity, i) {
1262 nulls += 1;
1263 continue;
1264 }
1265 mn = Some(mn.map_or(*v, |m| m.min(*v)));
1266 mx = Some(mx.map_or(*v, |m| m.max(*v)));
1267 }
1268 (
1269 mn.map(|v| v.to_be_bytes().to_vec()),
1270 mx.map(|v| v.to_be_bytes().to_vec()),
1271 nulls,
1272 )
1273 }
1274 NativeColumn::Float64 { data, validity } => {
1275 let (mut mn, mut mx, mut nulls) = (None::<f64>, None::<f64>, 0u64);
1276 for (i, v) in data.iter().enumerate() {
1277 if !validity_bit(validity, i) || v.is_nan() {
1278 nulls += 1;
1279 continue;
1280 }
1281 mn = Some(mn.map_or(*v, |m| m.min(*v)));
1282 mx = Some(mx.map_or(*v, |m| m.max(*v)));
1283 }
1284 (
1285 mn.map(|v| v.to_bits().to_be_bytes().to_vec()),
1286 mx.map(|v| v.to_bits().to_be_bytes().to_vec()),
1287 nulls,
1288 )
1289 }
1290 NativeColumn::Bool { data, validity } => {
1291 let (mut any_t, mut any_f, mut nulls) = (false, false, 0u64);
1292 for (i, v) in data.iter().enumerate() {
1293 if !validity_bit(validity, i) {
1294 nulls += 1;
1295 continue;
1296 }
1297 if *v != 0 {
1298 any_t = true;
1299 } else {
1300 any_f = true;
1301 }
1302 }
1303 let min = if any_f || any_t {
1304 Some(vec![if any_f { 0 } else { 1 }])
1305 } else {
1306 None
1307 };
1308 let max = if any_t || any_f {
1309 Some(vec![if any_t { 1 } else { 0 }])
1310 } else {
1311 None
1312 };
1313 (min, max, nulls)
1314 }
1315 NativeColumn::Bytes {
1316 offsets,
1317 values,
1318 validity,
1319 } => {
1320 let mut mn: Option<&[u8]> = None;
1321 let mut mx: Option<&[u8]> = None;
1322 let mut nulls = 0u64;
1323 for i in 0..offsets.len().saturating_sub(1) {
1324 if !validity_bit(validity, i) {
1325 nulls += 1;
1326 continue;
1327 }
1328 let s = &values[offsets[i] as usize..offsets[i + 1] as usize];
1329 mn = Some(match mn {
1330 None => s,
1331 Some(m) if s < m => s,
1332 Some(m) => m,
1333 });
1334 mx = Some(match mx {
1335 None => s,
1336 Some(m) if s > m => s,
1337 Some(m) => m,
1338 });
1339 }
1340 (mn.map(|s| s.to_vec()), mx.map(|s| s.to_vec()), nulls)
1341 }
1342 }
1343}
1344
1345pub fn page_stat_for(
1349 ty: TypeId,
1350 col: &NativeColumn,
1351 first_row_id: u64,
1352 last_row_id: u64,
1353) -> crate::page::PageStat {
1354 let (min, max, null_count) = native_min_max(ty, col);
1355 crate::page::PageStat {
1356 first_row_id,
1357 last_row_id,
1358 null_count,
1359 row_count: col.len() as u32,
1360 min,
1361 max,
1362 offset: 0,
1363 compressed_len: 0,
1364 uncompressed_len: 0,
1365 }
1366}
1367
1368pub fn encode_key_native(_ty: TypeId, col: &NativeColumn, i: usize) -> Option<Vec<u8>> {
1375 match col {
1376 NativeColumn::Int64 { data, validity } if validity_bit(validity, i) => {
1377 Some(data[i].to_be_bytes().to_vec())
1378 }
1379 NativeColumn::Float64 { data, validity } if validity_bit(validity, i) => {
1380 Some(data[i].to_bits().to_be_bytes().to_vec())
1381 }
1382 NativeColumn::Bool { data, validity } if validity_bit(validity, i) => Some(vec![data[i]]),
1383 NativeColumn::Bytes {
1384 offsets,
1385 values,
1386 validity,
1387 } if validity_bit(validity, i) => {
1388 let lo = offsets[i] as usize;
1389 let hi = offsets[i + 1] as usize;
1390 Some(values[lo..hi].to_vec())
1391 }
1392 _ => None,
1393 }
1394}
1395
1396pub fn native_bytes_at(col: &NativeColumn, i: usize) -> Option<&[u8]> {
1400 match col {
1401 NativeColumn::Bytes {
1402 offsets,
1403 values,
1404 validity,
1405 } if validity_bit(validity, i) => {
1406 let lo = offsets[i] as usize;
1407 let hi = offsets[i + 1] as usize;
1408 Some(&values[lo..hi])
1409 }
1410 _ => None,
1411 }
1412}
1413
1414pub fn values_to_native(ty: TypeId, values: &[Value]) -> NativeColumn {
1417 let n = values.len();
1418 let mut non_null = vec![false; n];
1419 match ty {
1420 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date64 | TypeId::Time64 => {
1421 let mut data = Vec::with_capacity(n);
1422 for (i, v) in values.iter().enumerate() {
1423 match v {
1424 Value::Int64(x) => {
1425 non_null[i] = true;
1426 data.push(*x);
1427 }
1428 _ => data.push(0),
1429 }
1430 }
1431 NativeColumn::Int64 {
1432 data,
1433 validity: validity_bitmap_from(non_null),
1434 }
1435 }
1436 TypeId::Float64 => {
1437 let mut data = Vec::with_capacity(n);
1438 for (i, v) in values.iter().enumerate() {
1439 match v {
1440 Value::Float64(x) => {
1441 non_null[i] = true;
1442 data.push(*x);
1443 }
1444 _ => data.push(0.0),
1445 }
1446 }
1447 NativeColumn::Float64 {
1448 data,
1449 validity: validity_bitmap_from(non_null),
1450 }
1451 }
1452 TypeId::Bool => {
1453 let mut data = Vec::with_capacity(n);
1454 for (i, v) in values.iter().enumerate() {
1455 match v {
1456 Value::Bool(x) => {
1457 non_null[i] = true;
1458 data.push(if *x { 1 } else { 0 });
1459 }
1460 _ => data.push(0),
1461 }
1462 }
1463 NativeColumn::Bool {
1464 data,
1465 validity: validity_bitmap_from(non_null),
1466 }
1467 }
1468 _ => {
1469 let mut offsets = Vec::with_capacity(n + 1);
1470 let mut vals = Vec::new();
1471 offsets.push(0u32);
1472 for (i, v) in values.iter().enumerate() {
1473 if let Value::Bytes(b) = v {
1474 non_null[i] = true;
1475 vals.extend_from_slice(b);
1476 }
1477 offsets.push(vals.len() as u32);
1478 }
1479 NativeColumn::Bytes {
1480 offsets,
1481 values: vals,
1482 validity: validity_bitmap_from(non_null),
1483 }
1484 }
1485 }
1486}
1487
1488pub fn rows_to_native(ty: TypeId, rows: &[Vec<(u16, Value)>], column_id: u16) -> NativeColumn {
1497 let n = rows.len();
1498 let mut non_null = vec![false; n];
1499 fn at(row: &[(u16, Value)], column_id: u16) -> Option<&Value> {
1500 row.iter().find(|(id, _)| *id == column_id).map(|(_, v)| v)
1501 }
1502 match ty {
1503 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date64 | TypeId::Time64 => {
1504 let mut data = Vec::with_capacity(n);
1505 for (i, row) in rows.iter().enumerate() {
1506 match at(row, column_id) {
1507 Some(Value::Int64(x)) => {
1508 non_null[i] = true;
1509 data.push(*x);
1510 }
1511 _ => data.push(0),
1512 }
1513 }
1514 NativeColumn::Int64 {
1515 data,
1516 validity: validity_bitmap_from(non_null),
1517 }
1518 }
1519 TypeId::Float64 => {
1520 let mut data = Vec::with_capacity(n);
1521 for (i, row) in rows.iter().enumerate() {
1522 match at(row, column_id) {
1523 Some(Value::Float64(x)) => {
1524 non_null[i] = true;
1525 data.push(*x);
1526 }
1527 _ => data.push(0.0),
1528 }
1529 }
1530 NativeColumn::Float64 {
1531 data,
1532 validity: validity_bitmap_from(non_null),
1533 }
1534 }
1535 TypeId::Bool => {
1536 let mut data = Vec::with_capacity(n);
1537 for (i, row) in rows.iter().enumerate() {
1538 match at(row, column_id) {
1539 Some(Value::Bool(x)) => {
1540 non_null[i] = true;
1541 data.push(if *x { 1 } else { 0 });
1542 }
1543 _ => data.push(0),
1544 }
1545 }
1546 NativeColumn::Bool {
1547 data,
1548 validity: validity_bitmap_from(non_null),
1549 }
1550 }
1551 _ => {
1552 let mut offsets = Vec::with_capacity(n + 1);
1553 let mut vals = Vec::new();
1554 offsets.push(0u32);
1555 for (i, row) in rows.iter().enumerate() {
1556 if let Some(Value::Bytes(b)) = at(row, column_id) {
1557 non_null[i] = true;
1558 vals.extend_from_slice(b);
1559 }
1560 offsets.push(vals.len() as u32);
1561 }
1562 NativeColumn::Bytes {
1563 offsets,
1564 values: vals,
1565 validity: validity_bitmap_from(non_null),
1566 }
1567 }
1568 }
1569}
1570
1571fn validity_bitmap_from(non_null: impl IntoIterator<Item = bool>) -> Vec<u8> {
1572 let bits: Vec<bool> = non_null.into_iter().collect();
1573 let n = bits.len();
1574 let mut out = vec![0u8; n.div_ceil(8)];
1575 for (i, &b) in bits.iter().enumerate() {
1576 if b {
1577 out[i / 8] |= 1 << (i % 8);
1578 }
1579 }
1580 out
1581}
1582
1583pub fn encode_page_native(
1590 ty: TypeId,
1591 col: &NativeColumn,
1592 encoding: Encoding,
1593 compress: Compress,
1594 le: bool,
1595) -> Result<Vec<u8>> {
1596 let raw = matches!(compress, Compress::Plain) || matches!(encoding, Encoding::Plain);
1597 Ok(match (ty.clone(), col) {
1598 (TypeId::Int64 | TypeId::TimestampNanos, NativeColumn::Int64 { data, validity }) => {
1599 if matches!(encoding, Encoding::Delta) && !raw {
1600 let mut payload = Vec::with_capacity(4 + validity.len() + data.len() * 8);
1601 payload.extend_from_slice(&(validity.len() as u32).to_be_bytes());
1602 payload.extend_from_slice(validity);
1603 let mut deltas = Vec::with_capacity(data.len());
1604 let mut prev = 0i64;
1605 for v in data {
1606 deltas.push(v - prev);
1607 prev = *v;
1608 }
1609 if le {
1610 append_i64_le(&mut payload, &deltas);
1611 } else {
1612 append_i64_be(&mut payload, &deltas);
1613 }
1614 compress_delta_payload(&payload, compress, le)?
1615 } else {
1616 native_plain_page(validity, compress, raw, le, |p| {
1617 if le {
1618 append_i64_le(p, data);
1619 } else {
1620 append_i64_be(p, data);
1621 }
1622 })
1623 }
1624 }
1625 (
1626 TypeId::Float64,
1627 NativeColumn::Float64 {
1628 data: fdata,
1629 validity,
1630 },
1631 ) => native_plain_page(validity, compress, raw, le, |p| {
1632 let bits: &[u64] = bytemuck::cast_slice::<f64, u64>(fdata);
1633 if le {
1634 append_u64_le(p, bits);
1635 } else {
1636 append_u64_be(p, bits);
1637 }
1638 }),
1639 (
1640 TypeId::Bool,
1641 NativeColumn::Bool {
1642 data: bdata,
1643 validity,
1644 },
1645 ) => native_plain_page(validity, compress, raw, le, |p| p.extend_from_slice(bdata)),
1646 (
1647 TypeId::Bytes,
1648 NativeColumn::Bytes {
1649 offsets,
1650 values,
1651 validity,
1652 },
1653 ) => {
1654 if matches!(encoding, Encoding::Dictionary) && !raw {
1655 let dict = dict_encode_bytes_native(offsets, values, validity);
1656 compress_dict_payload(&dict, compress, le)?
1657 } else {
1658 native_plain_page(validity, compress, raw, le, |p| {
1659 let offs: Vec<u64> = offsets.iter().map(|o| *o as u64).collect();
1660 if le {
1661 append_u64_le(p, &offs);
1662 } else {
1663 append_u64_be(p, &offs);
1664 }
1665 p.extend_from_slice(values);
1666 })
1667 }
1668 }
1669 _ => {
1670 return Err(MongrelError::InvalidArgument(format!(
1671 "encode_page_native: unsupported (ty={ty:?})"
1672 )))
1673 }
1674 })
1675}
1676
1677fn compress_delta_payload(payload: &[u8], compress: Compress, le: bool) -> Result<Vec<u8>> {
1681 Ok(match compress {
1682 Compress::Plain => {
1683 let mut out = vec![algo_with_le(ALGO_PLAIN, le)];
1684 out.extend_from_slice(payload);
1685 out
1686 }
1687 Compress::Zstd(level) => {
1688 let mut out = vec![algo_with_le(ALGO_ZSTD_DELTA, le)];
1689 out.extend(zstd_compress_level(payload, level)?);
1690 out
1691 }
1692 Compress::Lz4 => {
1693 let mut out = vec![algo_with_le(ALGO_LZ4_DELTA, le)];
1694 out.extend(lz4_compress(payload));
1695 out
1696 }
1697 })
1698}
1699
1700fn compress_dict_payload(payload: &[u8], compress: Compress, _le: bool) -> Result<Vec<u8>> {
1705 Ok(match compress {
1706 Compress::Plain => {
1707 let mut out = vec![ALGO_PLAIN];
1708 out.extend_from_slice(payload);
1709 out
1710 }
1711 Compress::Zstd(level) => {
1712 let mut out = vec![ALGO_ZSTD_DICT];
1713 out.extend(zstd_compress_level(payload, level)?);
1714 out
1715 }
1716 Compress::Lz4 => {
1717 let mut out = vec![ALGO_LZ4_DICT];
1718 out.extend(lz4_compress(payload));
1719 out
1720 }
1721 })
1722}
1723
1724fn native_plain_page(
1729 validity: &[u8],
1730 compress: Compress,
1731 raw: bool,
1732 le: bool,
1733 fill_payload: impl FnOnce(&mut Vec<u8>),
1734) -> Vec<u8> {
1735 let mut payload = Vec::new();
1736 payload.extend_from_slice(&(validity.len() as u32).to_be_bytes());
1737 payload.extend_from_slice(validity);
1738 fill_payload(&mut payload);
1739 if raw {
1740 let mut out = vec![algo_with_le(ALGO_PLAIN, le)];
1741 out.extend(payload);
1742 out
1743 } else {
1744 match compress {
1745 Compress::Zstd(level) => {
1746 let mut out = vec![algo_with_le(ALGO_ZSTD_PLAIN, le)];
1747 out.extend(zstd_compress_level(&payload, level).expect("zstd compress"));
1748 out
1749 }
1750 Compress::Lz4 => {
1751 let mut out = vec![algo_with_le(ALGO_LZ4_PLAIN, le)];
1752 out.extend(lz4_compress(&payload));
1753 out
1754 }
1755 Compress::Plain => {
1756 let mut out = vec![algo_with_le(ALGO_PLAIN, le)];
1757 out.extend(payload);
1758 out
1759 }
1760 }
1761 }
1762}
1763
1764pub fn decode_page_native(ty: TypeId, page: &[u8], n: usize) -> Result<NativeColumn> {
1769 use std::borrow::Cow;
1770 if page.is_empty() {
1771 return Err(MongrelError::InvalidArgument("empty page".into()));
1772 }
1773 let algo = page[0];
1774 let body = &page[1..];
1775 let le = algo & ALGO_LE_FLAG != 0;
1779 let base = algo & !ALGO_LE_FLAG;
1780 let plain_algo = matches!(base, ALGO_PLAIN | ALGO_ZSTD_PLAIN | ALGO_LZ4_PLAIN);
1781 let delta_algo = matches!(base, ALGO_ZSTD_DELTA | ALGO_LZ4_DELTA);
1782 let dict_algo = matches!(base, ALGO_ZSTD_DICT | ALGO_LZ4_DICT);
1783 if !plain_algo && !delta_algo && !dict_algo {
1784 return Err(MongrelError::InvalidArgument(format!(
1785 "decode_page_native: unsupported algo {algo} for ty {ty:?}"
1786 )));
1787 }
1788 let raw: Cow<[u8]> = match base {
1790 ALGO_PLAIN => Cow::Borrowed(body),
1791 ALGO_ZSTD_PLAIN | ALGO_ZSTD_DELTA | ALGO_ZSTD_DICT => Cow::Owned(zstd_decompress(
1792 body,
1793 max_decompressed_bytes(ty.clone(), n, base),
1794 )?),
1795 ALGO_LZ4_PLAIN | ALGO_LZ4_DELTA | ALGO_LZ4_DICT => Cow::Owned(lz4_decompress(
1796 body,
1797 max_decompressed_bytes(ty.clone(), n, base),
1798 )?),
1799 _ => unreachable!(),
1800 };
1801
1802 if dict_algo {
1805 return if matches!(ty, TypeId::Bytes | TypeId::Json | TypeId::Enum { .. }) {
1806 dict_decode_bytes_native(&raw, n)
1807 } else {
1808 Err(MongrelError::InvalidArgument(format!(
1809 "decode_page_native: dict algo {algo} only valid for variable bytes, got {ty:?}"
1810 )))
1811 };
1812 }
1813
1814 let take_i64 = |p: &[u8]| -> Result<Vec<i64>> {
1816 if le {
1817 take_i64_le(p, n)
1818 } else {
1819 take_i64_be(p, n)
1820 }
1821 };
1822 if delta_algo {
1824 if !matches!(ty, TypeId::Int64 | TypeId::TimestampNanos) {
1825 return Err(MongrelError::InvalidArgument(format!(
1826 "decode_page_native: delta algo {algo} only valid for Int64, got {ty:?}"
1827 )));
1828 }
1829 let (validity, p) = split_validity(&raw)?;
1830 let deltas = take_i64(p)?;
1831 let data = delta_prefix_sum_i64(&deltas);
1832 return Ok(NativeColumn::Int64 { data, validity });
1833 }
1834
1835 match ty {
1837 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date64 | TypeId::Time64 => {
1838 let (validity, p) = split_validity(&raw)?;
1839 Ok(NativeColumn::Int64 {
1840 data: take_i64(p)?,
1841 validity,
1842 })
1843 }
1844 TypeId::Float64 => {
1845 let (validity, p) = split_validity(&raw)?;
1846 let bits = if le {
1847 take_u64_le(p, n)?
1848 } else {
1849 take_u64_be(p, n)?
1850 };
1851 let data: Vec<f64> = bits.into_iter().map(f64::from_bits).collect();
1852 Ok(NativeColumn::Float64 { data, validity })
1853 }
1854 TypeId::Bool => {
1855 let (validity, p) = split_validity(&raw)?;
1856 if p.len() < n {
1857 return Err(MongrelError::InvalidArgument(
1858 "bool payload truncated".into(),
1859 ));
1860 }
1861 Ok(NativeColumn::Bool {
1862 data: p[..n].to_vec(),
1863 validity,
1864 })
1865 }
1866 TypeId::Bytes | TypeId::Json | TypeId::Enum { .. } => {
1867 decode_bytes_plain_payload(&raw, n, le)
1868 }
1869 _ => Err(MongrelError::InvalidArgument(format!(
1870 "decode_page_native: unsupported ty {ty:?}"
1871 ))),
1872 }
1873}
1874
1875fn split_validity(raw: &[u8]) -> Result<(Vec<u8>, &[u8])> {
1876 if raw.len() < 4 {
1877 return Err(MongrelError::InvalidArgument("page validity header".into()));
1878 }
1879 let vlen = u32::from_be_bytes([raw[0], raw[1], raw[2], raw[3]]) as usize;
1880 if 4 + vlen > raw.len() {
1881 return Err(MongrelError::InvalidArgument("page validity range".into()));
1882 }
1883 Ok((raw[4..4 + vlen].to_vec(), &raw[4 + vlen..]))
1884}
1885
1886fn decode_bytes_plain_payload(raw: &[u8], n: usize, le: bool) -> Result<NativeColumn> {
1889 let (validity, p) = split_validity(raw)?;
1890 let table = (n + 1) * 8;
1891 if p.len() < table {
1892 return Err(MongrelError::InvalidArgument(
1893 "bytes offsets truncated".into(),
1894 ));
1895 }
1896 let offsets_be: Vec<u64> = if le {
1897 take_u64_le(p, n + 1)?
1898 } else {
1899 take_u64_be(p, n + 1)?
1900 };
1901 let offsets: Vec<u32> = offsets_be.into_iter().map(|o| o as u32).collect();
1902 let values = p[table..].to_vec();
1903 Ok(NativeColumn::Bytes {
1904 offsets,
1905 values,
1906 validity,
1907 })
1908}
1909
1910fn take_i64_be(p: &[u8], n: usize) -> Result<Vec<i64>> {
1911 if p.len() < n * 8 {
1912 return Err(MongrelError::InvalidArgument(
1913 "int64 payload truncated".into(),
1914 ));
1915 }
1916 Ok(take_u64_be(p, n)?.into_iter().map(|u| u as i64).collect())
1917}
1918
1919fn delta_prefix_sum_i64(deltas: &[i64]) -> Vec<i64> {
1927 let mut out = vec![0i64; deltas.len()];
1928 prefix_sum_i64_into(deltas, &mut out);
1929 out
1930}
1931
1932fn prefix_sum_i64_into(deltas: &[i64], out: &mut [i64]) {
1933 debug_assert_eq!(deltas.len(), out.len());
1934 #[cfg(target_arch = "x86_64")]
1935 {
1936 if is_x86_feature_detected!("avx2") && deltas.len() >= 4 {
1937 unsafe {
1940 prefix_sum_avx2(deltas, out);
1941 }
1942 return;
1943 }
1944 }
1945 prefix_sum_scalar(deltas, out);
1946}
1947
1948fn prefix_sum_scalar(deltas: &[i64], out: &mut [i64]) {
1949 let mut acc = 0i64;
1950 for (i, &d) in deltas.iter().enumerate() {
1951 acc = acc.wrapping_add(d);
1952 out[i] = acc;
1953 }
1954}
1955
1956#[cfg(target_arch = "x86_64")]
1957#[target_feature(enable = "avx2")]
1958unsafe fn prefix_sum_avx2(deltas: &[i64], out: &mut [i64]) {
1959 use std::arch::x86_64::*;
1960 let n = deltas.len();
1961 let mut running = _mm256_setzero_si256(); let mut i = 0usize;
1963 while i + 4 <= n {
1964 let mut x = _mm256_loadu_si256(deltas.as_ptr().add(i) as *const __m256i);
1965 let s1 = _mm256_slli_si256(x, 8); x = _mm256_add_epi64(x, s1); let bc = _mm256_permute4x64_epi64(x, 0x50); let mask = _mm256_set_epi64x(-1, -1, 0, 0); let carry = _mm256_and_si256(bc, mask); x = _mm256_add_epi64(x, carry); x = _mm256_add_epi64(x, running); _mm256_storeu_si256(out.as_mut_ptr().add(i) as *mut __m256i, x);
1974 running = _mm256_permute4x64_epi64(x, 0xFF); i += 4;
1976 }
1977 let mut acc = if i == 0 { 0 } else { out[i - 1] };
1980 while i < n {
1981 acc = acc.wrapping_add(deltas[i]);
1982 out[i] = acc;
1983 i += 1;
1984 }
1985}
1986
1987fn append_u64_be(out: &mut Vec<u8>, data: &[u64]) {
1989 if cfg!(target_endian = "little") {
1990 let swapped: Vec<u64> = data.iter().map(|v| v.swap_bytes()).collect();
1991 out.extend_from_slice(bytemuck::cast_slice::<u64, u8>(&swapped));
1992 } else {
1993 out.extend_from_slice(bytemuck::cast_slice::<u64, u8>(data));
1994 }
1995}
1996
1997fn append_i64_be(out: &mut Vec<u8>, data: &[i64]) {
1999 if cfg!(target_endian = "little") {
2000 let swapped: Vec<i64> = data.iter().map(|v| v.swap_bytes()).collect();
2001 out.extend_from_slice(bytemuck::cast_slice::<i64, u8>(&swapped));
2002 } else {
2003 out.extend_from_slice(bytemuck::cast_slice::<i64, u8>(data));
2004 }
2005}
2006
2007fn take_u64_be(p: &[u8], n: usize) -> Result<Vec<u64>> {
2011 if p.len() < n * 8 {
2012 return Err(MongrelError::InvalidArgument(
2013 "u64 payload truncated".into(),
2014 ));
2015 }
2016 let bytes = &p[..n * 8];
2017 if let Ok(native) = bytemuck::try_cast_slice::<u8, u64>(bytes) {
2018 Ok(native.iter().map(|v| v.swap_bytes()).collect())
2019 } else {
2020 let mut out = Vec::with_capacity(n);
2021 for chunk in bytes.chunks_exact(8) {
2022 out.push(u64::from_be_bytes(chunk.try_into().unwrap()));
2023 }
2024 Ok(out)
2025 }
2026}
2027
2028fn append_u64_le(out: &mut Vec<u8>, data: &[u64]) {
2032 if cfg!(target_endian = "little") {
2033 out.extend_from_slice(bytemuck::cast_slice::<u64, u8>(data));
2034 } else {
2035 let swapped: Vec<u64> = data.iter().map(|v| v.swap_bytes()).collect();
2036 out.extend_from_slice(bytemuck::cast_slice::<u64, u8>(&swapped));
2037 }
2038}
2039
2040fn append_i64_le(out: &mut Vec<u8>, data: &[i64]) {
2042 if cfg!(target_endian = "little") {
2043 out.extend_from_slice(bytemuck::cast_slice::<i64, u8>(data));
2044 } else {
2045 let swapped: Vec<i64> = data.iter().map(|v| v.swap_bytes()).collect();
2046 out.extend_from_slice(bytemuck::cast_slice::<i64, u8>(&swapped));
2047 }
2048}
2049
2050fn take_u64_le(p: &[u8], n: usize) -> Result<Vec<u64>> {
2056 if p.len() < n * 8 {
2057 return Err(MongrelError::InvalidArgument(
2058 "u64 payload truncated".into(),
2059 ));
2060 }
2061 let bytes = &p[..n * 8];
2062 if cfg!(target_endian = "little") {
2063 if let Ok(native) = bytemuck::try_cast_slice::<u8, u64>(bytes) {
2064 return Ok(native.to_vec());
2066 }
2067 Ok(bytes
2068 .chunks_exact(8)
2069 .map(|c| u64::from_le_bytes(c.try_into().unwrap()))
2070 .collect())
2071 } else {
2072 Ok(bytes
2073 .chunks_exact(8)
2074 .map(|c| u64::from_le_bytes(c.try_into().unwrap()))
2075 .collect())
2076 }
2077}
2078
2079fn take_i64_le(p: &[u8], n: usize) -> Result<Vec<i64>> {
2081 if p.len() < n * 8 {
2082 return Err(MongrelError::InvalidArgument(
2083 "int64 payload truncated".into(),
2084 ));
2085 }
2086 Ok(take_u64_le(p, n)?.into_iter().map(|u| u as i64).collect())
2087}
2088
2089fn dict_encode_bytes_native(offsets: &[u32], values: &[u8], validity: &[u8]) -> Vec<u8> {
2090 let n = offsets.len() - 1;
2091 let mut table: Vec<Vec<u8>> = Vec::new();
2092 let mut index_of: std::collections::HashMap<&[u8], u32> = std::collections::HashMap::new();
2093 let mut indices: Vec<u32> = Vec::with_capacity(n);
2094 for i in 0..n {
2095 let lo = offsets[i] as usize;
2096 let hi = offsets[i + 1] as usize;
2097 let slice = &values[lo..hi];
2098 let idx = if let Some(&idx) = index_of.get(slice) {
2099 idx
2100 } else {
2101 let idx = table.len() as u32;
2102 index_of.insert(slice, idx);
2103 table.push(slice.to_vec());
2104 idx
2105 };
2106 indices.push(idx);
2107 }
2108 let mut out = Vec::new();
2109 out.extend_from_slice(&(validity.len() as u32).to_be_bytes());
2110 out.extend_from_slice(validity);
2111 out.extend_from_slice(&(indices.len() as u32).to_be_bytes());
2112 for i in &indices {
2113 out.extend_from_slice(&i.to_be_bytes());
2114 }
2115 out.extend_from_slice(&(table.len() as u32).to_be_bytes());
2116 for entry in &table {
2117 out.extend_from_slice(&(entry.len() as u32).to_be_bytes());
2118 out.extend_from_slice(entry);
2119 }
2120 out
2121}
2122
2123fn dict_decode_bytes_native(data: &[u8], n: usize) -> Result<NativeColumn> {
2124 let mut cur = 0usize;
2125 let vlen = read_u32_be(data, &mut cur)? as usize;
2126 let validity = checked_slice(data, &mut cur, vlen)?.to_vec();
2127 let index_count = read_u32_be(data, &mut cur)? as usize;
2128 if index_count < n {
2129 return Err(MongrelError::InvalidArgument("dict index_count < n".into()));
2130 }
2131 let mut indices = Vec::with_capacity(index_count.min(n));
2132 for _ in 0..index_count {
2133 indices.push(read_u32_be(data, &mut cur)?);
2134 }
2135 let table_count = read_u32_be(data, &mut cur)? as usize;
2136 let mut table: Vec<(usize, usize)> = Vec::with_capacity(table_count); let mut values = Vec::new();
2138 for _ in 0..table_count {
2139 let len = read_u32_be(data, &mut cur)? as usize;
2140 let chunk = checked_slice(data, &mut cur, len)?;
2141 let start = values.len();
2142 values.extend_from_slice(chunk);
2143 table.push((start, len));
2144 }
2145 let mut merged = Vec::new();
2146 let mut offs = vec![0u32];
2147 for (i, &idx) in indices.iter().enumerate().take(n) {
2148 let non_null = (validity.get(i / 8).copied().unwrap_or(0) >> (i % 8)) & 1 == 1;
2152 if non_null {
2153 let (start, len) = table
2154 .get(idx as usize)
2155 .copied()
2156 .ok_or_else(|| MongrelError::InvalidArgument("dict index out of range".into()))?;
2157 merged.extend_from_slice(&values[start..start + len]);
2158 }
2159 offs.push(merged.len() as u32);
2160 }
2161 Ok(NativeColumn::Bytes {
2162 offsets: offs,
2163 values: merged,
2164 validity,
2165 })
2166}
2167
2168#[cfg(test)]
2169mod native_tests {
2170 use super::*;
2171
2172 #[test]
2173 fn native_int64_plain_round_trip() {
2174 let col = NativeColumn::Int64 {
2175 data: (0..1000).collect(),
2176 validity: full_validity(1000),
2177 };
2178 let page = encode_page_native(
2179 TypeId::Int64,
2180 &col,
2181 Encoding::Zstd,
2182 Compress::Zstd(3),
2183 false,
2184 )
2185 .unwrap();
2186 let back = decode_page_native(TypeId::Int64, &page, 1000).unwrap();
2187 match back {
2188 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2189 _ => panic!(),
2190 }
2191 }
2192
2193 #[test]
2194 fn native_int64_delta_crushes_sequential() {
2195 let col = NativeColumn::int64_sequence(0, 100_000);
2196 let plain = encode_page_native(
2197 TypeId::Int64,
2198 &col,
2199 Encoding::Zstd,
2200 Compress::Zstd(3),
2201 false,
2202 )
2203 .unwrap();
2204 let delta = encode_page_native(
2205 TypeId::Int64,
2206 &col,
2207 Encoding::Delta,
2208 Compress::Zstd(3),
2209 false,
2210 )
2211 .unwrap();
2212 assert!(
2213 delta.len() < plain.len() / 5,
2214 "delta must crush sequential ints"
2215 );
2216 let back = decode_page_native(TypeId::Int64, &delta, 100_000).unwrap();
2217 match back {
2218 NativeColumn::Int64 { data, .. } => assert_eq!(data.len(), 100_000),
2219 _ => panic!(),
2220 }
2221 }
2222
2223 #[test]
2224 fn native_bytes_dict_round_trip() {
2225 let n = 500;
2226 let mut offsets = vec![0u32];
2227 let mut values = Vec::new();
2228 for i in 0..n {
2229 let s = ["red", "green", "blue"][i % 3];
2230 values.extend_from_slice(s.as_bytes());
2231 offsets.push(values.len() as u32);
2232 }
2233 let col = NativeColumn::Bytes {
2234 offsets,
2235 values,
2236 validity: full_validity(n),
2237 };
2238 let page = encode_page_native(
2239 TypeId::Bytes,
2240 &col,
2241 Encoding::Dictionary,
2242 Compress::Zstd(3),
2243 false,
2244 )
2245 .unwrap();
2246 assert!(page.len() < 100, "dict page tiny, got {}", page.len());
2247 let back = decode_page_native(TypeId::Bytes, &page, n).unwrap();
2248 assert_eq!(back.len(), n);
2249 }
2250
2251 #[test]
2252 fn native_gather_picks_indices() {
2253 let col = NativeColumn::Int64 {
2254 data: vec![10, 20, 30, 40],
2255 validity: full_validity(4),
2256 };
2257 let g = col.gather(&[0, 2, 3]);
2258 match g {
2259 NativeColumn::Int64 { data, .. } => assert_eq!(data, vec![10, 30, 40]),
2260 _ => panic!(),
2261 }
2262 }
2263
2264 #[test]
2267 fn native_plain_no_zstd_round_trips_all_types() {
2268 let i = NativeColumn::Int64 {
2269 data: (0..1000).collect(),
2270 validity: full_validity(1000),
2271 };
2272 let p =
2273 encode_page_native(TypeId::Int64, &i, Encoding::Plain, Compress::Plain, false).unwrap();
2274 assert_eq!(p[0], ALGO_PLAIN, "Int64 plain must be ALGO_PLAIN");
2275 match decode_page_native(TypeId::Int64, &p, 1000).unwrap() {
2276 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2277 _ => panic!(),
2278 }
2279
2280 let f = NativeColumn::Float64 {
2281 data: (0..500).map(|x| x as f64 * 1.5).collect(),
2282 validity: full_validity(500),
2283 };
2284 let p = encode_page_native(TypeId::Float64, &f, Encoding::Plain, Compress::Plain, false)
2285 .unwrap();
2286 assert_eq!(p[0], ALGO_PLAIN);
2287 match decode_page_native(TypeId::Float64, &p, 500).unwrap() {
2288 NativeColumn::Float64 { data, .. } => {
2289 assert_eq!(data, (0..500).map(|x| x as f64 * 1.5).collect::<Vec<_>>())
2290 }
2291 _ => panic!(),
2292 }
2293
2294 let b = NativeColumn::Bool {
2295 data: (0..64).map(|i| (i % 2) as u8).collect(),
2296 validity: full_validity(64),
2297 };
2298 let p =
2299 encode_page_native(TypeId::Bool, &b, Encoding::Plain, Compress::Plain, false).unwrap();
2300 assert_eq!(p[0], ALGO_PLAIN);
2301 match decode_page_native(TypeId::Bool, &p, 64).unwrap() {
2302 NativeColumn::Bool { data, .. } => {
2303 assert_eq!(data, (0..64).map(|i| (i % 2) as u8).collect::<Vec<_>>())
2304 }
2305 _ => panic!(),
2306 }
2307
2308 let mut offsets = vec![0u32];
2309 let mut values = Vec::new();
2310 for i in 0..200u32 {
2311 values.extend_from_slice(format!("v{i}").as_bytes());
2312 offsets.push(values.len() as u32);
2313 }
2314 let s = NativeColumn::Bytes {
2315 offsets,
2316 values,
2317 validity: full_validity(200),
2318 };
2319 let p =
2320 encode_page_native(TypeId::Bytes, &s, Encoding::Plain, Compress::Plain, false).unwrap();
2321 assert_eq!(p[0], ALGO_PLAIN);
2322 match decode_page_native(TypeId::Bytes, &p, 200).unwrap() {
2323 NativeColumn::Bytes {
2324 offsets: o,
2325 values: v,
2326 ..
2327 } => {
2328 assert_eq!(o.len(), 201);
2329 for i in 0..200 {
2330 let lo = o[i] as usize;
2331 let hi = o[i + 1] as usize;
2332 assert_eq!(&v[lo..hi], format!("v{i}").as_bytes());
2333 }
2334 }
2335 _ => panic!(),
2336 }
2337 }
2338
2339 #[test]
2342 fn lz4_pages_round_trip_all_types() {
2343 let i = NativeColumn::int64_sequence(0, 1000);
2345 let p =
2346 encode_page_native(TypeId::Int64, &i, Encoding::Delta, Compress::Lz4, false).unwrap();
2347 assert_eq!(p[0], ALGO_LZ4_DELTA);
2348 match decode_page_native(TypeId::Int64, &p, 1000).unwrap() {
2349 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2350 _ => panic!(),
2351 }
2352 let p =
2354 encode_page_native(TypeId::Int64, &i, Encoding::Zstd, Compress::Lz4, false).unwrap();
2355 assert_eq!(p[0], ALGO_LZ4_PLAIN);
2356 match decode_page_native(TypeId::Int64, &p, 1000).unwrap() {
2357 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2358 _ => panic!(),
2359 }
2360 let f = NativeColumn::Float64 {
2362 data: (0..500).map(|x| x as f64 * 1.5).collect(),
2363 validity: full_validity(500),
2364 };
2365 let p =
2366 encode_page_native(TypeId::Float64, &f, Encoding::Zstd, Compress::Lz4, false).unwrap();
2367 assert_eq!(p[0], ALGO_LZ4_PLAIN);
2368 match decode_page_native(TypeId::Float64, &p, 500).unwrap() {
2369 NativeColumn::Float64 { data, .. } => {
2370 assert_eq!(data, (0..500).map(|x| x as f64 * 1.5).collect::<Vec<_>>())
2371 }
2372 _ => panic!(),
2373 }
2374 let mut offsets = vec![0u32];
2376 let mut values = Vec::new();
2377 for i in 0..300u32 {
2378 values.extend_from_slice(["red", "green", "blue"][(i % 3) as usize].as_bytes());
2379 offsets.push(values.len() as u32);
2380 }
2381 let s = NativeColumn::Bytes {
2382 offsets,
2383 values,
2384 validity: full_validity(300),
2385 };
2386 let p = encode_page_native(
2387 TypeId::Bytes,
2388 &s,
2389 Encoding::Dictionary,
2390 Compress::Lz4,
2391 false,
2392 )
2393 .unwrap();
2394 assert_eq!(p[0], ALGO_LZ4_DICT);
2395 match decode_page_native(TypeId::Bytes, &p, 300).unwrap() {
2396 NativeColumn::Bytes { offsets: o, .. } => assert_eq!(o.len(), 301),
2397 _ => panic!(),
2398 }
2399 }
2400
2401 #[test]
2405 fn le_pages_round_trip_all_types() {
2406 let assert_le = |page: &[u8]| assert_ne!(page[0] & ALGO_LE_FLAG, 0, "LE flag must be set");
2407
2408 let i = NativeColumn::Int64 {
2410 data: (0..1000).collect(),
2411 validity: full_validity(1000),
2412 };
2413 let p =
2414 encode_page_native(TypeId::Int64, &i, Encoding::Plain, Compress::Plain, true).unwrap();
2415 assert_eq!(p[0], ALGO_LE_FLAG, "raw LE Int64 algo = flag only");
2416 match decode_page_native(TypeId::Int64, &p, 1000).unwrap() {
2417 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2418 _ => panic!(),
2419 }
2420
2421 let p =
2423 encode_page_native(TypeId::Int64, &i, Encoding::Zstd, Compress::Zstd(3), true).unwrap();
2424 assert_le(&p);
2425 match decode_page_native(TypeId::Int64, &p, 1000).unwrap() {
2426 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2427 _ => panic!(),
2428 }
2429
2430 let seq = NativeColumn::int64_sequence(0, 1000);
2432 let p =
2433 encode_page_native(TypeId::Int64, &seq, Encoding::Delta, Compress::Lz4, true).unwrap();
2434 assert_eq!(p[0], ALGO_LZ4_DELTA | ALGO_LE_FLAG);
2435 match decode_page_native(TypeId::Int64, &p, 1000).unwrap() {
2436 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2437 _ => panic!(),
2438 }
2439
2440 let f = NativeColumn::Float64 {
2442 data: (0..500).map(|x| x as f64 * 1.5).collect(),
2443 validity: full_validity(500),
2444 };
2445 let p = encode_page_native(TypeId::Float64, &f, Encoding::Zstd, Compress::Zstd(3), true)
2446 .unwrap();
2447 assert_le(&p);
2448 match decode_page_native(TypeId::Float64, &p, 500).unwrap() {
2449 NativeColumn::Float64 { data, .. } => {
2450 assert_eq!(data, (0..500).map(|x| x as f64 * 1.5).collect::<Vec<_>>())
2451 }
2452 _ => panic!(),
2453 }
2454 }
2455
2456 #[test]
2460 fn malicious_lz4_size_prefix_is_rejected() {
2461 let mut evil = vec![ALGO_LZ4_PLAIN];
2463 evil.extend_from_slice(&0xFFFF_FFFFu32.to_le_bytes());
2464 evil.extend_from_slice(b"junk");
2465 let err =
2466 decode_page_native(TypeId::Int64, &evil, 1000).expect_err("must reject oversize lz4");
2467 let msg = format!("{err}");
2468 assert!(msg.contains("exceeds page limit"), "got: {msg}");
2469
2470 let err = decode_page(TypeId::Int64, &evil, 1000).expect_err("must reject oversize lz4");
2472 assert!(format!("{err}").contains("exceeds page limit"));
2473 }
2474
2475 #[test]
2479 fn truncated_dict_payload_does_not_panic() {
2480 let mut body = vec![ALGO_ZSTD_DICT];
2483 body.extend_from_slice(&zstd_compress(&[0u8; 2]).unwrap()[..]);
2484 decode_page(TypeId::Bytes, &body, 4).expect_err("value dict trunc must Err");
2485
2486 let truncated: &[u8] = &[
2488 0x00, 0x00, 0x00, 0x10, ];
2490 dict_decode_bytes_native(truncated, 2).expect_err("native dict trunc must Err");
2491 dict_decode_bytes(truncated, 2).expect_err("value dict trunc must Err");
2492
2493 let mut bad = vec![0u8; 0];
2496 bad.extend_from_slice(&1u32.to_be_bytes()); bad.push(0xFF); bad.extend_from_slice(&1u32.to_be_bytes()); bad.extend_from_slice(&9u32.to_be_bytes()); bad.extend_from_slice(&0u32.to_be_bytes()); dict_decode_bytes(&bad, 1).expect_err("oob index must Err");
2502 dict_decode_bytes_native(&bad, 1).expect_err("oob index must Err");
2503 }
2504
2505 #[test]
2509 fn delta_prefix_sum_matches_scalar_all_lengths() {
2510 let mut deltas: Vec<i64> = Vec::with_capacity(2000);
2512 let mut s: u64 = 0x9E37_79B9_7F4A_7C15;
2513 for _ in 0..2000 {
2514 s = s
2515 .wrapping_mul(6364136223846793005)
2516 .wrapping_add(1442695040888963407);
2517 deltas.push((s as i32) as i64); }
2519 for &len in &[
2520 0usize, 1, 2, 3, 4, 5, 6, 7, 8, 9, 15, 16, 17, 31, 32, 33, 63, 64, 65, 127, 128, 129,
2521 1000,
2522 ] {
2523 let input = &deltas[..len];
2524 let simd = delta_prefix_sum_i64(input);
2525 let mut ref_out = vec![0i64; len];
2526 prefix_sum_scalar(input, &mut ref_out);
2527 assert_eq!(simd, ref_out, "len {len}: SIMD prefix sum diverged");
2528 }
2529
2530 let mut ids = vec![0i64; 333];
2532 for x in ids.iter_mut().skip(1) {
2533 *x = 1;
2534 }
2535 let got = delta_prefix_sum_i64(&ids);
2536 assert_eq!(got, (0..333).map(|x| x as i64).collect::<Vec<_>>());
2537 }
2538
2539 #[test]
2542 fn delta_page_decodes_with_null_via_vectorized_path() {
2543 let mut validity = full_validity(17);
2544 validity[9 / 8] &= !(1 << (9 % 8)); let col = NativeColumn::Int64 {
2546 data: (0..17).collect(),
2547 validity: validity.clone(),
2548 };
2549 for (name, comp) in [("zstd", Compress::Zstd(3)), ("lz4", Compress::Lz4)] {
2550 let page =
2551 encode_page_native(TypeId::Int64, &col, Encoding::Delta, comp, false).unwrap();
2552 assert!(matches!(page[0], ALGO_ZSTD_DELTA | ALGO_LZ4_DELTA));
2554 let back = decode_page_native(TypeId::Int64, &page, 17).unwrap();
2555 match back {
2556 NativeColumn::Int64 { data, validity: v } => {
2557 assert_eq!(data, (0..17).collect::<Vec<_>>(), "{name} data");
2558 assert_eq!(v, validity, "{name} validity");
2559 }
2560 _ => panic!(),
2561 }
2562 }
2563 }
2564}