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