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