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 let full = n / 8;
781 let mut set = validity[..full.min(validity.len())]
782 .iter()
783 .map(|b| b.count_ones() as usize)
784 .sum::<usize>();
785 let tail_bits = n % 8;
786 if tail_bits > 0 {
787 if let Some(&b) = validity.get(full) {
788 set += (b & ((1u8 << tail_bits) - 1)).count_ones() as usize;
789 }
790 }
791 n - set
792 }
793
794 pub fn approx_bytes(&self) -> u64 {
796 match self {
797 NativeColumn::Int64 { data, validity } => {
798 (data.len() as u64) * 8 + validity.len() as u64
799 }
800 NativeColumn::Float64 { data, validity } => {
801 (data.len() as u64) * 8 + validity.len() as u64
802 }
803 NativeColumn::Bool { data, validity } => data.len() as u64 + validity.len() as u64,
804 NativeColumn::Bytes {
805 offsets,
806 values,
807 validity,
808 } => values.len() as u64 + (offsets.len() as u64) * 4 + validity.len() as u64,
809 }
810 }
811
812 pub fn int64_sequence(start: i64, n: usize) -> Self {
814 NativeColumn::Int64 {
815 data: (0..n).map(|i| start + i as i64).collect(),
816 validity: full_validity(n),
817 }
818 }
819
820 pub fn int64_constant(value: i64, n: usize) -> Self {
822 NativeColumn::Int64 {
823 data: vec![value; n],
824 validity: full_validity(n),
825 }
826 }
827
828 pub fn bool_constant(value: bool, n: usize) -> Self {
830 NativeColumn::Bool {
831 data: vec![if value { 1 } else { 0 }; n],
832 validity: full_validity(n),
833 }
834 }
835
836 pub fn gather(&self, indices: &[usize]) -> NativeColumn {
839 let bit = |v: &[u8], i: usize| (v.get(i / 8).copied().unwrap_or(0) >> (i % 8)) & 1 == 1;
840 match self {
841 NativeColumn::Int64 { data, validity } => NativeColumn::Int64 {
842 data: indices.iter().map(|&i| data[i]).collect(),
843 validity: validity_bitmap_from(indices.iter().map(|&i| bit(validity, i))),
844 },
845 NativeColumn::Float64 {
846 data: fdata,
847 validity: fval,
848 } => NativeColumn::Float64 {
849 data: indices.iter().map(|&i| fdata[i]).collect(),
850 validity: validity_bitmap_from(indices.iter().map(|&i| bit(fval, i))),
851 },
852 NativeColumn::Bool {
853 data: bdata,
854 validity: bval,
855 } => NativeColumn::Bool {
856 data: indices.iter().map(|&i| bdata[i]).collect(),
857 validity: validity_bitmap_from(indices.iter().map(|&i| bit(bval, i))),
858 },
859 NativeColumn::Bytes {
860 offsets,
861 values,
862 validity,
863 } => {
864 let mut out_offsets = Vec::with_capacity(indices.len() + 1);
865 let mut out_values = Vec::new();
866 out_offsets.push(0);
867 for &i in indices {
868 let lo = offsets[i] as usize;
869 let hi = offsets[i + 1] as usize;
870 out_values.extend_from_slice(&values[lo..hi]);
871 out_offsets.push(out_values.len() as u32);
872 }
873 NativeColumn::Bytes {
874 offsets: out_offsets,
875 values: out_values,
876 validity: validity_bitmap_from(indices.iter().map(|&i| bit(validity, i))),
877 }
878 }
879 }
880 }
881
882 pub fn value_at(&self, idx: usize) -> Option<Value> {
887 match self {
888 NativeColumn::Int64 { data, validity } => {
889 if !validity_bit(validity, idx) {
890 return None;
891 }
892 data.get(idx).copied().map(Value::Int64)
893 }
894 NativeColumn::Float64 { data, validity } => {
895 if !validity_bit(validity, idx) {
896 return None;
897 }
898 data.get(idx).copied().map(Value::Float64)
899 }
900 NativeColumn::Bool { data, validity } => {
901 if !validity_bit(validity, idx) {
902 return None;
903 }
904 data.get(idx).copied().map(|b| Value::Bool(b != 0))
905 }
906 NativeColumn::Bytes {
907 offsets,
908 values,
909 validity,
910 } => {
911 if !validity_bit(validity, idx) {
912 return None;
913 }
914 if idx + 1 >= offsets.len() {
915 return None;
916 }
917 let lo = offsets[idx] as usize;
918 let hi = offsets[idx + 1] as usize;
919 Some(Value::Bytes(values[lo..hi].to_vec()))
920 }
921 }
922 }
923
924 pub fn slice_range(&self, start: usize, end: usize) -> NativeColumn {
927 let mk_validity = |v: &[u8]| -> Vec<bool> {
928 (0..(end - start))
929 .map(|i| validity_bit(v, start + i))
930 .collect()
931 };
932 match self {
933 NativeColumn::Int64 { data, validity } => NativeColumn::Int64 {
934 data: data[start..end].to_vec(),
935 validity: validity_bitmap_from(mk_validity(validity)),
936 },
937 NativeColumn::Float64 { data, validity } => NativeColumn::Float64 {
938 data: data[start..end].to_vec(),
939 validity: validity_bitmap_from(mk_validity(validity)),
940 },
941 NativeColumn::Bool { data, validity } => NativeColumn::Bool {
942 data: data[start..end].to_vec(),
943 validity: validity_bitmap_from(mk_validity(validity)),
944 },
945 NativeColumn::Bytes {
946 offsets,
947 values,
948 validity,
949 } => {
950 let lo = offsets[start] as usize;
951 let hi = offsets[end] as usize;
952 let new_offsets: Vec<u32> = offsets[start..=end]
953 .iter()
954 .map(|o| *o - offsets[start])
955 .collect();
956 NativeColumn::Bytes {
957 offsets: new_offsets,
958 values: values[lo..hi].to_vec(),
959 validity: validity_bitmap_from(mk_validity(validity)),
960 }
961 }
962 }
963 }
964
965 pub fn concat(parts: &[NativeColumn]) -> NativeColumn {
968 match parts.first() {
969 Some(NativeColumn::Int64 { .. }) => {
970 let mut data = Vec::new();
971 let mut non_null: Vec<bool> = Vec::new();
972 for p in parts {
973 if let NativeColumn::Int64 { data: d, validity } = p {
974 data.extend_from_slice(d);
975 non_null.extend((0..d.len()).map(|i| validity_bit(validity, i)));
976 }
977 }
978 NativeColumn::Int64 {
979 data,
980 validity: validity_bitmap_from(non_null),
981 }
982 }
983 Some(NativeColumn::Float64 { .. }) => {
984 let mut data = Vec::new();
985 let mut non_null: Vec<bool> = Vec::new();
986 for p in parts {
987 if let NativeColumn::Float64 { data: d, validity } = p {
988 data.extend_from_slice(d);
989 non_null.extend((0..d.len()).map(|i| validity_bit(validity, i)));
990 }
991 }
992 NativeColumn::Float64 {
993 data,
994 validity: validity_bitmap_from(non_null),
995 }
996 }
997 Some(NativeColumn::Bool { .. }) => {
998 let mut data = Vec::new();
999 let mut non_null: Vec<bool> = Vec::new();
1000 for p in parts {
1001 if let NativeColumn::Bool { data: d, validity } = p {
1002 data.extend_from_slice(d);
1003 non_null.extend((0..d.len()).map(|i| validity_bit(validity, i)));
1004 }
1005 }
1006 NativeColumn::Bool {
1007 data,
1008 validity: validity_bitmap_from(non_null),
1009 }
1010 }
1011 Some(NativeColumn::Bytes { .. }) => {
1012 let mut offsets: Vec<u32> = vec![0];
1013 let mut values = Vec::new();
1014 let mut non_null: Vec<bool> = Vec::new();
1015 for p in parts {
1016 if let NativeColumn::Bytes {
1017 offsets: off,
1018 values: val,
1019 validity,
1020 } = p
1021 {
1022 for w in off.windows(2) {
1023 values.extend_from_slice(&val[w[0] as usize..w[1] as usize]);
1024 offsets.push(values.len() as u32);
1025 }
1026 non_null.extend((0..off.len() - 1).map(|i| validity_bit(validity, i)));
1027 }
1028 }
1029 NativeColumn::Bytes {
1030 offsets,
1031 values,
1032 validity: validity_bitmap_from(non_null),
1033 }
1034 }
1035 None => NativeColumn::Bytes {
1036 offsets: vec![0],
1037 values: Vec::new(),
1038 validity: Vec::new(),
1039 },
1040 }
1041 }
1042}
1043
1044fn full_validity(n: usize) -> Vec<u8> {
1045 validity_bitmap_from(std::iter::repeat(true).take(n))
1046}
1047
1048pub fn null_native(ty: TypeId, n: usize) -> NativeColumn {
1051 let validity = vec![0u8; n.div_ceil(8)];
1052 match ty {
1053 TypeId::Int64 | TypeId::TimestampNanos => NativeColumn::Int64 {
1054 data: vec![0; n],
1055 validity,
1056 },
1057 TypeId::Float64 => NativeColumn::Float64 {
1058 data: vec![0.0; n],
1059 validity,
1060 },
1061 TypeId::Bool => NativeColumn::Bool {
1062 data: vec![0; n],
1063 validity,
1064 },
1065 _ => NativeColumn::Bytes {
1066 offsets: vec![0u32; n + 1],
1067 values: Vec::new(),
1068 validity,
1069 },
1070 }
1071}
1072
1073#[inline]
1075pub fn validity_bit(validity: &[u8], i: usize) -> bool {
1076 (validity.get(i / 8).copied().unwrap_or(0) >> (i % 8)) & 1 == 1
1077}
1078
1079pub fn all_non_null(validity: &[u8], n: usize) -> bool {
1082 if n == 0 {
1083 return true;
1084 }
1085 let full = n / 8;
1086 if !validity[..full].iter().all(|&b| b == 0xFF) {
1087 return false;
1088 }
1089 if n % 8 != 0 {
1090 let mask = (1u8 << (n % 8)) - 1;
1091 (validity.get(full).copied().unwrap_or(0) & mask) == mask
1092 } else {
1093 true
1094 }
1095}
1096
1097pub fn native_min_max(ty: TypeId, col: &NativeColumn) -> (Option<Vec<u8>>, Option<Vec<u8>>, u64) {
1102 let _ = ty;
1103 match col {
1104 NativeColumn::Int64 { data, validity } => {
1105 let (mut mn, mut mx, mut nulls) = (None::<i64>, None::<i64>, 0u64);
1106 for (i, v) in data.iter().enumerate() {
1107 if !validity_bit(validity, i) {
1108 nulls += 1;
1109 continue;
1110 }
1111 mn = Some(mn.map_or(*v, |m| m.min(*v)));
1112 mx = Some(mx.map_or(*v, |m| m.max(*v)));
1113 }
1114 (
1115 mn.map(|v| v.to_be_bytes().to_vec()),
1116 mx.map(|v| v.to_be_bytes().to_vec()),
1117 nulls,
1118 )
1119 }
1120 NativeColumn::Float64 { data, validity } => {
1121 let (mut mn, mut mx, mut nulls) = (None::<f64>, None::<f64>, 0u64);
1122 for (i, v) in data.iter().enumerate() {
1123 if !validity_bit(validity, i) || v.is_nan() {
1124 nulls += 1;
1125 continue;
1126 }
1127 mn = Some(mn.map_or(*v, |m| m.min(*v)));
1128 mx = Some(mx.map_or(*v, |m| m.max(*v)));
1129 }
1130 (
1131 mn.map(|v| v.to_bits().to_be_bytes().to_vec()),
1132 mx.map(|v| v.to_bits().to_be_bytes().to_vec()),
1133 nulls,
1134 )
1135 }
1136 NativeColumn::Bool { data, validity } => {
1137 let (mut any_t, mut any_f, mut nulls) = (false, false, 0u64);
1138 for (i, v) in data.iter().enumerate() {
1139 if !validity_bit(validity, i) {
1140 nulls += 1;
1141 continue;
1142 }
1143 if *v != 0 {
1144 any_t = true;
1145 } else {
1146 any_f = true;
1147 }
1148 }
1149 let min = if any_f || any_t {
1150 Some(vec![if any_f { 0 } else { 1 }])
1151 } else {
1152 None
1153 };
1154 let max = if any_t || any_f {
1155 Some(vec![if any_t { 1 } else { 0 }])
1156 } else {
1157 None
1158 };
1159 (min, max, nulls)
1160 }
1161 NativeColumn::Bytes {
1162 offsets,
1163 values,
1164 validity,
1165 } => {
1166 let mut mn: Option<&[u8]> = None;
1167 let mut mx: Option<&[u8]> = None;
1168 let mut nulls = 0u64;
1169 for i in 0..offsets.len().saturating_sub(1) {
1170 if !validity_bit(validity, i) {
1171 nulls += 1;
1172 continue;
1173 }
1174 let s = &values[offsets[i] as usize..offsets[i + 1] as usize];
1175 mn = Some(match mn {
1176 None => s,
1177 Some(m) if s < m => s,
1178 Some(m) => m,
1179 });
1180 mx = Some(match mx {
1181 None => s,
1182 Some(m) if s > m => s,
1183 Some(m) => m,
1184 });
1185 }
1186 (mn.map(|s| s.to_vec()), mx.map(|s| s.to_vec()), nulls)
1187 }
1188 }
1189}
1190
1191pub fn page_stat_for(
1195 ty: TypeId,
1196 col: &NativeColumn,
1197 first_row_id: u64,
1198 last_row_id: u64,
1199) -> crate::page::PageStat {
1200 let (min, max, null_count) = native_min_max(ty, col);
1201 crate::page::PageStat {
1202 first_row_id,
1203 last_row_id,
1204 null_count,
1205 row_count: col.len() as u32,
1206 min,
1207 max,
1208 offset: 0,
1209 compressed_len: 0,
1210 uncompressed_len: 0,
1211 }
1212}
1213
1214pub fn encode_key_native(_ty: TypeId, col: &NativeColumn, i: usize) -> Option<Vec<u8>> {
1221 match col {
1222 NativeColumn::Int64 { data, validity } if validity_bit(validity, i) => {
1223 Some(data[i].to_be_bytes().to_vec())
1224 }
1225 NativeColumn::Float64 { data, validity } if validity_bit(validity, i) => {
1226 Some(data[i].to_bits().to_be_bytes().to_vec())
1227 }
1228 NativeColumn::Bool { data, validity } if validity_bit(validity, i) => Some(vec![data[i]]),
1229 NativeColumn::Bytes {
1230 offsets,
1231 values,
1232 validity,
1233 } if validity_bit(validity, i) => {
1234 let lo = offsets[i] as usize;
1235 let hi = offsets[i + 1] as usize;
1236 Some(values[lo..hi].to_vec())
1237 }
1238 _ => None,
1239 }
1240}
1241
1242pub fn native_bytes_at(col: &NativeColumn, i: usize) -> Option<&[u8]> {
1246 match col {
1247 NativeColumn::Bytes {
1248 offsets,
1249 values,
1250 validity,
1251 } if validity_bit(validity, i) => {
1252 let lo = offsets[i] as usize;
1253 let hi = offsets[i + 1] as usize;
1254 Some(&values[lo..hi])
1255 }
1256 _ => None,
1257 }
1258}
1259
1260pub fn values_to_native(ty: TypeId, values: &[Value]) -> NativeColumn {
1263 let n = values.len();
1264 let mut non_null = vec![false; n];
1265 match ty {
1266 TypeId::Int64 | TypeId::TimestampNanos => {
1267 let mut data = Vec::with_capacity(n);
1268 for (i, v) in values.iter().enumerate() {
1269 match v {
1270 Value::Int64(x) => {
1271 non_null[i] = true;
1272 data.push(*x);
1273 }
1274 _ => data.push(0),
1275 }
1276 }
1277 NativeColumn::Int64 {
1278 data,
1279 validity: validity_bitmap_from(non_null),
1280 }
1281 }
1282 TypeId::Float64 => {
1283 let mut data = Vec::with_capacity(n);
1284 for (i, v) in values.iter().enumerate() {
1285 match v {
1286 Value::Float64(x) => {
1287 non_null[i] = true;
1288 data.push(*x);
1289 }
1290 _ => data.push(0.0),
1291 }
1292 }
1293 NativeColumn::Float64 {
1294 data,
1295 validity: validity_bitmap_from(non_null),
1296 }
1297 }
1298 TypeId::Bool => {
1299 let mut data = Vec::with_capacity(n);
1300 for (i, v) in values.iter().enumerate() {
1301 match v {
1302 Value::Bool(x) => {
1303 non_null[i] = true;
1304 data.push(if *x { 1 } else { 0 });
1305 }
1306 _ => data.push(0),
1307 }
1308 }
1309 NativeColumn::Bool {
1310 data,
1311 validity: validity_bitmap_from(non_null),
1312 }
1313 }
1314 _ => {
1315 let mut offsets = Vec::with_capacity(n + 1);
1316 let mut vals = Vec::new();
1317 offsets.push(0u32);
1318 for (i, v) in values.iter().enumerate() {
1319 if let Value::Bytes(b) = v {
1320 non_null[i] = true;
1321 vals.extend_from_slice(b);
1322 }
1323 offsets.push(vals.len() as u32);
1324 }
1325 NativeColumn::Bytes {
1326 offsets,
1327 values: vals,
1328 validity: validity_bitmap_from(non_null),
1329 }
1330 }
1331 }
1332}
1333
1334pub fn rows_to_native(ty: TypeId, rows: &[Vec<(u16, Value)>], column_id: u16) -> NativeColumn {
1343 let n = rows.len();
1344 let mut non_null = vec![false; n];
1345 fn at(row: &[(u16, Value)], column_id: u16) -> Option<&Value> {
1346 row.iter().find(|(id, _)| *id == column_id).map(|(_, v)| v)
1347 }
1348 match ty {
1349 TypeId::Int64 | TypeId::TimestampNanos => {
1350 let mut data = Vec::with_capacity(n);
1351 for (i, row) in rows.iter().enumerate() {
1352 match at(row, column_id) {
1353 Some(Value::Int64(x)) => {
1354 non_null[i] = true;
1355 data.push(*x);
1356 }
1357 _ => data.push(0),
1358 }
1359 }
1360 NativeColumn::Int64 {
1361 data,
1362 validity: validity_bitmap_from(non_null),
1363 }
1364 }
1365 TypeId::Float64 => {
1366 let mut data = Vec::with_capacity(n);
1367 for (i, row) in rows.iter().enumerate() {
1368 match at(row, column_id) {
1369 Some(Value::Float64(x)) => {
1370 non_null[i] = true;
1371 data.push(*x);
1372 }
1373 _ => data.push(0.0),
1374 }
1375 }
1376 NativeColumn::Float64 {
1377 data,
1378 validity: validity_bitmap_from(non_null),
1379 }
1380 }
1381 TypeId::Bool => {
1382 let mut data = Vec::with_capacity(n);
1383 for (i, row) in rows.iter().enumerate() {
1384 match at(row, column_id) {
1385 Some(Value::Bool(x)) => {
1386 non_null[i] = true;
1387 data.push(if *x { 1 } else { 0 });
1388 }
1389 _ => data.push(0),
1390 }
1391 }
1392 NativeColumn::Bool {
1393 data,
1394 validity: validity_bitmap_from(non_null),
1395 }
1396 }
1397 _ => {
1398 let mut offsets = Vec::with_capacity(n + 1);
1399 let mut vals = Vec::new();
1400 offsets.push(0u32);
1401 for (i, row) in rows.iter().enumerate() {
1402 if let Some(Value::Bytes(b)) = at(row, column_id) {
1403 non_null[i] = true;
1404 vals.extend_from_slice(b);
1405 }
1406 offsets.push(vals.len() as u32);
1407 }
1408 NativeColumn::Bytes {
1409 offsets,
1410 values: vals,
1411 validity: validity_bitmap_from(non_null),
1412 }
1413 }
1414 }
1415}
1416
1417fn validity_bitmap_from(non_null: impl IntoIterator<Item = bool>) -> Vec<u8> {
1418 let bits: Vec<bool> = non_null.into_iter().collect();
1419 let n = bits.len();
1420 let mut out = vec![0u8; n.div_ceil(8)];
1421 for (i, &b) in bits.iter().enumerate() {
1422 if b {
1423 out[i / 8] |= 1 << (i % 8);
1424 }
1425 }
1426 out
1427}
1428
1429pub fn encode_page_native(
1436 ty: TypeId,
1437 col: &NativeColumn,
1438 encoding: Encoding,
1439 compress: Compress,
1440 le: bool,
1441) -> Result<Vec<u8>> {
1442 let raw = matches!(compress, Compress::Plain) || matches!(encoding, Encoding::Plain);
1443 Ok(match (ty, col) {
1444 (TypeId::Int64 | TypeId::TimestampNanos, NativeColumn::Int64 { data, validity }) => {
1445 if matches!(encoding, Encoding::Delta) && !raw {
1446 let mut payload = Vec::with_capacity(4 + validity.len() + data.len() * 8);
1447 payload.extend_from_slice(&(validity.len() as u32).to_be_bytes());
1448 payload.extend_from_slice(validity);
1449 let mut deltas = Vec::with_capacity(data.len());
1450 let mut prev = 0i64;
1451 for v in data {
1452 deltas.push(v - prev);
1453 prev = *v;
1454 }
1455 if le {
1456 append_i64_le(&mut payload, &deltas);
1457 } else {
1458 append_i64_be(&mut payload, &deltas);
1459 }
1460 compress_delta_payload(&payload, compress, le)?
1461 } else {
1462 native_plain_page(validity, compress, raw, le, |p| {
1463 if le {
1464 append_i64_le(p, data);
1465 } else {
1466 append_i64_be(p, data);
1467 }
1468 })
1469 }
1470 }
1471 (
1472 TypeId::Float64,
1473 NativeColumn::Float64 {
1474 data: fdata,
1475 validity,
1476 },
1477 ) => native_plain_page(validity, compress, raw, le, |p| {
1478 let bits: &[u64] = bytemuck::cast_slice::<f64, u64>(fdata);
1479 if le {
1480 append_u64_le(p, bits);
1481 } else {
1482 append_u64_be(p, bits);
1483 }
1484 }),
1485 (
1486 TypeId::Bool,
1487 NativeColumn::Bool {
1488 data: bdata,
1489 validity,
1490 },
1491 ) => native_plain_page(validity, compress, raw, le, |p| p.extend_from_slice(bdata)),
1492 (
1493 TypeId::Bytes,
1494 NativeColumn::Bytes {
1495 offsets,
1496 values,
1497 validity,
1498 },
1499 ) => {
1500 if matches!(encoding, Encoding::Dictionary) && !raw {
1501 let dict = dict_encode_bytes_native(offsets, values, validity);
1502 compress_dict_payload(&dict, compress, le)?
1503 } else {
1504 native_plain_page(validity, compress, raw, le, |p| {
1505 let offs: Vec<u64> = offsets.iter().map(|o| *o as u64).collect();
1506 if le {
1507 append_u64_le(p, &offs);
1508 } else {
1509 append_u64_be(p, &offs);
1510 }
1511 p.extend_from_slice(values);
1512 })
1513 }
1514 }
1515 _ => {
1516 return Err(MongrelError::InvalidArgument(format!(
1517 "encode_page_native: unsupported (ty={ty:?})"
1518 )))
1519 }
1520 })
1521}
1522
1523fn compress_delta_payload(payload: &[u8], compress: Compress, le: bool) -> Result<Vec<u8>> {
1527 Ok(match compress {
1528 Compress::Plain => {
1529 let mut out = vec![algo_with_le(ALGO_PLAIN, le)];
1530 out.extend_from_slice(payload);
1531 out
1532 }
1533 Compress::Zstd(level) => {
1534 let mut out = vec![algo_with_le(ALGO_ZSTD_DELTA, le)];
1535 out.extend(zstd_compress_level(payload, level)?);
1536 out
1537 }
1538 Compress::Lz4 => {
1539 let mut out = vec![algo_with_le(ALGO_LZ4_DELTA, le)];
1540 out.extend(lz4_compress(payload));
1541 out
1542 }
1543 })
1544}
1545
1546fn compress_dict_payload(payload: &[u8], compress: Compress, _le: bool) -> Result<Vec<u8>> {
1551 Ok(match compress {
1552 Compress::Plain => {
1553 let mut out = vec![ALGO_PLAIN];
1554 out.extend_from_slice(payload);
1555 out
1556 }
1557 Compress::Zstd(level) => {
1558 let mut out = vec![ALGO_ZSTD_DICT];
1559 out.extend(zstd_compress_level(payload, level)?);
1560 out
1561 }
1562 Compress::Lz4 => {
1563 let mut out = vec![ALGO_LZ4_DICT];
1564 out.extend(lz4_compress(payload));
1565 out
1566 }
1567 })
1568}
1569
1570fn native_plain_page(
1575 validity: &[u8],
1576 compress: Compress,
1577 raw: bool,
1578 le: bool,
1579 fill_payload: impl FnOnce(&mut Vec<u8>),
1580) -> Vec<u8> {
1581 let mut payload = Vec::new();
1582 payload.extend_from_slice(&(validity.len() as u32).to_be_bytes());
1583 payload.extend_from_slice(validity);
1584 fill_payload(&mut payload);
1585 if raw {
1586 let mut out = vec![algo_with_le(ALGO_PLAIN, le)];
1587 out.extend(payload);
1588 out
1589 } else {
1590 match compress {
1591 Compress::Zstd(level) => {
1592 let mut out = vec![algo_with_le(ALGO_ZSTD_PLAIN, le)];
1593 out.extend(zstd_compress_level(&payload, level).expect("zstd compress"));
1594 out
1595 }
1596 Compress::Lz4 => {
1597 let mut out = vec![algo_with_le(ALGO_LZ4_PLAIN, le)];
1598 out.extend(lz4_compress(&payload));
1599 out
1600 }
1601 Compress::Plain => {
1602 let mut out = vec![algo_with_le(ALGO_PLAIN, le)];
1603 out.extend(payload);
1604 out
1605 }
1606 }
1607 }
1608}
1609
1610pub fn decode_page_native(ty: TypeId, page: &[u8], n: usize) -> Result<NativeColumn> {
1615 use std::borrow::Cow;
1616 if page.is_empty() {
1617 return Err(MongrelError::InvalidArgument("empty page".into()));
1618 }
1619 let algo = page[0];
1620 let body = &page[1..];
1621 let le = algo & ALGO_LE_FLAG != 0;
1625 let base = algo & !ALGO_LE_FLAG;
1626 let plain_algo = matches!(base, ALGO_PLAIN | ALGO_ZSTD_PLAIN | ALGO_LZ4_PLAIN);
1627 let delta_algo = matches!(base, ALGO_ZSTD_DELTA | ALGO_LZ4_DELTA);
1628 let dict_algo = matches!(base, ALGO_ZSTD_DICT | ALGO_LZ4_DICT);
1629 if !plain_algo && !delta_algo && !dict_algo {
1630 return Err(MongrelError::InvalidArgument(format!(
1631 "decode_page_native: unsupported algo {algo} for ty {ty:?}"
1632 )));
1633 }
1634 let raw: Cow<[u8]> = match base {
1636 ALGO_PLAIN => Cow::Borrowed(body),
1637 ALGO_ZSTD_PLAIN | ALGO_ZSTD_DELTA | ALGO_ZSTD_DICT => {
1638 Cow::Owned(zstd_decompress(body, max_decompressed_bytes(ty, n, base))?)
1639 }
1640 ALGO_LZ4_PLAIN | ALGO_LZ4_DELTA | ALGO_LZ4_DICT => {
1641 Cow::Owned(lz4_decompress(body, max_decompressed_bytes(ty, n, base))?)
1642 }
1643 _ => unreachable!(),
1644 };
1645
1646 if dict_algo {
1649 return if matches!(ty, TypeId::Bytes) {
1650 dict_decode_bytes_native(&raw, n)
1651 } else {
1652 Err(MongrelError::InvalidArgument(format!(
1653 "decode_page_native: dict algo {algo} only valid for Bytes, got {ty:?}"
1654 )))
1655 };
1656 }
1657
1658 let take_i64 = |p: &[u8]| -> Result<Vec<i64>> {
1660 if le {
1661 take_i64_le(p, n)
1662 } else {
1663 take_i64_be(p, n)
1664 }
1665 };
1666 if delta_algo {
1668 if !matches!(ty, TypeId::Int64 | TypeId::TimestampNanos) {
1669 return Err(MongrelError::InvalidArgument(format!(
1670 "decode_page_native: delta algo {algo} only valid for Int64, got {ty:?}"
1671 )));
1672 }
1673 let (validity, p) = split_validity(&raw)?;
1674 let deltas = take_i64(p)?;
1675 let data = delta_prefix_sum_i64(&deltas);
1676 return Ok(NativeColumn::Int64 { data, validity });
1677 }
1678
1679 match ty {
1681 TypeId::Int64 | TypeId::TimestampNanos => {
1682 let (validity, p) = split_validity(&raw)?;
1683 Ok(NativeColumn::Int64 {
1684 data: take_i64(p)?,
1685 validity,
1686 })
1687 }
1688 TypeId::Float64 => {
1689 let (validity, p) = split_validity(&raw)?;
1690 let bits = if le {
1691 take_u64_le(p, n)?
1692 } else {
1693 take_u64_be(p, n)?
1694 };
1695 let data: Vec<f64> = bits.into_iter().map(f64::from_bits).collect();
1696 Ok(NativeColumn::Float64 { data, validity })
1697 }
1698 TypeId::Bool => {
1699 let (validity, p) = split_validity(&raw)?;
1700 if p.len() < n {
1701 return Err(MongrelError::InvalidArgument(
1702 "bool payload truncated".into(),
1703 ));
1704 }
1705 Ok(NativeColumn::Bool {
1706 data: p[..n].to_vec(),
1707 validity,
1708 })
1709 }
1710 TypeId::Bytes => decode_bytes_plain_payload(&raw, n, le),
1711 _ => Err(MongrelError::InvalidArgument(format!(
1712 "decode_page_native: unsupported ty {ty:?}"
1713 ))),
1714 }
1715}
1716
1717fn split_validity(raw: &[u8]) -> Result<(Vec<u8>, &[u8])> {
1718 if raw.len() < 4 {
1719 return Err(MongrelError::InvalidArgument("page validity header".into()));
1720 }
1721 let vlen = u32::from_be_bytes([raw[0], raw[1], raw[2], raw[3]]) as usize;
1722 if 4 + vlen > raw.len() {
1723 return Err(MongrelError::InvalidArgument("page validity range".into()));
1724 }
1725 Ok((raw[4..4 + vlen].to_vec(), &raw[4 + vlen..]))
1726}
1727
1728fn decode_bytes_plain_payload(raw: &[u8], n: usize, le: bool) -> Result<NativeColumn> {
1731 let (validity, p) = split_validity(raw)?;
1732 let table = (n + 1) * 8;
1733 if p.len() < table {
1734 return Err(MongrelError::InvalidArgument(
1735 "bytes offsets truncated".into(),
1736 ));
1737 }
1738 let offsets_be: Vec<u64> = if le {
1739 take_u64_le(p, n + 1)?
1740 } else {
1741 take_u64_be(p, n + 1)?
1742 };
1743 let offsets: Vec<u32> = offsets_be.into_iter().map(|o| o as u32).collect();
1744 let values = p[table..].to_vec();
1745 Ok(NativeColumn::Bytes {
1746 offsets,
1747 values,
1748 validity,
1749 })
1750}
1751
1752fn take_i64_be(p: &[u8], n: usize) -> Result<Vec<i64>> {
1753 if p.len() < n * 8 {
1754 return Err(MongrelError::InvalidArgument(
1755 "int64 payload truncated".into(),
1756 ));
1757 }
1758 Ok(take_u64_be(p, n)?.into_iter().map(|u| u as i64).collect())
1759}
1760
1761fn delta_prefix_sum_i64(deltas: &[i64]) -> Vec<i64> {
1769 let mut out = vec![0i64; deltas.len()];
1770 prefix_sum_i64_into(deltas, &mut out);
1771 out
1772}
1773
1774fn prefix_sum_i64_into(deltas: &[i64], out: &mut [i64]) {
1775 debug_assert_eq!(deltas.len(), out.len());
1776 #[cfg(target_arch = "x86_64")]
1777 {
1778 if is_x86_feature_detected!("avx2") && deltas.len() >= 4 {
1779 unsafe {
1782 prefix_sum_avx2(deltas, out);
1783 }
1784 return;
1785 }
1786 }
1787 prefix_sum_scalar(deltas, out);
1788}
1789
1790fn prefix_sum_scalar(deltas: &[i64], out: &mut [i64]) {
1791 let mut acc = 0i64;
1792 for (i, &d) in deltas.iter().enumerate() {
1793 acc = acc.wrapping_add(d);
1794 out[i] = acc;
1795 }
1796}
1797
1798#[cfg(target_arch = "x86_64")]
1799#[target_feature(enable = "avx2")]
1800unsafe fn prefix_sum_avx2(deltas: &[i64], out: &mut [i64]) {
1801 use std::arch::x86_64::*;
1802 let n = deltas.len();
1803 let mut running = _mm256_setzero_si256(); let mut i = 0usize;
1805 while i + 4 <= n {
1806 let mut x = _mm256_loadu_si256(deltas.as_ptr().add(i) as *const __m256i);
1807 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);
1816 running = _mm256_permute4x64_epi64(x, 0xFF); i += 4;
1818 }
1819 let mut acc = if i == 0 { 0 } else { out[i - 1] };
1822 while i < n {
1823 acc = acc.wrapping_add(deltas[i]);
1824 out[i] = acc;
1825 i += 1;
1826 }
1827}
1828
1829fn append_u64_be(out: &mut Vec<u8>, data: &[u64]) {
1831 if cfg!(target_endian = "little") {
1832 let swapped: Vec<u64> = data.iter().map(|v| v.swap_bytes()).collect();
1833 out.extend_from_slice(bytemuck::cast_slice::<u64, u8>(&swapped));
1834 } else {
1835 out.extend_from_slice(bytemuck::cast_slice::<u64, u8>(data));
1836 }
1837}
1838
1839fn append_i64_be(out: &mut Vec<u8>, data: &[i64]) {
1841 if cfg!(target_endian = "little") {
1842 let swapped: Vec<i64> = data.iter().map(|v| v.swap_bytes()).collect();
1843 out.extend_from_slice(bytemuck::cast_slice::<i64, u8>(&swapped));
1844 } else {
1845 out.extend_from_slice(bytemuck::cast_slice::<i64, u8>(data));
1846 }
1847}
1848
1849fn take_u64_be(p: &[u8], n: usize) -> Result<Vec<u64>> {
1853 if p.len() < n * 8 {
1854 return Err(MongrelError::InvalidArgument(
1855 "u64 payload truncated".into(),
1856 ));
1857 }
1858 let bytes = &p[..n * 8];
1859 if let Ok(native) = bytemuck::try_cast_slice::<u8, u64>(bytes) {
1860 Ok(native.iter().map(|v| v.swap_bytes()).collect())
1861 } else {
1862 let mut out = Vec::with_capacity(n);
1863 for chunk in bytes.chunks_exact(8) {
1864 out.push(u64::from_be_bytes(chunk.try_into().unwrap()));
1865 }
1866 Ok(out)
1867 }
1868}
1869
1870fn append_u64_le(out: &mut Vec<u8>, data: &[u64]) {
1874 if cfg!(target_endian = "little") {
1875 out.extend_from_slice(bytemuck::cast_slice::<u64, u8>(data));
1876 } else {
1877 let swapped: Vec<u64> = data.iter().map(|v| v.swap_bytes()).collect();
1878 out.extend_from_slice(bytemuck::cast_slice::<u64, u8>(&swapped));
1879 }
1880}
1881
1882fn append_i64_le(out: &mut Vec<u8>, data: &[i64]) {
1884 if cfg!(target_endian = "little") {
1885 out.extend_from_slice(bytemuck::cast_slice::<i64, u8>(data));
1886 } else {
1887 let swapped: Vec<i64> = data.iter().map(|v| v.swap_bytes()).collect();
1888 out.extend_from_slice(bytemuck::cast_slice::<i64, u8>(&swapped));
1889 }
1890}
1891
1892fn take_u64_le(p: &[u8], n: usize) -> Result<Vec<u64>> {
1898 if p.len() < n * 8 {
1899 return Err(MongrelError::InvalidArgument(
1900 "u64 payload truncated".into(),
1901 ));
1902 }
1903 let bytes = &p[..n * 8];
1904 if cfg!(target_endian = "little") {
1905 if let Ok(native) = bytemuck::try_cast_slice::<u8, u64>(bytes) {
1906 return Ok(native.to_vec());
1908 }
1909 Ok(bytes
1910 .chunks_exact(8)
1911 .map(|c| u64::from_le_bytes(c.try_into().unwrap()))
1912 .collect())
1913 } else {
1914 Ok(bytes
1915 .chunks_exact(8)
1916 .map(|c| u64::from_le_bytes(c.try_into().unwrap()))
1917 .collect())
1918 }
1919}
1920
1921fn take_i64_le(p: &[u8], n: usize) -> Result<Vec<i64>> {
1923 if p.len() < n * 8 {
1924 return Err(MongrelError::InvalidArgument(
1925 "int64 payload truncated".into(),
1926 ));
1927 }
1928 Ok(take_u64_le(p, n)?.into_iter().map(|u| u as i64).collect())
1929}
1930
1931fn dict_encode_bytes_native(offsets: &[u32], values: &[u8], validity: &[u8]) -> Vec<u8> {
1932 let n = offsets.len() - 1;
1933 let mut table: Vec<Vec<u8>> = Vec::new();
1934 let mut index_of: std::collections::HashMap<&[u8], u32> = std::collections::HashMap::new();
1935 let mut indices: Vec<u32> = Vec::with_capacity(n);
1936 for i in 0..n {
1937 let lo = offsets[i] as usize;
1938 let hi = offsets[i + 1] as usize;
1939 let slice = &values[lo..hi];
1940 let idx = if let Some(&idx) = index_of.get(slice) {
1941 idx
1942 } else {
1943 let idx = table.len() as u32;
1944 index_of.insert(slice, idx);
1945 table.push(slice.to_vec());
1946 idx
1947 };
1948 indices.push(idx);
1949 }
1950 let mut out = Vec::new();
1951 out.extend_from_slice(&(validity.len() as u32).to_be_bytes());
1952 out.extend_from_slice(validity);
1953 out.extend_from_slice(&(indices.len() as u32).to_be_bytes());
1954 for i in &indices {
1955 out.extend_from_slice(&i.to_be_bytes());
1956 }
1957 out.extend_from_slice(&(table.len() as u32).to_be_bytes());
1958 for entry in &table {
1959 out.extend_from_slice(&(entry.len() as u32).to_be_bytes());
1960 out.extend_from_slice(entry);
1961 }
1962 out
1963}
1964
1965fn dict_decode_bytes_native(data: &[u8], n: usize) -> Result<NativeColumn> {
1966 let mut cur = 0usize;
1967 let vlen = read_u32_be(data, &mut cur)? as usize;
1968 let validity = checked_slice(data, &mut cur, vlen)?.to_vec();
1969 let index_count = read_u32_be(data, &mut cur)? as usize;
1970 if index_count < n {
1971 return Err(MongrelError::InvalidArgument("dict index_count < n".into()));
1972 }
1973 let mut indices = Vec::with_capacity(index_count.min(n));
1974 for _ in 0..index_count {
1975 indices.push(read_u32_be(data, &mut cur)?);
1976 }
1977 let table_count = read_u32_be(data, &mut cur)? as usize;
1978 let mut table: Vec<(usize, usize)> = Vec::with_capacity(table_count); let mut values = Vec::new();
1980 for _ in 0..table_count {
1981 let len = read_u32_be(data, &mut cur)? as usize;
1982 let chunk = checked_slice(data, &mut cur, len)?;
1983 let start = values.len();
1984 values.extend_from_slice(chunk);
1985 table.push((start, len));
1986 }
1987 let mut merged = Vec::new();
1988 let mut offs = vec![0u32];
1989 for (i, &idx) in indices.iter().enumerate().take(n) {
1990 let non_null = (validity.get(i / 8).copied().unwrap_or(0) >> (i % 8)) & 1 == 1;
1994 if non_null {
1995 let (start, len) = table
1996 .get(idx as usize)
1997 .copied()
1998 .ok_or_else(|| MongrelError::InvalidArgument("dict index out of range".into()))?;
1999 merged.extend_from_slice(&values[start..start + len]);
2000 }
2001 offs.push(merged.len() as u32);
2002 }
2003 Ok(NativeColumn::Bytes {
2004 offsets: offs,
2005 values: merged,
2006 validity,
2007 })
2008}
2009
2010#[cfg(test)]
2011mod native_tests {
2012 use super::*;
2013
2014 #[test]
2015 fn native_int64_plain_round_trip() {
2016 let col = NativeColumn::Int64 {
2017 data: (0..1000).collect(),
2018 validity: full_validity(1000),
2019 };
2020 let page = encode_page_native(
2021 TypeId::Int64,
2022 &col,
2023 Encoding::Zstd,
2024 Compress::Zstd(3),
2025 false,
2026 )
2027 .unwrap();
2028 let back = decode_page_native(TypeId::Int64, &page, 1000).unwrap();
2029 match back {
2030 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2031 _ => panic!(),
2032 }
2033 }
2034
2035 #[test]
2036 fn native_int64_delta_crushes_sequential() {
2037 let col = NativeColumn::int64_sequence(0, 100_000);
2038 let plain = encode_page_native(
2039 TypeId::Int64,
2040 &col,
2041 Encoding::Zstd,
2042 Compress::Zstd(3),
2043 false,
2044 )
2045 .unwrap();
2046 let delta = encode_page_native(
2047 TypeId::Int64,
2048 &col,
2049 Encoding::Delta,
2050 Compress::Zstd(3),
2051 false,
2052 )
2053 .unwrap();
2054 assert!(
2055 delta.len() < plain.len() / 5,
2056 "delta must crush sequential ints"
2057 );
2058 let back = decode_page_native(TypeId::Int64, &delta, 100_000).unwrap();
2059 match back {
2060 NativeColumn::Int64 { data, .. } => assert_eq!(data.len(), 100_000),
2061 _ => panic!(),
2062 }
2063 }
2064
2065 #[test]
2066 fn native_bytes_dict_round_trip() {
2067 let n = 500;
2068 let mut offsets = vec![0u32];
2069 let mut values = Vec::new();
2070 for i in 0..n {
2071 let s = ["red", "green", "blue"][i % 3];
2072 values.extend_from_slice(s.as_bytes());
2073 offsets.push(values.len() as u32);
2074 }
2075 let col = NativeColumn::Bytes {
2076 offsets,
2077 values,
2078 validity: full_validity(n),
2079 };
2080 let page = encode_page_native(
2081 TypeId::Bytes,
2082 &col,
2083 Encoding::Dictionary,
2084 Compress::Zstd(3),
2085 false,
2086 )
2087 .unwrap();
2088 assert!(page.len() < 100, "dict page tiny, got {}", page.len());
2089 let back = decode_page_native(TypeId::Bytes, &page, n).unwrap();
2090 assert_eq!(back.len(), n);
2091 }
2092
2093 #[test]
2094 fn native_gather_picks_indices() {
2095 let col = NativeColumn::Int64 {
2096 data: vec![10, 20, 30, 40],
2097 validity: full_validity(4),
2098 };
2099 let g = col.gather(&[0, 2, 3]);
2100 match g {
2101 NativeColumn::Int64 { data, .. } => assert_eq!(data, vec![10, 30, 40]),
2102 _ => panic!(),
2103 }
2104 }
2105
2106 #[test]
2109 fn native_plain_no_zstd_round_trips_all_types() {
2110 let i = NativeColumn::Int64 {
2111 data: (0..1000).collect(),
2112 validity: full_validity(1000),
2113 };
2114 let p =
2115 encode_page_native(TypeId::Int64, &i, Encoding::Plain, Compress::Plain, false).unwrap();
2116 assert_eq!(p[0], ALGO_PLAIN, "Int64 plain must be ALGO_PLAIN");
2117 match decode_page_native(TypeId::Int64, &p, 1000).unwrap() {
2118 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2119 _ => panic!(),
2120 }
2121
2122 let f = NativeColumn::Float64 {
2123 data: (0..500).map(|x| x as f64 * 1.5).collect(),
2124 validity: full_validity(500),
2125 };
2126 let p = encode_page_native(TypeId::Float64, &f, Encoding::Plain, Compress::Plain, false)
2127 .unwrap();
2128 assert_eq!(p[0], ALGO_PLAIN);
2129 match decode_page_native(TypeId::Float64, &p, 500).unwrap() {
2130 NativeColumn::Float64 { data, .. } => {
2131 assert_eq!(data, (0..500).map(|x| x as f64 * 1.5).collect::<Vec<_>>())
2132 }
2133 _ => panic!(),
2134 }
2135
2136 let b = NativeColumn::Bool {
2137 data: (0..64).map(|i| (i % 2) as u8).collect(),
2138 validity: full_validity(64),
2139 };
2140 let p =
2141 encode_page_native(TypeId::Bool, &b, Encoding::Plain, Compress::Plain, false).unwrap();
2142 assert_eq!(p[0], ALGO_PLAIN);
2143 match decode_page_native(TypeId::Bool, &p, 64).unwrap() {
2144 NativeColumn::Bool { data, .. } => {
2145 assert_eq!(data, (0..64).map(|i| (i % 2) as u8).collect::<Vec<_>>())
2146 }
2147 _ => panic!(),
2148 }
2149
2150 let mut offsets = vec![0u32];
2151 let mut values = Vec::new();
2152 for i in 0..200u32 {
2153 values.extend_from_slice(format!("v{i}").as_bytes());
2154 offsets.push(values.len() as u32);
2155 }
2156 let s = NativeColumn::Bytes {
2157 offsets,
2158 values,
2159 validity: full_validity(200),
2160 };
2161 let p =
2162 encode_page_native(TypeId::Bytes, &s, Encoding::Plain, Compress::Plain, false).unwrap();
2163 assert_eq!(p[0], ALGO_PLAIN);
2164 match decode_page_native(TypeId::Bytes, &p, 200).unwrap() {
2165 NativeColumn::Bytes {
2166 offsets: o,
2167 values: v,
2168 ..
2169 } => {
2170 assert_eq!(o.len(), 201);
2171 for i in 0..200 {
2172 let lo = o[i] as usize;
2173 let hi = o[i + 1] as usize;
2174 assert_eq!(&v[lo..hi], format!("v{i}").as_bytes());
2175 }
2176 }
2177 _ => panic!(),
2178 }
2179 }
2180
2181 #[test]
2184 fn lz4_pages_round_trip_all_types() {
2185 let i = NativeColumn::int64_sequence(0, 1000);
2187 let p =
2188 encode_page_native(TypeId::Int64, &i, Encoding::Delta, Compress::Lz4, false).unwrap();
2189 assert_eq!(p[0], ALGO_LZ4_DELTA);
2190 match decode_page_native(TypeId::Int64, &p, 1000).unwrap() {
2191 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2192 _ => panic!(),
2193 }
2194 let p =
2196 encode_page_native(TypeId::Int64, &i, Encoding::Zstd, Compress::Lz4, false).unwrap();
2197 assert_eq!(p[0], ALGO_LZ4_PLAIN);
2198 match decode_page_native(TypeId::Int64, &p, 1000).unwrap() {
2199 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2200 _ => panic!(),
2201 }
2202 let f = NativeColumn::Float64 {
2204 data: (0..500).map(|x| x as f64 * 1.5).collect(),
2205 validity: full_validity(500),
2206 };
2207 let p =
2208 encode_page_native(TypeId::Float64, &f, Encoding::Zstd, Compress::Lz4, false).unwrap();
2209 assert_eq!(p[0], ALGO_LZ4_PLAIN);
2210 match decode_page_native(TypeId::Float64, &p, 500).unwrap() {
2211 NativeColumn::Float64 { data, .. } => {
2212 assert_eq!(data, (0..500).map(|x| x as f64 * 1.5).collect::<Vec<_>>())
2213 }
2214 _ => panic!(),
2215 }
2216 let mut offsets = vec![0u32];
2218 let mut values = Vec::new();
2219 for i in 0..300u32 {
2220 values.extend_from_slice(["red", "green", "blue"][(i % 3) as usize].as_bytes());
2221 offsets.push(values.len() as u32);
2222 }
2223 let s = NativeColumn::Bytes {
2224 offsets,
2225 values,
2226 validity: full_validity(300),
2227 };
2228 let p = encode_page_native(
2229 TypeId::Bytes,
2230 &s,
2231 Encoding::Dictionary,
2232 Compress::Lz4,
2233 false,
2234 )
2235 .unwrap();
2236 assert_eq!(p[0], ALGO_LZ4_DICT);
2237 match decode_page_native(TypeId::Bytes, &p, 300).unwrap() {
2238 NativeColumn::Bytes { offsets: o, .. } => assert_eq!(o.len(), 301),
2239 _ => panic!(),
2240 }
2241 }
2242
2243 #[test]
2247 fn le_pages_round_trip_all_types() {
2248 let assert_le = |page: &[u8]| assert_ne!(page[0] & ALGO_LE_FLAG, 0, "LE flag must be set");
2249
2250 let i = NativeColumn::Int64 {
2252 data: (0..1000).collect(),
2253 validity: full_validity(1000),
2254 };
2255 let p =
2256 encode_page_native(TypeId::Int64, &i, Encoding::Plain, Compress::Plain, true).unwrap();
2257 assert_eq!(p[0], ALGO_LE_FLAG, "raw LE Int64 algo = flag only");
2258 match decode_page_native(TypeId::Int64, &p, 1000).unwrap() {
2259 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2260 _ => panic!(),
2261 }
2262
2263 let p =
2265 encode_page_native(TypeId::Int64, &i, Encoding::Zstd, Compress::Zstd(3), true).unwrap();
2266 assert_le(&p);
2267 match decode_page_native(TypeId::Int64, &p, 1000).unwrap() {
2268 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2269 _ => panic!(),
2270 }
2271
2272 let seq = NativeColumn::int64_sequence(0, 1000);
2274 let p =
2275 encode_page_native(TypeId::Int64, &seq, Encoding::Delta, Compress::Lz4, true).unwrap();
2276 assert_eq!(p[0], ALGO_LZ4_DELTA | ALGO_LE_FLAG);
2277 match decode_page_native(TypeId::Int64, &p, 1000).unwrap() {
2278 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2279 _ => panic!(),
2280 }
2281
2282 let f = NativeColumn::Float64 {
2284 data: (0..500).map(|x| x as f64 * 1.5).collect(),
2285 validity: full_validity(500),
2286 };
2287 let p = encode_page_native(TypeId::Float64, &f, Encoding::Zstd, Compress::Zstd(3), true)
2288 .unwrap();
2289 assert_le(&p);
2290 match decode_page_native(TypeId::Float64, &p, 500).unwrap() {
2291 NativeColumn::Float64 { data, .. } => {
2292 assert_eq!(data, (0..500).map(|x| x as f64 * 1.5).collect::<Vec<_>>())
2293 }
2294 _ => panic!(),
2295 }
2296 }
2297
2298 #[test]
2302 fn malicious_lz4_size_prefix_is_rejected() {
2303 let mut evil = vec![ALGO_LZ4_PLAIN];
2305 evil.extend_from_slice(&0xFFFF_FFFFu32.to_le_bytes());
2306 evil.extend_from_slice(b"junk");
2307 let err =
2308 decode_page_native(TypeId::Int64, &evil, 1000).expect_err("must reject oversize lz4");
2309 let msg = format!("{err}");
2310 assert!(msg.contains("exceeds page limit"), "got: {msg}");
2311
2312 let err = decode_page(TypeId::Int64, &evil, 1000).expect_err("must reject oversize lz4");
2314 assert!(format!("{err}").contains("exceeds page limit"));
2315 }
2316
2317 #[test]
2321 fn truncated_dict_payload_does_not_panic() {
2322 let mut body = vec![ALGO_ZSTD_DICT];
2325 body.extend_from_slice(&zstd_compress(&[0u8; 2]).unwrap()[..]);
2326 decode_page(TypeId::Bytes, &body, 4).expect_err("value dict trunc must Err");
2327
2328 let truncated: &[u8] = &[
2330 0x00, 0x00, 0x00, 0x10, ];
2332 dict_decode_bytes_native(truncated, 2).expect_err("native dict trunc must Err");
2333 dict_decode_bytes(truncated, 2).expect_err("value dict trunc must Err");
2334
2335 let mut bad = vec![0u8; 0];
2338 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");
2344 dict_decode_bytes_native(&bad, 1).expect_err("oob index must Err");
2345 }
2346
2347 #[test]
2351 fn delta_prefix_sum_matches_scalar_all_lengths() {
2352 let mut deltas: Vec<i64> = Vec::with_capacity(2000);
2354 let mut s: u64 = 0x9E37_79B9_7F4A_7C15;
2355 for _ in 0..2000 {
2356 s = s
2357 .wrapping_mul(6364136223846793005)
2358 .wrapping_add(1442695040888963407);
2359 deltas.push((s as i32) as i64); }
2361 for &len in &[
2362 0usize, 1, 2, 3, 4, 5, 6, 7, 8, 9, 15, 16, 17, 31, 32, 33, 63, 64, 65, 127, 128, 129,
2363 1000,
2364 ] {
2365 let input = &deltas[..len];
2366 let simd = delta_prefix_sum_i64(input);
2367 let mut ref_out = vec![0i64; len];
2368 prefix_sum_scalar(input, &mut ref_out);
2369 assert_eq!(simd, ref_out, "len {len}: SIMD prefix sum diverged");
2370 }
2371
2372 let mut ids = vec![0i64; 333];
2374 for x in ids.iter_mut().skip(1) {
2375 *x = 1;
2376 }
2377 let got = delta_prefix_sum_i64(&ids);
2378 assert_eq!(got, (0..333).map(|x| x as i64).collect::<Vec<_>>());
2379 }
2380
2381 #[test]
2384 fn delta_page_decodes_with_null_via_vectorized_path() {
2385 let mut validity = full_validity(17);
2386 validity[9 / 8] &= !(1 << (9 % 8)); let col = NativeColumn::Int64 {
2388 data: (0..17).collect(),
2389 validity: validity.clone(),
2390 };
2391 for (name, comp) in [("zstd", Compress::Zstd(3)), ("lz4", Compress::Lz4)] {
2392 let page =
2393 encode_page_native(TypeId::Int64, &col, Encoding::Delta, comp, false).unwrap();
2394 assert!(matches!(page[0], ALGO_ZSTD_DELTA | ALGO_LZ4_DELTA));
2396 let back = decode_page_native(TypeId::Int64, &page, 17).unwrap();
2397 match back {
2398 NativeColumn::Int64 { data, validity: v } => {
2399 assert_eq!(data, (0..17).collect::<Vec<_>>(), "{name} data");
2400 assert_eq!(v, validity, "{name} validity");
2401 }
2402 _ => panic!(),
2403 }
2404 }
2405 }
2406}