1use crate::error::{MongrelError, Result};
16use crate::memtable::Value;
17use crate::page::Encoding;
18use crate::schema::TypeId;
19use serde::{Deserialize, Serialize};
20
21pub fn encode_column(ty: TypeId, values: &[Value]) -> Result<Vec<u8>> {
23 let n = values.len();
24 let validity = validity_bitmap(values);
25 let payload = match ty {
26 TypeId::Int64 => fixed_encode(values, 8, |v| match v {
27 Value::Int64(x) => Ok(x.to_be_bytes().to_vec()),
28 Value::Null => Ok(vec![0; 8]),
29 _ => Err(type_mismatch(ty, v)),
30 })?,
31 TypeId::Float64 => fixed_encode(values, 8, |v| match v {
32 Value::Float64(f) => Ok(f.to_bits().to_be_bytes().to_vec()),
33 Value::Null => Ok(vec![0; 8]),
34 _ => Err(type_mismatch(ty, v)),
35 })?,
36 TypeId::TimestampNanos => fixed_encode(values, 8, |v| match v {
37 Value::Int64(x) => Ok(x.to_be_bytes().to_vec()),
38 Value::Null => Ok(vec![0; 8]),
39 _ => Err(type_mismatch(ty, v)),
40 })?,
41 TypeId::Bool => fixed_encode(values, 1, |v| match v {
42 Value::Bool(b) => Ok(vec![*b as u8]),
43 Value::Null => Ok(vec![0]),
44 _ => Err(type_mismatch(ty, v)),
45 })?,
46 TypeId::Int32 | TypeId::UInt32 | TypeId::Date32 => fixed_encode(values, 4, |v| match v {
47 Value::Int64(x) => Ok((*x as i32).to_be_bytes().to_vec()),
48 Value::Null => Ok(vec![0; 4]),
49 _ => Err(type_mismatch(ty, v)),
50 })?,
51 TypeId::Bytes => bytes_encode(values)?,
52 TypeId::Embedding { dim } => embedding_encode(values, dim)?,
53 other => {
54 return Err(MongrelError::Schema(format!(
55 "encoding for type {other:?} not implemented yet"
56 )))
57 }
58 };
59 let mut page = Vec::with_capacity(4 + validity.len() + payload.len());
60 page.extend_from_slice(&(validity.len() as u32).to_be_bytes());
61 page.extend_from_slice(&validity);
62 page.extend_from_slice(&payload);
63 let _ = n;
64 Ok(page)
65}
66
67pub fn decode_column(ty: TypeId, page: &[u8], n: usize, le: bool) -> Result<Vec<Value>> {
72 if page.len() < 4 {
73 return Err(MongrelError::InvalidArgument(
74 "page too short for header".into(),
75 ));
76 }
77 let vlen = u32::from_be_bytes([page[0], page[1], page[2], page[3]]) as usize;
78 if 4 + vlen > page.len() {
79 return Err(MongrelError::InvalidArgument(
80 "page validity out of range".into(),
81 ));
82 }
83 let validity = &page[4..4 + vlen];
84 let payload = &page[4 + vlen..];
85
86 let mut out = Vec::with_capacity(n);
87 let mut cur = 0usize;
88 let i64_at = |b: &[u8]| -> i64 {
89 if le {
90 i64::from_le_bytes(b.try_into().unwrap())
91 } else {
92 i64::from_be_bytes(b.try_into().unwrap())
93 }
94 };
95 let u64_at = |b: &[u8]| -> u64 {
96 if le {
97 u64::from_le_bytes(b.try_into().unwrap())
98 } else {
99 u64::from_be_bytes(b.try_into().unwrap())
100 }
101 };
102 for i in 0..n {
103 let non_null = (validity.get(i / 8).copied().unwrap_or(0) >> (i % 8)) & 1 == 1;
104 if !non_null {
105 out.push(Value::Null);
106 advance_null(&ty, payload, &mut cur)?; continue;
108 }
109 let val = match ty {
110 TypeId::Int64 | TypeId::TimestampNanos => {
111 let b = take(payload, &mut cur, 8)?;
112 Value::Int64(i64_at(&b))
113 }
114 TypeId::Float64 => {
115 let b = take(payload, &mut cur, 8)?;
116 Value::Float64(f64::from_bits(u64_at(&b)))
117 }
118 TypeId::Bool => {
119 let b = take(payload, &mut cur, 1)?;
120 Value::Bool(b[0] != 0)
121 }
122 TypeId::Int32 | TypeId::UInt32 | TypeId::Date32 => {
123 let b = take(payload, &mut cur, 4)?;
124 let v = if le {
125 i32::from_le_bytes(b.try_into().unwrap())
126 } else {
127 i32::from_be_bytes(b.try_into().unwrap())
128 };
129 Value::Int64(v as i64)
130 }
131 TypeId::Bytes => {
132 let bytes_start = (n + 1) * 8;
133 let lo = read_off(payload, i, le);
134 let hi = read_off(payload, i + 1, le);
135 Value::Bytes(payload[bytes_start + lo..bytes_start + hi].to_vec())
136 }
137 TypeId::Embedding { dim } => {
138 let mut acc = Vec::with_capacity(dim as usize);
139 for _ in 0..dim {
140 let b = take(payload, &mut cur, 4)?;
141 acc.push(f32::from_bits(u32::from_be_bytes(b.try_into().unwrap())));
142 }
143 Value::Embedding(acc)
144 }
145 other => {
146 return Err(MongrelError::Schema(format!(
147 "decoding for type {other:?} not implemented yet"
148 )))
149 }
150 };
151 out.push(val);
152 }
153 Ok(out)
154}
155
156fn validity_bitmap(values: &[Value]) -> Vec<u8> {
159 let n = values.len();
160 let mut bits = vec![0u8; n.div_ceil(8)];
161 for (i, v) in values.iter().enumerate() {
162 if !matches!(v, Value::Null) {
163 bits[i / 8] |= 1 << (i % 8);
164 }
165 }
166 bits
167}
168
169fn fixed_encode(
170 values: &[Value],
171 _width: usize,
172 mut enc: impl FnMut(&Value) -> Result<Vec<u8>>,
173) -> Result<Vec<u8>> {
174 let mut out = Vec::new();
175 for v in values {
176 out.extend_from_slice(&enc(v)?);
177 }
178 Ok(out)
179}
180
181fn embedding_encode(values: &[Value], dim: u32) -> Result<Vec<u8>> {
182 let mut out = Vec::with_capacity(values.len() * dim as usize * 4);
183 for v in values {
184 match v {
185 Value::Embedding(vec) => {
186 if vec.len() != dim as usize {
187 return Err(MongrelError::Schema(format!(
188 "embedding dimension mismatch: expected {dim}, got {}",
189 vec.len()
190 )));
191 }
192 for x in vec {
193 out.extend_from_slice(&x.to_bits().to_be_bytes());
194 }
195 }
196 Value::Null => {
197 for _ in 0..dim {
198 out.extend_from_slice(&0u32.to_be_bytes());
199 }
200 }
201 _ => return Err(type_mismatch(TypeId::Embedding { dim }, v)),
202 }
203 }
204 Ok(out)
205}
206
207fn bytes_encode(values: &[Value]) -> Result<Vec<u8>> {
208 let n = values.len();
209 let mut offsets = Vec::with_capacity((n + 1) * 8);
210 let mut data = Vec::new();
211 let mut off = 0u64;
212 offsets.extend_from_slice(&off.to_be_bytes()); for v in values {
214 if let Value::Bytes(b) = v {
215 data.extend_from_slice(b);
216 off = off
217 .checked_add(b.len() as u64)
218 .ok_or_else(|| MongrelError::InvalidArgument("bytes length overflow".into()))?;
219 }
220 offsets.extend_from_slice(&off.to_be_bytes());
221 }
222 let mut out = offsets;
223 out.extend_from_slice(&data);
224 Ok(out)
225}
226
227fn read_off(payload: &[u8], idx: usize, le: bool) -> usize {
228 let s = idx * 8;
229 if le {
230 u64::from_le_bytes(payload[s..s + 8].try_into().unwrap()) as usize
231 } else {
232 u64::from_be_bytes(payload[s..s + 8].try_into().unwrap()) as usize
233 }
234}
235
236fn take(payload: &[u8], cur: &mut usize, n: usize) -> Result<Vec<u8>> {
237 if *cur + n > payload.len() {
238 return Err(MongrelError::InvalidArgument("payload truncated".into()));
239 }
240 let s = &payload[*cur..*cur + n];
241 *cur += n;
242 Ok(s.to_vec())
243}
244
245fn advance_null(ty: &TypeId, payload: &[u8], cur: &mut usize) -> Result<()> {
247 let w = match ty {
248 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos => 8,
249 TypeId::Int32 | TypeId::UInt32 | TypeId::Date32 => 4,
250 TypeId::Bool => 1,
251 TypeId::Bytes | TypeId::Embedding { .. } => return Ok(()),
254 _ => return Ok(()),
255 };
256 if *cur + w > payload.len() {
257 return Err(MongrelError::InvalidArgument(
258 "payload truncated at null".into(),
259 ));
260 }
261 *cur += w;
262 Ok(())
263}
264
265fn type_mismatch(ty: TypeId, v: &Value) -> MongrelError {
266 MongrelError::Schema(format!("type mismatch: column {ty:?}, value {v:?}"))
267}
268
269const 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;
290
291#[inline]
292fn algo_with_le(algo: u8, le: bool) -> u8 {
293 if le {
294 algo | ALGO_LE_FLAG
295 } else {
296 algo
297 }
298}
299
300#[derive(Debug, Clone, Copy)]
304pub enum Compress {
305 Plain,
307 Zstd(i32),
309 Lz4,
312}
313
314fn zstd_compress(data: &[u8]) -> Result<Vec<u8>> {
315 zstd_compress_level(data, 3)
316}
317
318fn zstd_compress_level(data: &[u8], level: i32) -> Result<Vec<u8>> {
323 zstd::encode_all(data, level)
324 .map_err(|e| MongrelError::InvalidArgument(format!("zstd compress: {e}")))
325}
326
327fn lz4_compress(data: &[u8]) -> Vec<u8> {
333 lz4_flex::block::compress_prepend_size(data)
334}
335
336const MAX_VAR_BYTES_PER_ROW: usize = 1 << 14; fn max_decompressed_bytes(ty: TypeId, n: usize, algo: u8) -> usize {
343 let validity = 4 + n.div_ceil(8);
345 if matches!(algo, ALGO_ZSTD_DELTA | ALGO_LZ4_DELTA) {
346 return validity + n.saturating_mul(8);
347 }
348 if matches!(algo, ALGO_ZSTD_DICT | ALGO_LZ4_DICT) {
349 return validity + 8 + n.saturating_mul(4) + n.saturating_mul(4 + MAX_VAR_BYTES_PER_ROW);
352 }
353 let payload = match ty {
354 TypeId::Bytes => (n + 1).saturating_mul(8) + n.saturating_mul(MAX_VAR_BYTES_PER_ROW),
355 TypeId::Embedding { dim } => (dim as usize).saturating_mul(8).saturating_mul(n),
356 _ => n.saturating_mul(ty.fixed_size().unwrap_or(8)),
357 };
358 validity + payload
359}
360
361fn lz4_decompress(data: &[u8], max_bytes: usize) -> Result<Vec<u8>> {
362 if data.len() < 4 {
367 return Err(MongrelError::InvalidArgument("lz4 page truncated".into()));
368 }
369 let declared = u32::from_le_bytes([data[0], data[1], data[2], data[3]]) as usize;
370 if declared > max_bytes {
371 return Err(MongrelError::InvalidArgument(format!(
372 "lz4 declared size {declared} exceeds page limit {max_bytes}"
373 )));
374 }
375 lz4_flex::block::decompress(&data[4..], declared)
376 .map_err(|e| MongrelError::InvalidArgument(format!("lz4 decompress: {e}")))
377}
378
379fn zstd_decompress(data: &[u8], max_bytes: usize) -> Result<Vec<u8>> {
380 let out = zstd::decode_all(data)
383 .map_err(|e| MongrelError::InvalidArgument(format!("zstd decompress: {e}")))?;
384 if out.len() > max_bytes {
385 return Err(MongrelError::InvalidArgument(format!(
386 "zstd output {} exceeds page limit {max_bytes}",
387 out.len()
388 )));
389 }
390 Ok(out)
391}
392
393pub fn encode_page(ty: TypeId, values: &[Value], encoding: Encoding) -> Result<Vec<u8>> {
396 Ok(match encoding {
397 Encoding::Plain => {
398 let mut out = vec![ALGO_PLAIN];
399 out.extend(encode_column(ty, values)?);
400 out
401 }
402 Encoding::Dictionary if matches!(ty, TypeId::Bytes) => {
403 let dict = dict_encode_bytes(values);
404 let mut out = vec![ALGO_ZSTD_DICT];
405 out.extend(zstd_compress(&dict)?);
406 out
407 }
408 _ => {
409 let mut out = vec![ALGO_ZSTD_PLAIN];
410 let enc = encode_column(ty, values)?;
411 out.extend(zstd_compress(&enc)?);
412 out
413 }
414 })
415}
416
417pub fn decode_page(ty: TypeId, page: &[u8], n: usize) -> Result<Vec<Value>> {
419 if page.is_empty() {
420 return Err(MongrelError::InvalidArgument("empty page".into()));
421 }
422 let algo = page[0];
423 let le = algo & ALGO_LE_FLAG != 0;
428 let base = algo & !ALGO_LE_FLAG;
429 let body = &page[1..];
430 let raw_owned;
432 let raw: &[u8] = match base {
433 ALGO_PLAIN => body,
434 ALGO_ZSTD_PLAIN | ALGO_ZSTD_DICT | ALGO_ZSTD_DELTA => {
435 raw_owned = zstd_decompress(body, max_decompressed_bytes(ty, n, base))?;
436 &raw_owned
437 }
438 ALGO_LZ4_PLAIN | ALGO_LZ4_DICT | ALGO_LZ4_DELTA => {
439 raw_owned = lz4_decompress(body, max_decompressed_bytes(ty, n, base))?;
440 &raw_owned
441 }
442 other => {
443 return Err(MongrelError::InvalidArgument(format!(
444 "unknown page algo {other}"
445 )))
446 }
447 };
448 match base {
449 ALGO_PLAIN | ALGO_ZSTD_PLAIN | ALGO_LZ4_PLAIN => decode_column(ty, raw, n, le),
450 ALGO_ZSTD_DICT | ALGO_LZ4_DICT => dict_decode_bytes(raw, n),
451 ALGO_ZSTD_DELTA | ALGO_LZ4_DELTA => decode_int64_delta_values(raw, ty, n, le),
452 _ => unreachable!(),
453 }
454}
455
456fn decode_int64_delta_values(raw: &[u8], ty: TypeId, n: usize, le: bool) -> Result<Vec<Value>> {
459 if !matches!(ty, TypeId::Int64 | TypeId::TimestampNanos) {
460 return Err(MongrelError::InvalidArgument(format!(
461 "delta page not valid for {ty:?}"
462 )));
463 }
464 let (validity, p) = split_validity(raw)?;
465 let deltas = if le {
466 take_i64_le(p, n)?
467 } else {
468 take_i64_be(p, n)?
469 };
470 let data = delta_prefix_sum_i64(&deltas);
471 let mut out = Vec::with_capacity(n);
472 for (i, &v) in data.iter().enumerate() {
473 let non_null = (validity.get(i / 8).copied().unwrap_or(0) >> (i % 8)) & 1 == 1;
474 out.push(if non_null {
475 Value::Int64(v)
476 } else {
477 Value::Null
478 });
479 }
480 Ok(out)
481}
482
483fn dict_encode_bytes(values: &[Value]) -> Vec<u8> {
487 let validity = validity_bitmap(values);
488 let mut table: Vec<Vec<u8>> = Vec::new();
489 let mut index_of: std::collections::HashMap<&[u8], u32> = std::collections::HashMap::new();
490 let mut indices: Vec<u32> = Vec::with_capacity(values.len());
491 for v in values {
492 let idx = match v {
493 Value::Bytes(b) => {
494 if let Some(&i) = index_of.get(b.as_slice()) {
495 i
496 } else {
497 let i = table.len() as u32;
498 index_of.insert(b.as_slice(), i);
499 table.push(b.clone());
500 i
501 }
502 }
503 _ => 0,
504 };
505 indices.push(idx);
506 }
507
508 let mut out = Vec::new();
509 out.extend_from_slice(&(validity.len() as u32).to_be_bytes());
510 out.extend_from_slice(&validity);
511 out.extend_from_slice(&(indices.len() as u32).to_be_bytes());
512 for i in &indices {
513 out.extend_from_slice(&i.to_be_bytes());
514 }
515 out.extend_from_slice(&(table.len() as u32).to_be_bytes());
516 for entry in &table {
517 out.extend_from_slice(&(entry.len() as u32).to_be_bytes());
518 out.extend_from_slice(entry);
519 }
520 out
521}
522
523fn dict_decode_bytes(data: &[u8], n: usize) -> Result<Vec<Value>> {
524 let mut cur = 0usize;
525 let vlen = read_u32_be(data, &mut cur)? as usize;
526 let validity = checked_slice(data, &mut cur, vlen)?;
527 let index_count = read_u32_be(data, &mut cur)? as usize;
528 let mut indices = Vec::with_capacity(index_count.min(n));
529 for _ in 0..index_count {
530 indices.push(read_u32_be(data, &mut cur)?);
531 }
532 let table_count = read_u32_be(data, &mut cur)? as usize;
533 let mut table: Vec<Vec<u8>> = Vec::with_capacity(table_count);
534 for _ in 0..table_count {
535 let len = read_u32_be(data, &mut cur)? as usize;
536 table.push(checked_slice(data, &mut cur, len)?.to_vec());
537 }
538
539 let mut out = Vec::with_capacity(n);
540 for (i, &idx) in indices.iter().enumerate().take(n) {
541 let non_null = (validity.get(i / 8).copied().unwrap_or(0) >> (i % 8)) & 1 == 1;
542 if !non_null {
543 out.push(Value::Null);
544 } else {
545 let entry = table
546 .get(idx as usize)
547 .cloned()
548 .ok_or_else(|| MongrelError::InvalidArgument("dict index out of range".into()))?;
549 out.push(Value::Bytes(entry));
550 }
551 }
552 Ok(out)
553}
554
555fn read_u32_be(data: &[u8], cur: &mut usize) -> Result<u32> {
556 if *cur + 4 > data.len() {
557 return Err(MongrelError::InvalidArgument(
558 "dict payload truncated".into(),
559 ));
560 }
561 let v = u32::from_be_bytes([data[*cur], data[*cur + 1], data[*cur + 2], data[*cur + 3]]);
562 *cur += 4;
563 Ok(v)
564}
565
566fn checked_slice<'a>(data: &'a [u8], cur: &mut usize, len: usize) -> Result<&'a [u8]> {
569 if *cur + len > data.len() {
570 return Err(MongrelError::InvalidArgument(
571 "dict payload truncated".into(),
572 ));
573 }
574 let s = &data[*cur..*cur + len];
575 *cur += len;
576 Ok(s)
577}
578
579#[cfg(test)]
580mod compressed_tests {
581 use super::*;
582
583 #[test]
584 fn zstd_plain_round_trip_int64() {
585 let vals: Vec<Value> = (0..1000).map(Value::Int64).collect();
586 let page = encode_page(TypeId::Int64, &vals, Encoding::Zstd).unwrap();
587 assert!(
588 page.len() < vals.len() * 8,
589 "zstd must shrink sequential ints"
590 );
591 let back = decode_page(TypeId::Int64, &page, vals.len()).unwrap();
592 assert_eq!(back, vals);
593 }
594
595 #[test]
596 fn dictionary_round_trip_low_card_bytes() {
597 let palette: &[&[u8]] = &[b"red", b"green", b"blue", b"red"];
598 let vals: Vec<Value> = (0..500)
599 .map(|i| Value::Bytes(palette[i % palette.len()].to_vec()))
600 .collect();
601 let page = encode_page(TypeId::Bytes, &vals, Encoding::Dictionary).unwrap();
602 assert!(
603 page.len() < 100,
604 "4 distinct strings over 500 rows must compress to a tiny page, got {}",
605 page.len()
606 );
607 let back = decode_page(TypeId::Bytes, &page, vals.len()).unwrap();
608 assert_eq!(back, vals);
609 }
610
611 #[test]
612 fn plain_page_still_round_trips() {
613 let vals = vec![Value::Int64(1), Value::Null, Value::Int64(9)];
614 let page = encode_page(TypeId::Int64, &vals, Encoding::Plain).unwrap();
615 assert_eq!(page[0], ALGO_PLAIN);
616 assert_eq!(decode_page(TypeId::Int64, &page, 3).unwrap(), vals);
617 }
618}
619
620#[cfg(test)]
621mod tests {
622 use super::*;
623
624 #[test]
625 fn round_trips_int64_with_nulls() {
626 let vals = vec![
627 Value::Int64(1),
628 Value::Null,
629 Value::Int64(-5),
630 Value::Int64(1 << 40),
631 ];
632 let page = encode_column(TypeId::Int64, &vals).unwrap();
633 let back = decode_column(TypeId::Int64, &page, vals.len(), false).unwrap();
634 assert_eq!(back, vals);
635 }
636
637 #[test]
638 fn round_trips_bytes() {
639 let vals = vec![
640 Value::Bytes(b"hello".to_vec()),
641 Value::Null,
642 Value::Bytes(b"".to_vec()),
643 Value::Bytes(b"wide \x00 byte".to_vec()),
644 ];
645 let page = encode_column(TypeId::Bytes, &vals).unwrap();
646 let back = decode_column(TypeId::Bytes, &page, vals.len(), false).unwrap();
647 assert_eq!(back, vals);
648 }
649
650 #[test]
651 fn round_trips_embedding() {
652 let vals = vec![
653 Value::Embedding(vec![1.0, -2.5, 3.0]),
654 Value::Null,
655 Value::Embedding(vec![0.0; 3]),
656 ];
657 let page = encode_column(TypeId::Embedding { dim: 3 }, &vals).unwrap();
658 let back = decode_column(TypeId::Embedding { dim: 3 }, &page, vals.len(), false).unwrap();
659 assert_eq!(back, vals);
660 }
661
662 #[test]
663 fn round_trips_bool() {
664 let vals = vec![
665 Value::Bool(true),
666 Value::Bool(false),
667 Value::Null,
668 Value::Bool(true),
669 ];
670 let page = encode_column(TypeId::Bool, &vals).unwrap();
671 assert_eq!(
672 decode_column(TypeId::Bool, &page, vals.len(), false).unwrap(),
673 vals
674 );
675 }
676}
677
678#[derive(Debug, Clone, Serialize, Deserialize)]
687pub enum NativeColumn {
688 Int64 {
689 data: Vec<i64>,
690 validity: Vec<u8>,
691 },
692 Float64 {
693 data: Vec<f64>,
694 validity: Vec<u8>,
695 },
696 Bool {
698 data: Vec<u8>,
699 validity: Vec<u8>,
700 },
701 Bytes {
703 offsets: Vec<u32>,
704 values: Vec<u8>,
705 validity: Vec<u8>,
706 },
707}
708
709impl NativeColumn {
710 pub fn len(&self) -> usize {
711 match self {
712 NativeColumn::Int64 { data, .. } => data.len(),
713 NativeColumn::Float64 { data, .. } => data.len(),
714 NativeColumn::Bool { data, .. } => data.len(),
715 NativeColumn::Bytes { offsets, .. } => offsets.len().saturating_sub(1),
716 }
717 }
718
719 pub fn is_empty(&self) -> bool {
720 self.len() == 0
721 }
722
723 pub fn validity(&self) -> &[u8] {
726 match self {
727 NativeColumn::Int64 { validity, .. }
728 | NativeColumn::Float64 { validity, .. }
729 | NativeColumn::Bool { validity, .. }
730 | NativeColumn::Bytes { validity, .. } => validity,
731 }
732 }
733
734 pub fn validate(&self) -> bool {
738 match self {
739 NativeColumn::Int64 { data, validity } => {
740 validity.len() == data.len().div_ceil(8) || validity.is_empty()
741 }
742 NativeColumn::Float64 { data, validity } => {
743 validity.len() == data.len().div_ceil(8) || validity.is_empty()
744 }
745 NativeColumn::Bool { data, validity } => {
746 validity.len() == data.len().div_ceil(8) || validity.is_empty()
747 }
748 NativeColumn::Bytes {
749 offsets,
750 values,
751 validity,
752 } => {
753 let n = offsets.len().saturating_sub(1);
754 (validity.len() == n.div_ceil(8) || validity.is_empty())
755 && offsets
756 .last()
757 .map(|&last| (last as usize) <= values.len())
758 .unwrap_or(true)
759 }
760 }
761 }
762
763 pub fn null_count(&self, n: usize) -> usize {
766 if n == 0 {
767 return 0;
768 }
769 let validity = match self {
770 NativeColumn::Int64 { validity, .. }
771 | NativeColumn::Float64 { validity, .. }
772 | NativeColumn::Bool { validity, .. }
773 | NativeColumn::Bytes { validity, .. } => validity,
774 };
775 if validity.is_empty() {
776 return 0;
777 }
778 (0..n).filter(|&i| !validity_bit(validity, i)).count()
779 }
780
781 pub fn approx_bytes(&self) -> u64 {
783 match self {
784 NativeColumn::Int64 { data, validity } => {
785 (data.len() as u64) * 8 + validity.len() as u64
786 }
787 NativeColumn::Float64 { data, validity } => {
788 (data.len() as u64) * 8 + validity.len() as u64
789 }
790 NativeColumn::Bool { data, validity } => data.len() as u64 + validity.len() as u64,
791 NativeColumn::Bytes {
792 offsets,
793 values,
794 validity,
795 } => values.len() as u64 + (offsets.len() as u64) * 4 + validity.len() as u64,
796 }
797 }
798
799 pub fn int64_sequence(start: i64, n: usize) -> Self {
801 NativeColumn::Int64 {
802 data: (0..n).map(|i| start + i as i64).collect(),
803 validity: full_validity(n),
804 }
805 }
806
807 pub fn int64_constant(value: i64, n: usize) -> Self {
809 NativeColumn::Int64 {
810 data: vec![value; n],
811 validity: full_validity(n),
812 }
813 }
814
815 pub fn bool_constant(value: bool, n: usize) -> Self {
817 NativeColumn::Bool {
818 data: vec![if value { 1 } else { 0 }; n],
819 validity: full_validity(n),
820 }
821 }
822
823 pub fn gather(&self, indices: &[usize]) -> NativeColumn {
826 let bit = |v: &[u8], i: usize| (v.get(i / 8).copied().unwrap_or(0) >> (i % 8)) & 1 == 1;
827 match self {
828 NativeColumn::Int64 { data, validity } => NativeColumn::Int64 {
829 data: indices.iter().map(|&i| data[i]).collect(),
830 validity: validity_bitmap_from(indices.iter().map(|&i| bit(validity, i))),
831 },
832 NativeColumn::Float64 {
833 data: fdata,
834 validity: fval,
835 } => NativeColumn::Float64 {
836 data: indices.iter().map(|&i| fdata[i]).collect(),
837 validity: validity_bitmap_from(indices.iter().map(|&i| bit(fval, i))),
838 },
839 NativeColumn::Bool {
840 data: bdata,
841 validity: bval,
842 } => NativeColumn::Bool {
843 data: indices.iter().map(|&i| bdata[i]).collect(),
844 validity: validity_bitmap_from(indices.iter().map(|&i| bit(bval, i))),
845 },
846 NativeColumn::Bytes {
847 offsets,
848 values,
849 validity,
850 } => {
851 let mut out_offsets = Vec::with_capacity(indices.len() + 1);
852 let mut out_values = Vec::new();
853 out_offsets.push(0);
854 for &i in indices {
855 let lo = offsets[i] as usize;
856 let hi = offsets[i + 1] as usize;
857 out_values.extend_from_slice(&values[lo..hi]);
858 out_offsets.push(out_values.len() as u32);
859 }
860 NativeColumn::Bytes {
861 offsets: out_offsets,
862 values: out_values,
863 validity: validity_bitmap_from(indices.iter().map(|&i| bit(validity, i))),
864 }
865 }
866 }
867 }
868
869 pub fn value_at(&self, idx: usize) -> Option<Value> {
874 match self {
875 NativeColumn::Int64 { data, validity } => {
876 if !validity_bit(validity, idx) {
877 return None;
878 }
879 data.get(idx).copied().map(Value::Int64)
880 }
881 NativeColumn::Float64 { data, validity } => {
882 if !validity_bit(validity, idx) {
883 return None;
884 }
885 data.get(idx).copied().map(Value::Float64)
886 }
887 NativeColumn::Bool { data, validity } => {
888 if !validity_bit(validity, idx) {
889 return None;
890 }
891 data.get(idx).copied().map(|b| Value::Bool(b != 0))
892 }
893 NativeColumn::Bytes {
894 offsets,
895 values,
896 validity,
897 } => {
898 if !validity_bit(validity, idx) {
899 return None;
900 }
901 if idx + 1 >= offsets.len() {
902 return None;
903 }
904 let lo = offsets[idx] as usize;
905 let hi = offsets[idx + 1] as usize;
906 Some(Value::Bytes(values[lo..hi].to_vec()))
907 }
908 }
909 }
910
911 pub fn slice_range(&self, start: usize, end: usize) -> NativeColumn {
914 let mk_validity = |v: &[u8]| -> Vec<bool> {
915 (0..(end - start))
916 .map(|i| validity_bit(v, start + i))
917 .collect()
918 };
919 match self {
920 NativeColumn::Int64 { data, validity } => NativeColumn::Int64 {
921 data: data[start..end].to_vec(),
922 validity: validity_bitmap_from(mk_validity(validity)),
923 },
924 NativeColumn::Float64 { data, validity } => NativeColumn::Float64 {
925 data: data[start..end].to_vec(),
926 validity: validity_bitmap_from(mk_validity(validity)),
927 },
928 NativeColumn::Bool { data, validity } => NativeColumn::Bool {
929 data: data[start..end].to_vec(),
930 validity: validity_bitmap_from(mk_validity(validity)),
931 },
932 NativeColumn::Bytes {
933 offsets,
934 values,
935 validity,
936 } => {
937 let lo = offsets[start] as usize;
938 let hi = offsets[end] as usize;
939 let new_offsets: Vec<u32> = offsets[start..=end]
940 .iter()
941 .map(|o| *o - offsets[start])
942 .collect();
943 NativeColumn::Bytes {
944 offsets: new_offsets,
945 values: values[lo..hi].to_vec(),
946 validity: validity_bitmap_from(mk_validity(validity)),
947 }
948 }
949 }
950 }
951
952 pub fn concat(parts: &[NativeColumn]) -> NativeColumn {
955 match parts.first() {
956 Some(NativeColumn::Int64 { .. }) => {
957 let mut data = Vec::new();
958 let mut non_null: Vec<bool> = Vec::new();
959 for p in parts {
960 if let NativeColumn::Int64 { data: d, validity } = p {
961 data.extend_from_slice(d);
962 non_null.extend((0..d.len()).map(|i| validity_bit(validity, i)));
963 }
964 }
965 NativeColumn::Int64 {
966 data,
967 validity: validity_bitmap_from(non_null),
968 }
969 }
970 Some(NativeColumn::Float64 { .. }) => {
971 let mut data = Vec::new();
972 let mut non_null: Vec<bool> = Vec::new();
973 for p in parts {
974 if let NativeColumn::Float64 { data: d, validity } = p {
975 data.extend_from_slice(d);
976 non_null.extend((0..d.len()).map(|i| validity_bit(validity, i)));
977 }
978 }
979 NativeColumn::Float64 {
980 data,
981 validity: validity_bitmap_from(non_null),
982 }
983 }
984 Some(NativeColumn::Bool { .. }) => {
985 let mut data = Vec::new();
986 let mut non_null: Vec<bool> = Vec::new();
987 for p in parts {
988 if let NativeColumn::Bool { data: d, validity } = p {
989 data.extend_from_slice(d);
990 non_null.extend((0..d.len()).map(|i| validity_bit(validity, i)));
991 }
992 }
993 NativeColumn::Bool {
994 data,
995 validity: validity_bitmap_from(non_null),
996 }
997 }
998 Some(NativeColumn::Bytes { .. }) => {
999 let mut offsets: Vec<u32> = vec![0];
1000 let mut values = Vec::new();
1001 let mut non_null: Vec<bool> = Vec::new();
1002 for p in parts {
1003 if let NativeColumn::Bytes {
1004 offsets: off,
1005 values: val,
1006 validity,
1007 } = p
1008 {
1009 for w in off.windows(2) {
1010 values.extend_from_slice(&val[w[0] as usize..w[1] as usize]);
1011 offsets.push(values.len() as u32);
1012 }
1013 non_null.extend((0..off.len() - 1).map(|i| validity_bit(validity, i)));
1014 }
1015 }
1016 NativeColumn::Bytes {
1017 offsets,
1018 values,
1019 validity: validity_bitmap_from(non_null),
1020 }
1021 }
1022 None => NativeColumn::Bytes {
1023 offsets: vec![0],
1024 values: Vec::new(),
1025 validity: Vec::new(),
1026 },
1027 }
1028 }
1029}
1030
1031fn full_validity(n: usize) -> Vec<u8> {
1032 validity_bitmap_from(std::iter::repeat(true).take(n))
1033}
1034
1035pub fn null_native(ty: TypeId, n: usize) -> NativeColumn {
1038 let validity = vec![0u8; n.div_ceil(8)];
1039 match ty {
1040 TypeId::Int64 | TypeId::TimestampNanos => NativeColumn::Int64 {
1041 data: vec![0; n],
1042 validity,
1043 },
1044 TypeId::Float64 => NativeColumn::Float64 {
1045 data: vec![0.0; n],
1046 validity,
1047 },
1048 TypeId::Bool => NativeColumn::Bool {
1049 data: vec![0; n],
1050 validity,
1051 },
1052 _ => NativeColumn::Bytes {
1053 offsets: vec![0u32; n + 1],
1054 values: Vec::new(),
1055 validity,
1056 },
1057 }
1058}
1059
1060#[inline]
1062pub fn validity_bit(validity: &[u8], i: usize) -> bool {
1063 (validity.get(i / 8).copied().unwrap_or(0) >> (i % 8)) & 1 == 1
1064}
1065
1066pub fn all_non_null(validity: &[u8], n: usize) -> bool {
1069 if n == 0 {
1070 return true;
1071 }
1072 let full = n / 8;
1073 if !validity[..full].iter().all(|&b| b == 0xFF) {
1074 return false;
1075 }
1076 if n % 8 != 0 {
1077 let mask = (1u8 << (n % 8)) - 1;
1078 (validity.get(full).copied().unwrap_or(0) & mask) == mask
1079 } else {
1080 true
1081 }
1082}
1083
1084pub fn native_min_max(ty: TypeId, col: &NativeColumn) -> (Option<Vec<u8>>, Option<Vec<u8>>, u64) {
1089 let _ = ty;
1090 match col {
1091 NativeColumn::Int64 { data, validity } => {
1092 let (mut mn, mut mx, mut nulls) = (None::<i64>, None::<i64>, 0u64);
1093 for (i, v) in data.iter().enumerate() {
1094 if !validity_bit(validity, i) {
1095 nulls += 1;
1096 continue;
1097 }
1098 mn = Some(mn.map_or(*v, |m| m.min(*v)));
1099 mx = Some(mx.map_or(*v, |m| m.max(*v)));
1100 }
1101 (
1102 mn.map(|v| v.to_be_bytes().to_vec()),
1103 mx.map(|v| v.to_be_bytes().to_vec()),
1104 nulls,
1105 )
1106 }
1107 NativeColumn::Float64 { data, validity } => {
1108 let (mut mn, mut mx, mut nulls) = (None::<f64>, None::<f64>, 0u64);
1109 for (i, v) in data.iter().enumerate() {
1110 if !validity_bit(validity, i) || v.is_nan() {
1111 nulls += 1;
1112 continue;
1113 }
1114 mn = Some(mn.map_or(*v, |m| m.min(*v)));
1115 mx = Some(mx.map_or(*v, |m| m.max(*v)));
1116 }
1117 (
1118 mn.map(|v| v.to_bits().to_be_bytes().to_vec()),
1119 mx.map(|v| v.to_bits().to_be_bytes().to_vec()),
1120 nulls,
1121 )
1122 }
1123 NativeColumn::Bool { data, validity } => {
1124 let (mut any_t, mut any_f, mut nulls) = (false, false, 0u64);
1125 for (i, v) in data.iter().enumerate() {
1126 if !validity_bit(validity, i) {
1127 nulls += 1;
1128 continue;
1129 }
1130 if *v != 0 {
1131 any_t = true;
1132 } else {
1133 any_f = true;
1134 }
1135 }
1136 let min = if any_f || any_t {
1137 Some(vec![if any_f { 0 } else { 1 }])
1138 } else {
1139 None
1140 };
1141 let max = if any_t || any_f {
1142 Some(vec![if any_t { 1 } else { 0 }])
1143 } else {
1144 None
1145 };
1146 (min, max, nulls)
1147 }
1148 NativeColumn::Bytes {
1149 offsets,
1150 values,
1151 validity,
1152 } => {
1153 let mut mn: Option<&[u8]> = None;
1154 let mut mx: Option<&[u8]> = None;
1155 let mut nulls = 0u64;
1156 for i in 0..offsets.len().saturating_sub(1) {
1157 if !validity_bit(validity, i) {
1158 nulls += 1;
1159 continue;
1160 }
1161 let s = &values[offsets[i] as usize..offsets[i + 1] as usize];
1162 mn = Some(match mn {
1163 None => s,
1164 Some(m) if s < m => s,
1165 Some(m) => m,
1166 });
1167 mx = Some(match mx {
1168 None => s,
1169 Some(m) if s > m => s,
1170 Some(m) => m,
1171 });
1172 }
1173 (mn.map(|s| s.to_vec()), mx.map(|s| s.to_vec()), nulls)
1174 }
1175 }
1176}
1177
1178pub fn page_stat_for(
1182 ty: TypeId,
1183 col: &NativeColumn,
1184 first_row_id: u64,
1185 last_row_id: u64,
1186) -> crate::page::PageStat {
1187 let (min, max, null_count) = native_min_max(ty, col);
1188 crate::page::PageStat {
1189 first_row_id,
1190 last_row_id,
1191 null_count,
1192 row_count: col.len() as u32,
1193 min,
1194 max,
1195 offset: 0,
1196 compressed_len: 0,
1197 uncompressed_len: 0,
1198 }
1199}
1200
1201pub fn encode_key_native(_ty: TypeId, col: &NativeColumn, i: usize) -> Option<Vec<u8>> {
1208 match col {
1209 NativeColumn::Int64 { data, validity } if validity_bit(validity, i) => {
1210 Some(data[i].to_be_bytes().to_vec())
1211 }
1212 NativeColumn::Float64 { data, validity } if validity_bit(validity, i) => {
1213 Some(data[i].to_bits().to_be_bytes().to_vec())
1214 }
1215 NativeColumn::Bool { data, validity } if validity_bit(validity, i) => Some(vec![data[i]]),
1216 NativeColumn::Bytes {
1217 offsets,
1218 values,
1219 validity,
1220 } if validity_bit(validity, i) => {
1221 let lo = offsets[i] as usize;
1222 let hi = offsets[i + 1] as usize;
1223 Some(values[lo..hi].to_vec())
1224 }
1225 _ => None,
1226 }
1227}
1228
1229pub fn native_bytes_at(col: &NativeColumn, i: usize) -> Option<&[u8]> {
1233 match col {
1234 NativeColumn::Bytes {
1235 offsets,
1236 values,
1237 validity,
1238 } if validity_bit(validity, i) => {
1239 let lo = offsets[i] as usize;
1240 let hi = offsets[i + 1] as usize;
1241 Some(&values[lo..hi])
1242 }
1243 _ => None,
1244 }
1245}
1246
1247pub fn values_to_native(ty: TypeId, values: &[Value]) -> NativeColumn {
1250 let n = values.len();
1251 let mut non_null = vec![false; n];
1252 match ty {
1253 TypeId::Int64 | TypeId::TimestampNanos => {
1254 let mut data = Vec::with_capacity(n);
1255 for (i, v) in values.iter().enumerate() {
1256 match v {
1257 Value::Int64(x) => {
1258 non_null[i] = true;
1259 data.push(*x);
1260 }
1261 _ => data.push(0),
1262 }
1263 }
1264 NativeColumn::Int64 {
1265 data,
1266 validity: validity_bitmap_from(non_null),
1267 }
1268 }
1269 TypeId::Float64 => {
1270 let mut data = Vec::with_capacity(n);
1271 for (i, v) in values.iter().enumerate() {
1272 match v {
1273 Value::Float64(x) => {
1274 non_null[i] = true;
1275 data.push(*x);
1276 }
1277 _ => data.push(0.0),
1278 }
1279 }
1280 NativeColumn::Float64 {
1281 data,
1282 validity: validity_bitmap_from(non_null),
1283 }
1284 }
1285 TypeId::Bool => {
1286 let mut data = Vec::with_capacity(n);
1287 for (i, v) in values.iter().enumerate() {
1288 match v {
1289 Value::Bool(x) => {
1290 non_null[i] = true;
1291 data.push(if *x { 1 } else { 0 });
1292 }
1293 _ => data.push(0),
1294 }
1295 }
1296 NativeColumn::Bool {
1297 data,
1298 validity: validity_bitmap_from(non_null),
1299 }
1300 }
1301 _ => {
1302 let mut offsets = Vec::with_capacity(n + 1);
1303 let mut vals = Vec::new();
1304 offsets.push(0u32);
1305 for (i, v) in values.iter().enumerate() {
1306 if let Value::Bytes(b) = v {
1307 non_null[i] = true;
1308 vals.extend_from_slice(b);
1309 }
1310 offsets.push(vals.len() as u32);
1311 }
1312 NativeColumn::Bytes {
1313 offsets,
1314 values: vals,
1315 validity: validity_bitmap_from(non_null),
1316 }
1317 }
1318 }
1319}
1320
1321fn validity_bitmap_from(non_null: impl IntoIterator<Item = bool>) -> Vec<u8> {
1322 let bits: Vec<bool> = non_null.into_iter().collect();
1323 let n = bits.len();
1324 let mut out = vec![0u8; n.div_ceil(8)];
1325 for (i, &b) in bits.iter().enumerate() {
1326 if b {
1327 out[i / 8] |= 1 << (i % 8);
1328 }
1329 }
1330 out
1331}
1332
1333pub fn encode_page_native(
1340 ty: TypeId,
1341 col: &NativeColumn,
1342 encoding: Encoding,
1343 compress: Compress,
1344 le: bool,
1345) -> Result<Vec<u8>> {
1346 let raw = matches!(compress, Compress::Plain) || matches!(encoding, Encoding::Plain);
1347 Ok(match (ty, col) {
1348 (TypeId::Int64 | TypeId::TimestampNanos, NativeColumn::Int64 { data, validity }) => {
1349 if matches!(encoding, Encoding::Delta) && !raw {
1350 let mut payload = Vec::with_capacity(4 + validity.len() + data.len() * 8);
1351 payload.extend_from_slice(&(validity.len() as u32).to_be_bytes());
1352 payload.extend_from_slice(validity);
1353 let mut deltas = Vec::with_capacity(data.len());
1354 let mut prev = 0i64;
1355 for v in data {
1356 deltas.push(v - prev);
1357 prev = *v;
1358 }
1359 if le {
1360 append_i64_le(&mut payload, &deltas);
1361 } else {
1362 append_i64_be(&mut payload, &deltas);
1363 }
1364 compress_delta_payload(&payload, compress, le)?
1365 } else {
1366 native_plain_page(validity, compress, raw, le, |p| {
1367 if le {
1368 append_i64_le(p, data);
1369 } else {
1370 append_i64_be(p, data);
1371 }
1372 })
1373 }
1374 }
1375 (
1376 TypeId::Float64,
1377 NativeColumn::Float64 {
1378 data: fdata,
1379 validity,
1380 },
1381 ) => native_plain_page(validity, compress, raw, le, |p| {
1382 let bits: &[u64] = bytemuck::cast_slice::<f64, u64>(fdata);
1383 if le {
1384 append_u64_le(p, bits);
1385 } else {
1386 append_u64_be(p, bits);
1387 }
1388 }),
1389 (
1390 TypeId::Bool,
1391 NativeColumn::Bool {
1392 data: bdata,
1393 validity,
1394 },
1395 ) => native_plain_page(validity, compress, raw, le, |p| p.extend_from_slice(bdata)),
1396 (
1397 TypeId::Bytes,
1398 NativeColumn::Bytes {
1399 offsets,
1400 values,
1401 validity,
1402 },
1403 ) => {
1404 if matches!(encoding, Encoding::Dictionary) && !raw {
1405 let dict = dict_encode_bytes_native(offsets, values, validity);
1406 compress_dict_payload(&dict, compress, le)?
1407 } else {
1408 native_plain_page(validity, compress, raw, le, |p| {
1409 let offs: Vec<u64> = offsets.iter().map(|o| *o as u64).collect();
1410 if le {
1411 append_u64_le(p, &offs);
1412 } else {
1413 append_u64_be(p, &offs);
1414 }
1415 p.extend_from_slice(values);
1416 })
1417 }
1418 }
1419 _ => {
1420 return Err(MongrelError::InvalidArgument(format!(
1421 "encode_page_native: unsupported (ty={ty:?})"
1422 )))
1423 }
1424 })
1425}
1426
1427fn compress_delta_payload(payload: &[u8], compress: Compress, le: bool) -> Result<Vec<u8>> {
1431 Ok(match compress {
1432 Compress::Plain => {
1433 let mut out = vec![algo_with_le(ALGO_PLAIN, le)];
1434 out.extend_from_slice(payload);
1435 out
1436 }
1437 Compress::Zstd(level) => {
1438 let mut out = vec![algo_with_le(ALGO_ZSTD_DELTA, le)];
1439 out.extend(zstd_compress_level(payload, level)?);
1440 out
1441 }
1442 Compress::Lz4 => {
1443 let mut out = vec![algo_with_le(ALGO_LZ4_DELTA, le)];
1444 out.extend(lz4_compress(payload));
1445 out
1446 }
1447 })
1448}
1449
1450fn compress_dict_payload(payload: &[u8], compress: Compress, _le: bool) -> Result<Vec<u8>> {
1455 Ok(match compress {
1456 Compress::Plain => {
1457 let mut out = vec![ALGO_PLAIN];
1458 out.extend_from_slice(payload);
1459 out
1460 }
1461 Compress::Zstd(level) => {
1462 let mut out = vec![ALGO_ZSTD_DICT];
1463 out.extend(zstd_compress_level(payload, level)?);
1464 out
1465 }
1466 Compress::Lz4 => {
1467 let mut out = vec![ALGO_LZ4_DICT];
1468 out.extend(lz4_compress(payload));
1469 out
1470 }
1471 })
1472}
1473
1474fn native_plain_page(
1479 validity: &[u8],
1480 compress: Compress,
1481 raw: bool,
1482 le: bool,
1483 fill_payload: impl FnOnce(&mut Vec<u8>),
1484) -> Vec<u8> {
1485 let mut payload = Vec::new();
1486 payload.extend_from_slice(&(validity.len() as u32).to_be_bytes());
1487 payload.extend_from_slice(validity);
1488 fill_payload(&mut payload);
1489 if raw {
1490 let mut out = vec![algo_with_le(ALGO_PLAIN, le)];
1491 out.extend(payload);
1492 out
1493 } else {
1494 match compress {
1495 Compress::Zstd(level) => {
1496 let mut out = vec![algo_with_le(ALGO_ZSTD_PLAIN, le)];
1497 out.extend(zstd_compress_level(&payload, level).expect("zstd compress"));
1498 out
1499 }
1500 Compress::Lz4 => {
1501 let mut out = vec![algo_with_le(ALGO_LZ4_PLAIN, le)];
1502 out.extend(lz4_compress(&payload));
1503 out
1504 }
1505 Compress::Plain => {
1506 let mut out = vec![algo_with_le(ALGO_PLAIN, le)];
1507 out.extend(payload);
1508 out
1509 }
1510 }
1511 }
1512}
1513
1514pub fn decode_page_native(ty: TypeId, page: &[u8], n: usize) -> Result<NativeColumn> {
1519 use std::borrow::Cow;
1520 if page.is_empty() {
1521 return Err(MongrelError::InvalidArgument("empty page".into()));
1522 }
1523 let algo = page[0];
1524 let body = &page[1..];
1525 let le = algo & ALGO_LE_FLAG != 0;
1529 let base = algo & !ALGO_LE_FLAG;
1530 let plain_algo = matches!(base, ALGO_PLAIN | ALGO_ZSTD_PLAIN | ALGO_LZ4_PLAIN);
1531 let delta_algo = matches!(base, ALGO_ZSTD_DELTA | ALGO_LZ4_DELTA);
1532 let dict_algo = matches!(base, ALGO_ZSTD_DICT | ALGO_LZ4_DICT);
1533 if !plain_algo && !delta_algo && !dict_algo {
1534 return Err(MongrelError::InvalidArgument(format!(
1535 "decode_page_native: unsupported algo {algo} for ty {ty:?}"
1536 )));
1537 }
1538 let raw: Cow<[u8]> = match base {
1540 ALGO_PLAIN => Cow::Borrowed(body),
1541 ALGO_ZSTD_PLAIN | ALGO_ZSTD_DELTA | ALGO_ZSTD_DICT => {
1542 Cow::Owned(zstd_decompress(body, max_decompressed_bytes(ty, n, base))?)
1543 }
1544 ALGO_LZ4_PLAIN | ALGO_LZ4_DELTA | ALGO_LZ4_DICT => {
1545 Cow::Owned(lz4_decompress(body, max_decompressed_bytes(ty, n, base))?)
1546 }
1547 _ => unreachable!(),
1548 };
1549
1550 if dict_algo {
1553 return if matches!(ty, TypeId::Bytes) {
1554 dict_decode_bytes_native(&raw, n)
1555 } else {
1556 Err(MongrelError::InvalidArgument(format!(
1557 "decode_page_native: dict algo {algo} only valid for Bytes, got {ty:?}"
1558 )))
1559 };
1560 }
1561
1562 let take_i64 = |p: &[u8]| -> Result<Vec<i64>> {
1564 if le {
1565 take_i64_le(p, n)
1566 } else {
1567 take_i64_be(p, n)
1568 }
1569 };
1570 if delta_algo {
1572 if !matches!(ty, TypeId::Int64 | TypeId::TimestampNanos) {
1573 return Err(MongrelError::InvalidArgument(format!(
1574 "decode_page_native: delta algo {algo} only valid for Int64, got {ty:?}"
1575 )));
1576 }
1577 let (validity, p) = split_validity(&raw)?;
1578 let deltas = take_i64(p)?;
1579 let data = delta_prefix_sum_i64(&deltas);
1580 return Ok(NativeColumn::Int64 { data, validity });
1581 }
1582
1583 match ty {
1585 TypeId::Int64 | TypeId::TimestampNanos => {
1586 let (validity, p) = split_validity(&raw)?;
1587 Ok(NativeColumn::Int64 {
1588 data: take_i64(p)?,
1589 validity,
1590 })
1591 }
1592 TypeId::Float64 => {
1593 let (validity, p) = split_validity(&raw)?;
1594 let bits = if le {
1595 take_u64_le(p, n)?
1596 } else {
1597 take_u64_be(p, n)?
1598 };
1599 let data: Vec<f64> = bits.into_iter().map(f64::from_bits).collect();
1600 Ok(NativeColumn::Float64 { data, validity })
1601 }
1602 TypeId::Bool => {
1603 let (validity, p) = split_validity(&raw)?;
1604 if p.len() < n {
1605 return Err(MongrelError::InvalidArgument(
1606 "bool payload truncated".into(),
1607 ));
1608 }
1609 Ok(NativeColumn::Bool {
1610 data: p[..n].to_vec(),
1611 validity,
1612 })
1613 }
1614 TypeId::Bytes => decode_bytes_plain_payload(&raw, n, le),
1615 _ => Err(MongrelError::InvalidArgument(format!(
1616 "decode_page_native: unsupported ty {ty:?}"
1617 ))),
1618 }
1619}
1620
1621fn split_validity(raw: &[u8]) -> Result<(Vec<u8>, &[u8])> {
1622 if raw.len() < 4 {
1623 return Err(MongrelError::InvalidArgument("page validity header".into()));
1624 }
1625 let vlen = u32::from_be_bytes([raw[0], raw[1], raw[2], raw[3]]) as usize;
1626 if 4 + vlen > raw.len() {
1627 return Err(MongrelError::InvalidArgument("page validity range".into()));
1628 }
1629 Ok((raw[4..4 + vlen].to_vec(), &raw[4 + vlen..]))
1630}
1631
1632fn decode_bytes_plain_payload(raw: &[u8], n: usize, le: bool) -> Result<NativeColumn> {
1635 let (validity, p) = split_validity(raw)?;
1636 let table = (n + 1) * 8;
1637 if p.len() < table {
1638 return Err(MongrelError::InvalidArgument(
1639 "bytes offsets truncated".into(),
1640 ));
1641 }
1642 let offsets_be: Vec<u64> = if le {
1643 take_u64_le(p, n + 1)?
1644 } else {
1645 take_u64_be(p, n + 1)?
1646 };
1647 let offsets: Vec<u32> = offsets_be.into_iter().map(|o| o as u32).collect();
1648 let values = p[table..].to_vec();
1649 Ok(NativeColumn::Bytes {
1650 offsets,
1651 values,
1652 validity,
1653 })
1654}
1655
1656fn take_i64_be(p: &[u8], n: usize) -> Result<Vec<i64>> {
1657 if p.len() < n * 8 {
1658 return Err(MongrelError::InvalidArgument(
1659 "int64 payload truncated".into(),
1660 ));
1661 }
1662 Ok(take_u64_be(p, n)?.into_iter().map(|u| u as i64).collect())
1663}
1664
1665fn delta_prefix_sum_i64(deltas: &[i64]) -> Vec<i64> {
1673 let mut out = vec![0i64; deltas.len()];
1674 prefix_sum_i64_into(deltas, &mut out);
1675 out
1676}
1677
1678fn prefix_sum_i64_into(deltas: &[i64], out: &mut [i64]) {
1679 debug_assert_eq!(deltas.len(), out.len());
1680 #[cfg(target_arch = "x86_64")]
1681 {
1682 if is_x86_feature_detected!("avx2") && deltas.len() >= 4 {
1683 unsafe {
1686 prefix_sum_avx2(deltas, out);
1687 }
1688 return;
1689 }
1690 }
1691 prefix_sum_scalar(deltas, out);
1692}
1693
1694fn prefix_sum_scalar(deltas: &[i64], out: &mut [i64]) {
1695 let mut acc = 0i64;
1696 for (i, &d) in deltas.iter().enumerate() {
1697 acc = acc.wrapping_add(d);
1698 out[i] = acc;
1699 }
1700}
1701
1702#[cfg(target_arch = "x86_64")]
1703#[target_feature(enable = "avx2")]
1704unsafe fn prefix_sum_avx2(deltas: &[i64], out: &mut [i64]) {
1705 use std::arch::x86_64::*;
1706 let n = deltas.len();
1707 let mut running = _mm256_setzero_si256(); let mut i = 0usize;
1709 while i + 4 <= n {
1710 let mut x = _mm256_loadu_si256(deltas.as_ptr().add(i) as *const __m256i);
1711 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);
1720 running = _mm256_permute4x64_epi64(x, 0xFF); i += 4;
1722 }
1723 let mut acc = if i == 0 { 0 } else { out[i - 1] };
1726 while i < n {
1727 acc = acc.wrapping_add(deltas[i]);
1728 out[i] = acc;
1729 i += 1;
1730 }
1731}
1732
1733fn append_u64_be(out: &mut Vec<u8>, data: &[u64]) {
1735 if cfg!(target_endian = "little") {
1736 let swapped: Vec<u64> = data.iter().map(|v| v.swap_bytes()).collect();
1737 out.extend_from_slice(bytemuck::cast_slice::<u64, u8>(&swapped));
1738 } else {
1739 out.extend_from_slice(bytemuck::cast_slice::<u64, u8>(data));
1740 }
1741}
1742
1743fn append_i64_be(out: &mut Vec<u8>, data: &[i64]) {
1745 if cfg!(target_endian = "little") {
1746 let swapped: Vec<i64> = data.iter().map(|v| v.swap_bytes()).collect();
1747 out.extend_from_slice(bytemuck::cast_slice::<i64, u8>(&swapped));
1748 } else {
1749 out.extend_from_slice(bytemuck::cast_slice::<i64, u8>(data));
1750 }
1751}
1752
1753fn take_u64_be(p: &[u8], n: usize) -> Result<Vec<u64>> {
1757 if p.len() < n * 8 {
1758 return Err(MongrelError::InvalidArgument(
1759 "u64 payload truncated".into(),
1760 ));
1761 }
1762 let bytes = &p[..n * 8];
1763 if let Ok(native) = bytemuck::try_cast_slice::<u8, u64>(bytes) {
1764 Ok(native.iter().map(|v| v.swap_bytes()).collect())
1765 } else {
1766 let mut out = Vec::with_capacity(n);
1767 for chunk in bytes.chunks_exact(8) {
1768 out.push(u64::from_be_bytes(chunk.try_into().unwrap()));
1769 }
1770 Ok(out)
1771 }
1772}
1773
1774fn append_u64_le(out: &mut Vec<u8>, data: &[u64]) {
1778 if cfg!(target_endian = "little") {
1779 out.extend_from_slice(bytemuck::cast_slice::<u64, u8>(data));
1780 } else {
1781 let swapped: Vec<u64> = data.iter().map(|v| v.swap_bytes()).collect();
1782 out.extend_from_slice(bytemuck::cast_slice::<u64, u8>(&swapped));
1783 }
1784}
1785
1786fn append_i64_le(out: &mut Vec<u8>, data: &[i64]) {
1788 if cfg!(target_endian = "little") {
1789 out.extend_from_slice(bytemuck::cast_slice::<i64, u8>(data));
1790 } else {
1791 let swapped: Vec<i64> = data.iter().map(|v| v.swap_bytes()).collect();
1792 out.extend_from_slice(bytemuck::cast_slice::<i64, u8>(&swapped));
1793 }
1794}
1795
1796fn take_u64_le(p: &[u8], n: usize) -> Result<Vec<u64>> {
1802 if p.len() < n * 8 {
1803 return Err(MongrelError::InvalidArgument(
1804 "u64 payload truncated".into(),
1805 ));
1806 }
1807 let bytes = &p[..n * 8];
1808 if cfg!(target_endian = "little") {
1809 if let Ok(native) = bytemuck::try_cast_slice::<u8, u64>(bytes) {
1810 return Ok(native.to_vec());
1812 }
1813 Ok(bytes
1814 .chunks_exact(8)
1815 .map(|c| u64::from_le_bytes(c.try_into().unwrap()))
1816 .collect())
1817 } else {
1818 Ok(bytes
1819 .chunks_exact(8)
1820 .map(|c| u64::from_le_bytes(c.try_into().unwrap()))
1821 .collect())
1822 }
1823}
1824
1825fn take_i64_le(p: &[u8], n: usize) -> Result<Vec<i64>> {
1827 if p.len() < n * 8 {
1828 return Err(MongrelError::InvalidArgument(
1829 "int64 payload truncated".into(),
1830 ));
1831 }
1832 Ok(take_u64_le(p, n)?.into_iter().map(|u| u as i64).collect())
1833}
1834
1835fn dict_encode_bytes_native(offsets: &[u32], values: &[u8], validity: &[u8]) -> Vec<u8> {
1836 let n = offsets.len() - 1;
1837 let mut table: Vec<Vec<u8>> = Vec::new();
1838 let mut index_of: std::collections::HashMap<&[u8], u32> = std::collections::HashMap::new();
1839 let mut indices: Vec<u32> = Vec::with_capacity(n);
1840 for i in 0..n {
1841 let lo = offsets[i] as usize;
1842 let hi = offsets[i + 1] as usize;
1843 let slice = &values[lo..hi];
1844 let idx = if let Some(&idx) = index_of.get(slice) {
1845 idx
1846 } else {
1847 let idx = table.len() as u32;
1848 index_of.insert(slice, idx);
1849 table.push(slice.to_vec());
1850 idx
1851 };
1852 indices.push(idx);
1853 }
1854 let mut out = Vec::new();
1855 out.extend_from_slice(&(validity.len() as u32).to_be_bytes());
1856 out.extend_from_slice(validity);
1857 out.extend_from_slice(&(indices.len() as u32).to_be_bytes());
1858 for i in &indices {
1859 out.extend_from_slice(&i.to_be_bytes());
1860 }
1861 out.extend_from_slice(&(table.len() as u32).to_be_bytes());
1862 for entry in &table {
1863 out.extend_from_slice(&(entry.len() as u32).to_be_bytes());
1864 out.extend_from_slice(entry);
1865 }
1866 out
1867}
1868
1869fn dict_decode_bytes_native(data: &[u8], n: usize) -> Result<NativeColumn> {
1870 let mut cur = 0usize;
1871 let vlen = read_u32_be(data, &mut cur)? as usize;
1872 let validity = checked_slice(data, &mut cur, vlen)?.to_vec();
1873 let index_count = read_u32_be(data, &mut cur)? as usize;
1874 if index_count < n {
1875 return Err(MongrelError::InvalidArgument("dict index_count < n".into()));
1876 }
1877 let mut indices = Vec::with_capacity(index_count.min(n));
1878 for _ in 0..index_count {
1879 indices.push(read_u32_be(data, &mut cur)?);
1880 }
1881 let table_count = read_u32_be(data, &mut cur)? as usize;
1882 let mut table: Vec<(usize, usize)> = Vec::with_capacity(table_count); let mut values = Vec::new();
1884 for _ in 0..table_count {
1885 let len = read_u32_be(data, &mut cur)? as usize;
1886 let chunk = checked_slice(data, &mut cur, len)?;
1887 let start = values.len();
1888 values.extend_from_slice(chunk);
1889 table.push((start, len));
1890 }
1891 let mut merged = Vec::new();
1892 let mut offs = vec![0u32];
1893 for (i, &idx) in indices.iter().enumerate().take(n) {
1894 let non_null = (validity.get(i / 8).copied().unwrap_or(0) >> (i % 8)) & 1 == 1;
1898 if non_null {
1899 let (start, len) = table
1900 .get(idx as usize)
1901 .copied()
1902 .ok_or_else(|| MongrelError::InvalidArgument("dict index out of range".into()))?;
1903 merged.extend_from_slice(&values[start..start + len]);
1904 }
1905 offs.push(merged.len() as u32);
1906 }
1907 Ok(NativeColumn::Bytes {
1908 offsets: offs,
1909 values: merged,
1910 validity,
1911 })
1912}
1913
1914#[cfg(test)]
1915mod native_tests {
1916 use super::*;
1917
1918 #[test]
1919 fn native_int64_plain_round_trip() {
1920 let col = NativeColumn::Int64 {
1921 data: (0..1000).collect(),
1922 validity: full_validity(1000),
1923 };
1924 let page = encode_page_native(
1925 TypeId::Int64,
1926 &col,
1927 Encoding::Zstd,
1928 Compress::Zstd(3),
1929 false,
1930 )
1931 .unwrap();
1932 let back = decode_page_native(TypeId::Int64, &page, 1000).unwrap();
1933 match back {
1934 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
1935 _ => panic!(),
1936 }
1937 }
1938
1939 #[test]
1940 fn native_int64_delta_crushes_sequential() {
1941 let col = NativeColumn::int64_sequence(0, 100_000);
1942 let plain = encode_page_native(
1943 TypeId::Int64,
1944 &col,
1945 Encoding::Zstd,
1946 Compress::Zstd(3),
1947 false,
1948 )
1949 .unwrap();
1950 let delta = encode_page_native(
1951 TypeId::Int64,
1952 &col,
1953 Encoding::Delta,
1954 Compress::Zstd(3),
1955 false,
1956 )
1957 .unwrap();
1958 assert!(
1959 delta.len() < plain.len() / 5,
1960 "delta must crush sequential ints"
1961 );
1962 let back = decode_page_native(TypeId::Int64, &delta, 100_000).unwrap();
1963 match back {
1964 NativeColumn::Int64 { data, .. } => assert_eq!(data.len(), 100_000),
1965 _ => panic!(),
1966 }
1967 }
1968
1969 #[test]
1970 fn native_bytes_dict_round_trip() {
1971 let n = 500;
1972 let mut offsets = vec![0u32];
1973 let mut values = Vec::new();
1974 for i in 0..n {
1975 let s = ["red", "green", "blue"][i % 3];
1976 values.extend_from_slice(s.as_bytes());
1977 offsets.push(values.len() as u32);
1978 }
1979 let col = NativeColumn::Bytes {
1980 offsets,
1981 values,
1982 validity: full_validity(n),
1983 };
1984 let page = encode_page_native(
1985 TypeId::Bytes,
1986 &col,
1987 Encoding::Dictionary,
1988 Compress::Zstd(3),
1989 false,
1990 )
1991 .unwrap();
1992 assert!(page.len() < 100, "dict page tiny, got {}", page.len());
1993 let back = decode_page_native(TypeId::Bytes, &page, n).unwrap();
1994 assert_eq!(back.len(), n);
1995 }
1996
1997 #[test]
1998 fn native_gather_picks_indices() {
1999 let col = NativeColumn::Int64 {
2000 data: vec![10, 20, 30, 40],
2001 validity: full_validity(4),
2002 };
2003 let g = col.gather(&[0, 2, 3]);
2004 match g {
2005 NativeColumn::Int64 { data, .. } => assert_eq!(data, vec![10, 30, 40]),
2006 _ => panic!(),
2007 }
2008 }
2009
2010 #[test]
2013 fn native_plain_no_zstd_round_trips_all_types() {
2014 let i = NativeColumn::Int64 {
2015 data: (0..1000).collect(),
2016 validity: full_validity(1000),
2017 };
2018 let p =
2019 encode_page_native(TypeId::Int64, &i, Encoding::Plain, Compress::Plain, false).unwrap();
2020 assert_eq!(p[0], ALGO_PLAIN, "Int64 plain must be ALGO_PLAIN");
2021 match decode_page_native(TypeId::Int64, &p, 1000).unwrap() {
2022 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2023 _ => panic!(),
2024 }
2025
2026 let f = NativeColumn::Float64 {
2027 data: (0..500).map(|x| x as f64 * 1.5).collect(),
2028 validity: full_validity(500),
2029 };
2030 let p = encode_page_native(TypeId::Float64, &f, Encoding::Plain, Compress::Plain, false)
2031 .unwrap();
2032 assert_eq!(p[0], ALGO_PLAIN);
2033 match decode_page_native(TypeId::Float64, &p, 500).unwrap() {
2034 NativeColumn::Float64 { data, .. } => {
2035 assert_eq!(data, (0..500).map(|x| x as f64 * 1.5).collect::<Vec<_>>())
2036 }
2037 _ => panic!(),
2038 }
2039
2040 let b = NativeColumn::Bool {
2041 data: (0..64).map(|i| (i % 2) as u8).collect(),
2042 validity: full_validity(64),
2043 };
2044 let p =
2045 encode_page_native(TypeId::Bool, &b, Encoding::Plain, Compress::Plain, false).unwrap();
2046 assert_eq!(p[0], ALGO_PLAIN);
2047 match decode_page_native(TypeId::Bool, &p, 64).unwrap() {
2048 NativeColumn::Bool { data, .. } => {
2049 assert_eq!(data, (0..64).map(|i| (i % 2) as u8).collect::<Vec<_>>())
2050 }
2051 _ => panic!(),
2052 }
2053
2054 let mut offsets = vec![0u32];
2055 let mut values = Vec::new();
2056 for i in 0..200u32 {
2057 values.extend_from_slice(format!("v{i}").as_bytes());
2058 offsets.push(values.len() as u32);
2059 }
2060 let s = NativeColumn::Bytes {
2061 offsets,
2062 values,
2063 validity: full_validity(200),
2064 };
2065 let p =
2066 encode_page_native(TypeId::Bytes, &s, Encoding::Plain, Compress::Plain, false).unwrap();
2067 assert_eq!(p[0], ALGO_PLAIN);
2068 match decode_page_native(TypeId::Bytes, &p, 200).unwrap() {
2069 NativeColumn::Bytes {
2070 offsets: o,
2071 values: v,
2072 ..
2073 } => {
2074 assert_eq!(o.len(), 201);
2075 for i in 0..200 {
2076 let lo = o[i] as usize;
2077 let hi = o[i + 1] as usize;
2078 assert_eq!(&v[lo..hi], format!("v{i}").as_bytes());
2079 }
2080 }
2081 _ => panic!(),
2082 }
2083 }
2084
2085 #[test]
2088 fn lz4_pages_round_trip_all_types() {
2089 let i = NativeColumn::int64_sequence(0, 1000);
2091 let p =
2092 encode_page_native(TypeId::Int64, &i, Encoding::Delta, Compress::Lz4, false).unwrap();
2093 assert_eq!(p[0], ALGO_LZ4_DELTA);
2094 match decode_page_native(TypeId::Int64, &p, 1000).unwrap() {
2095 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2096 _ => panic!(),
2097 }
2098 let p =
2100 encode_page_native(TypeId::Int64, &i, Encoding::Zstd, Compress::Lz4, false).unwrap();
2101 assert_eq!(p[0], ALGO_LZ4_PLAIN);
2102 match decode_page_native(TypeId::Int64, &p, 1000).unwrap() {
2103 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2104 _ => panic!(),
2105 }
2106 let f = NativeColumn::Float64 {
2108 data: (0..500).map(|x| x as f64 * 1.5).collect(),
2109 validity: full_validity(500),
2110 };
2111 let p =
2112 encode_page_native(TypeId::Float64, &f, Encoding::Zstd, Compress::Lz4, false).unwrap();
2113 assert_eq!(p[0], ALGO_LZ4_PLAIN);
2114 match decode_page_native(TypeId::Float64, &p, 500).unwrap() {
2115 NativeColumn::Float64 { data, .. } => {
2116 assert_eq!(data, (0..500).map(|x| x as f64 * 1.5).collect::<Vec<_>>())
2117 }
2118 _ => panic!(),
2119 }
2120 let mut offsets = vec![0u32];
2122 let mut values = Vec::new();
2123 for i in 0..300u32 {
2124 values.extend_from_slice(["red", "green", "blue"][(i % 3) as usize].as_bytes());
2125 offsets.push(values.len() as u32);
2126 }
2127 let s = NativeColumn::Bytes {
2128 offsets,
2129 values,
2130 validity: full_validity(300),
2131 };
2132 let p = encode_page_native(
2133 TypeId::Bytes,
2134 &s,
2135 Encoding::Dictionary,
2136 Compress::Lz4,
2137 false,
2138 )
2139 .unwrap();
2140 assert_eq!(p[0], ALGO_LZ4_DICT);
2141 match decode_page_native(TypeId::Bytes, &p, 300).unwrap() {
2142 NativeColumn::Bytes { offsets: o, .. } => assert_eq!(o.len(), 301),
2143 _ => panic!(),
2144 }
2145 }
2146
2147 #[test]
2151 fn le_pages_round_trip_all_types() {
2152 let assert_le = |page: &[u8]| assert_ne!(page[0] & ALGO_LE_FLAG, 0, "LE flag must be set");
2153
2154 let i = NativeColumn::Int64 {
2156 data: (0..1000).collect(),
2157 validity: full_validity(1000),
2158 };
2159 let p =
2160 encode_page_native(TypeId::Int64, &i, Encoding::Plain, Compress::Plain, true).unwrap();
2161 assert_eq!(p[0], ALGO_LE_FLAG, "raw LE Int64 algo = flag only");
2162 match decode_page_native(TypeId::Int64, &p, 1000).unwrap() {
2163 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2164 _ => panic!(),
2165 }
2166
2167 let p =
2169 encode_page_native(TypeId::Int64, &i, Encoding::Zstd, Compress::Zstd(3), true).unwrap();
2170 assert_le(&p);
2171 match decode_page_native(TypeId::Int64, &p, 1000).unwrap() {
2172 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2173 _ => panic!(),
2174 }
2175
2176 let seq = NativeColumn::int64_sequence(0, 1000);
2178 let p =
2179 encode_page_native(TypeId::Int64, &seq, Encoding::Delta, Compress::Lz4, true).unwrap();
2180 assert_eq!(p[0], ALGO_LZ4_DELTA | ALGO_LE_FLAG);
2181 match decode_page_native(TypeId::Int64, &p, 1000).unwrap() {
2182 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2183 _ => panic!(),
2184 }
2185
2186 let f = NativeColumn::Float64 {
2188 data: (0..500).map(|x| x as f64 * 1.5).collect(),
2189 validity: full_validity(500),
2190 };
2191 let p = encode_page_native(TypeId::Float64, &f, Encoding::Zstd, Compress::Zstd(3), true)
2192 .unwrap();
2193 assert_le(&p);
2194 match decode_page_native(TypeId::Float64, &p, 500).unwrap() {
2195 NativeColumn::Float64 { data, .. } => {
2196 assert_eq!(data, (0..500).map(|x| x as f64 * 1.5).collect::<Vec<_>>())
2197 }
2198 _ => panic!(),
2199 }
2200 }
2201
2202 #[test]
2206 fn malicious_lz4_size_prefix_is_rejected() {
2207 let mut evil = vec![ALGO_LZ4_PLAIN];
2209 evil.extend_from_slice(&0xFFFF_FFFFu32.to_le_bytes());
2210 evil.extend_from_slice(b"junk");
2211 let err =
2212 decode_page_native(TypeId::Int64, &evil, 1000).expect_err("must reject oversize lz4");
2213 let msg = format!("{err}");
2214 assert!(msg.contains("exceeds page limit"), "got: {msg}");
2215
2216 let err = decode_page(TypeId::Int64, &evil, 1000).expect_err("must reject oversize lz4");
2218 assert!(format!("{err}").contains("exceeds page limit"));
2219 }
2220
2221 #[test]
2225 fn truncated_dict_payload_does_not_panic() {
2226 let mut body = vec![ALGO_ZSTD_DICT];
2229 body.extend_from_slice(&zstd_compress(&[0u8; 2]).unwrap()[..]);
2230 decode_page(TypeId::Bytes, &body, 4).expect_err("value dict trunc must Err");
2231
2232 let truncated: &[u8] = &[
2234 0x00, 0x00, 0x00, 0x10, ];
2236 dict_decode_bytes_native(truncated, 2).expect_err("native dict trunc must Err");
2237 dict_decode_bytes(truncated, 2).expect_err("value dict trunc must Err");
2238
2239 let mut bad = vec![0u8; 0];
2242 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");
2248 dict_decode_bytes_native(&bad, 1).expect_err("oob index must Err");
2249 }
2250
2251 #[test]
2255 fn delta_prefix_sum_matches_scalar_all_lengths() {
2256 let mut deltas: Vec<i64> = Vec::with_capacity(2000);
2258 let mut s: u64 = 0x9E37_79B9_7F4A_7C15;
2259 for _ in 0..2000 {
2260 s = s
2261 .wrapping_mul(6364136223846793005)
2262 .wrapping_add(1442695040888963407);
2263 deltas.push((s as i32) as i64); }
2265 for &len in &[
2266 0usize, 1, 2, 3, 4, 5, 6, 7, 8, 9, 15, 16, 17, 31, 32, 33, 63, 64, 65, 127, 128, 129,
2267 1000,
2268 ] {
2269 let input = &deltas[..len];
2270 let simd = delta_prefix_sum_i64(input);
2271 let mut ref_out = vec![0i64; len];
2272 prefix_sum_scalar(input, &mut ref_out);
2273 assert_eq!(simd, ref_out, "len {len}: SIMD prefix sum diverged");
2274 }
2275
2276 let mut ids = vec![0i64; 333];
2278 for x in ids.iter_mut().skip(1) {
2279 *x = 1;
2280 }
2281 let got = delta_prefix_sum_i64(&ids);
2282 assert_eq!(got, (0..333).map(|x| x as i64).collect::<Vec<_>>());
2283 }
2284
2285 #[test]
2288 fn delta_page_decodes_with_null_via_vectorized_path() {
2289 let mut validity = full_validity(17);
2290 validity[9 / 8] &= !(1 << (9 % 8)); let col = NativeColumn::Int64 {
2292 data: (0..17).collect(),
2293 validity: validity.clone(),
2294 };
2295 for (name, comp) in [("zstd", Compress::Zstd(3)), ("lz4", Compress::Lz4)] {
2296 let page =
2297 encode_page_native(TypeId::Int64, &col, Encoding::Delta, comp, false).unwrap();
2298 assert!(matches!(page[0], ALGO_ZSTD_DELTA | ALGO_LZ4_DELTA));
2300 let back = decode_page_native(TypeId::Int64, &page, 17).unwrap();
2301 match back {
2302 NativeColumn::Int64 { data, validity: v } => {
2303 assert_eq!(data, (0..17).collect::<Vec<_>>(), "{name} data");
2304 assert_eq!(v, validity, "{name} validity");
2305 }
2306 _ => panic!(),
2307 }
2308 }
2309 }
2310}