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.clone(), 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.clone(), 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.clone(), 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.clone(), 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.clone(), 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.clone(), v)),
67 })?,
68 TypeId::Bytes | TypeId::Json | TypeId::Enum { .. } => 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.clone(), 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.clone(), 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 | TypeId::Enum { .. } => {
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 | TypeId::Enum { .. } => 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 | TypeId::Enum { .. } => {
414 (n + 1).saturating_mul(8) + n.saturating_mul(MAX_VAR_BYTES_PER_ROW)
415 }
416 TypeId::Embedding { dim } => (dim as usize).saturating_mul(8).saturating_mul(n),
417 _ => n.saturating_mul(ty.fixed_size().unwrap_or(8)),
418 };
419 validity + payload
420}
421
422fn lz4_decompress(data: &[u8], max_bytes: usize) -> Result<Vec<u8>> {
423 if data.len() < 4 {
428 return Err(MongrelError::InvalidArgument("lz4 page truncated".into()));
429 }
430 let declared = u32::from_le_bytes([data[0], data[1], data[2], data[3]]) as usize;
431 if declared > max_bytes {
432 return Err(MongrelError::InvalidArgument(format!(
433 "lz4 declared size {declared} exceeds page limit {max_bytes}"
434 )));
435 }
436 lz4_flex::block::decompress(&data[4..], declared)
437 .map_err(|e| MongrelError::InvalidArgument(format!("lz4 decompress: {e}")))
438}
439
440fn zstd_decompress(data: &[u8], max_bytes: usize) -> Result<Vec<u8>> {
441 let out = zstd::decode_all(data)
444 .map_err(|e| MongrelError::InvalidArgument(format!("zstd decompress: {e}")))?;
445 if out.len() > max_bytes {
446 return Err(MongrelError::InvalidArgument(format!(
447 "zstd output {} exceeds page limit {max_bytes}",
448 out.len()
449 )));
450 }
451 Ok(out)
452}
453
454pub fn encode_page(ty: TypeId, values: &[Value], encoding: Encoding) -> Result<Vec<u8>> {
457 Ok(match encoding {
458 Encoding::Plain => {
459 let mut out = vec![ALGO_PLAIN];
460 out.extend(encode_column(ty, values)?);
461 out
462 }
463 Encoding::Dictionary if matches!(ty, TypeId::Bytes | TypeId::Enum { .. }) => {
464 let dict = dict_encode_bytes(values);
465 let mut out = vec![ALGO_ZSTD_DICT];
466 out.extend(zstd_compress(&dict)?);
467 out
468 }
469 _ => {
470 let mut out = vec![ALGO_ZSTD_PLAIN];
471 let enc = encode_column(ty, values)?;
472 out.extend(zstd_compress(&enc)?);
473 out
474 }
475 })
476}
477
478pub fn decode_page(ty: TypeId, page: &[u8], n: usize) -> Result<Vec<Value>> {
480 if page.is_empty() {
481 return Err(MongrelError::InvalidArgument("empty page".into()));
482 }
483 let algo = page[0];
484 let le = algo & ALGO_LE_FLAG != 0;
489 let base = algo & !ALGO_LE_FLAG;
490 let body = &page[1..];
491 let raw_owned;
493 let raw: &[u8] = match base {
494 ALGO_PLAIN => body,
495 ALGO_ZSTD_PLAIN | ALGO_ZSTD_DICT | ALGO_ZSTD_DELTA => {
496 raw_owned = zstd_decompress(body, max_decompressed_bytes(ty.clone(), n, base))?;
497 &raw_owned
498 }
499 ALGO_LZ4_PLAIN | ALGO_LZ4_DICT | ALGO_LZ4_DELTA => {
500 raw_owned = lz4_decompress(body, max_decompressed_bytes(ty.clone(), n, base))?;
501 &raw_owned
502 }
503 other => {
504 return Err(MongrelError::InvalidArgument(format!(
505 "unknown page algo {other}"
506 )))
507 }
508 };
509 match base {
510 ALGO_PLAIN | ALGO_ZSTD_PLAIN | ALGO_LZ4_PLAIN => decode_column(ty, raw, n, le),
511 ALGO_ZSTD_DICT | ALGO_LZ4_DICT => dict_decode_bytes(raw, n),
512 ALGO_ZSTD_DELTA | ALGO_LZ4_DELTA => decode_int64_delta_values(raw, ty, n, le),
513 _ => unreachable!(),
514 }
515}
516
517fn decode_int64_delta_values(raw: &[u8], ty: TypeId, n: usize, le: bool) -> Result<Vec<Value>> {
520 if !matches!(ty, TypeId::Int64 | TypeId::TimestampNanos) {
521 return Err(MongrelError::InvalidArgument(format!(
522 "delta page not valid for {ty:?}"
523 )));
524 }
525 let (validity, p) = split_validity(raw)?;
526 let deltas = if le {
527 take_i64_le(p, n)?
528 } else {
529 take_i64_be(p, n)?
530 };
531 let data = delta_prefix_sum_i64(&deltas);
532 let mut out = Vec::with_capacity(n);
533 for (i, &v) in data.iter().enumerate() {
534 let non_null = (validity.get(i / 8).copied().unwrap_or(0) >> (i % 8)) & 1 == 1;
535 out.push(if non_null {
536 Value::Int64(v)
537 } else {
538 Value::Null
539 });
540 }
541 Ok(out)
542}
543
544fn dict_encode_bytes(values: &[Value]) -> Vec<u8> {
548 let validity = validity_bitmap(values);
549 let mut table: Vec<Vec<u8>> = Vec::new();
550 let mut index_of: std::collections::HashMap<&[u8], u32> = std::collections::HashMap::new();
551 let mut indices: Vec<u32> = Vec::with_capacity(values.len());
552 for v in values {
553 let idx = match v {
554 Value::Bytes(b) => {
555 if let Some(&i) = index_of.get(b.as_slice()) {
556 i
557 } else {
558 let i = table.len() as u32;
559 index_of.insert(b.as_slice(), i);
560 table.push(b.clone());
561 i
562 }
563 }
564 _ => 0,
565 };
566 indices.push(idx);
567 }
568
569 let mut out = Vec::new();
570 out.extend_from_slice(&(validity.len() as u32).to_be_bytes());
571 out.extend_from_slice(&validity);
572 out.extend_from_slice(&(indices.len() as u32).to_be_bytes());
573 for i in &indices {
574 out.extend_from_slice(&i.to_be_bytes());
575 }
576 out.extend_from_slice(&(table.len() as u32).to_be_bytes());
577 for entry in &table {
578 out.extend_from_slice(&(entry.len() as u32).to_be_bytes());
579 out.extend_from_slice(entry);
580 }
581 out
582}
583
584fn dict_decode_bytes(data: &[u8], n: usize) -> Result<Vec<Value>> {
585 let mut cur = 0usize;
586 let vlen = read_u32_be(data, &mut cur)? as usize;
587 let validity = checked_slice(data, &mut cur, vlen)?;
588 let index_count = read_u32_be(data, &mut cur)? as usize;
589 let mut indices = Vec::with_capacity(index_count.min(n));
590 for _ in 0..index_count {
591 indices.push(read_u32_be(data, &mut cur)?);
592 }
593 let table_count = read_u32_be(data, &mut cur)? as usize;
594 let mut table: Vec<Vec<u8>> = Vec::with_capacity(table_count);
595 for _ in 0..table_count {
596 let len = read_u32_be(data, &mut cur)? as usize;
597 table.push(checked_slice(data, &mut cur, len)?.to_vec());
598 }
599
600 let mut out = Vec::with_capacity(n);
601 for (i, &idx) in indices.iter().enumerate().take(n) {
602 let non_null = (validity.get(i / 8).copied().unwrap_or(0) >> (i % 8)) & 1 == 1;
603 if !non_null {
604 out.push(Value::Null);
605 } else {
606 let entry = table
607 .get(idx as usize)
608 .cloned()
609 .ok_or_else(|| MongrelError::InvalidArgument("dict index out of range".into()))?;
610 out.push(Value::Bytes(entry));
611 }
612 }
613 Ok(out)
614}
615
616fn read_u32_be(data: &[u8], cur: &mut usize) -> Result<u32> {
617 if *cur + 4 > data.len() {
618 return Err(MongrelError::InvalidArgument(
619 "dict payload truncated".into(),
620 ));
621 }
622 let v = u32::from_be_bytes([data[*cur], data[*cur + 1], data[*cur + 2], data[*cur + 3]]);
623 *cur += 4;
624 Ok(v)
625}
626
627fn checked_slice<'a>(data: &'a [u8], cur: &mut usize, len: usize) -> Result<&'a [u8]> {
630 if *cur + len > data.len() {
631 return Err(MongrelError::InvalidArgument(
632 "dict payload truncated".into(),
633 ));
634 }
635 let s = &data[*cur..*cur + len];
636 *cur += len;
637 Ok(s)
638}
639
640#[cfg(test)]
641mod compressed_tests {
642 use super::*;
643
644 #[test]
645 fn zstd_plain_round_trip_int64() {
646 let vals: Vec<Value> = (0..1000).map(Value::Int64).collect();
647 let page = encode_page(TypeId::Int64, &vals, Encoding::Zstd).unwrap();
648 assert!(
649 page.len() < vals.len() * 8,
650 "zstd must shrink sequential ints"
651 );
652 let back = decode_page(TypeId::Int64, &page, vals.len()).unwrap();
653 assert_eq!(back, vals);
654 }
655
656 #[test]
657 fn dictionary_round_trip_low_card_bytes() {
658 let palette: &[&[u8]] = &[b"red", b"green", b"blue", b"red"];
659 let vals: Vec<Value> = (0..500)
660 .map(|i| Value::Bytes(palette[i % palette.len()].to_vec()))
661 .collect();
662 let page = encode_page(TypeId::Bytes, &vals, Encoding::Dictionary).unwrap();
663 assert!(
664 page.len() < 100,
665 "4 distinct strings over 500 rows must compress to a tiny page, got {}",
666 page.len()
667 );
668 let back = decode_page(TypeId::Bytes, &page, vals.len()).unwrap();
669 assert_eq!(back, vals);
670 }
671
672 #[test]
673 fn plain_page_still_round_trips() {
674 let vals = vec![Value::Int64(1), Value::Null, Value::Int64(9)];
675 let page = encode_page(TypeId::Int64, &vals, Encoding::Plain).unwrap();
676 assert_eq!(page[0], ALGO_PLAIN);
677 assert_eq!(decode_page(TypeId::Int64, &page, 3).unwrap(), vals);
678 }
679}
680
681#[cfg(test)]
682mod tests {
683 use super::*;
684
685 #[test]
686 fn round_trips_int64_with_nulls() {
687 let vals = vec![
688 Value::Int64(1),
689 Value::Null,
690 Value::Int64(-5),
691 Value::Int64(1 << 40),
692 ];
693 let page = encode_column(TypeId::Int64, &vals).unwrap();
694 let back = decode_column(TypeId::Int64, &page, vals.len(), false).unwrap();
695 assert_eq!(back, vals);
696 }
697
698 #[test]
699 fn round_trips_bytes() {
700 let vals = vec![
701 Value::Bytes(b"hello".to_vec()),
702 Value::Null,
703 Value::Bytes(b"".to_vec()),
704 Value::Bytes(b"wide \x00 byte".to_vec()),
705 ];
706 let page = encode_column(TypeId::Bytes, &vals).unwrap();
707 let back = decode_column(TypeId::Bytes, &page, vals.len(), false).unwrap();
708 assert_eq!(back, vals);
709 }
710
711 #[test]
712 fn round_trips_embedding() {
713 let vals = vec![
714 Value::Embedding(vec![1.0, -2.5, 3.0]),
715 Value::Null,
716 Value::Embedding(vec![0.0; 3]),
717 ];
718 let page = encode_column(TypeId::Embedding { dim: 3 }, &vals).unwrap();
719 let back = decode_column(TypeId::Embedding { dim: 3 }, &page, vals.len(), false).unwrap();
720 assert_eq!(back, vals);
721 }
722
723 #[test]
724 fn round_trips_bool() {
725 let vals = vec![
726 Value::Bool(true),
727 Value::Bool(false),
728 Value::Null,
729 Value::Bool(true),
730 ];
731 let page = encode_column(TypeId::Bool, &vals).unwrap();
732 assert_eq!(
733 decode_column(TypeId::Bool, &page, vals.len(), false).unwrap(),
734 vals
735 );
736 }
737}
738
739#[derive(Debug, Clone, Serialize, Deserialize)]
748pub enum NativeColumn {
749 Int64 {
750 data: Vec<i64>,
751 validity: Vec<u8>,
752 },
753 Float64 {
754 data: Vec<f64>,
755 validity: Vec<u8>,
756 },
757 Bool {
759 data: Vec<u8>,
760 validity: Vec<u8>,
761 },
762 Bytes {
764 offsets: Vec<u32>,
765 values: Vec<u8>,
766 validity: Vec<u8>,
767 },
768}
769
770impl NativeColumn {
771 pub fn len(&self) -> usize {
772 match self {
773 NativeColumn::Int64 { data, .. } => data.len(),
774 NativeColumn::Float64 { data, .. } => data.len(),
775 NativeColumn::Bool { data, .. } => data.len(),
776 NativeColumn::Bytes { offsets, .. } => offsets.len().saturating_sub(1),
777 }
778 }
779
780 pub fn is_empty(&self) -> bool {
781 self.len() == 0
782 }
783
784 pub fn validity(&self) -> &[u8] {
787 match self {
788 NativeColumn::Int64 { validity, .. }
789 | NativeColumn::Float64 { validity, .. }
790 | NativeColumn::Bool { validity, .. }
791 | NativeColumn::Bytes { validity, .. } => validity,
792 }
793 }
794
795 pub fn validate(&self) -> bool {
799 match self {
800 NativeColumn::Int64 { data, validity } => {
801 validity.len() == data.len().div_ceil(8) || validity.is_empty()
802 }
803 NativeColumn::Float64 { data, validity } => {
804 validity.len() == data.len().div_ceil(8) || validity.is_empty()
805 }
806 NativeColumn::Bool { data, validity } => {
807 validity.len() == data.len().div_ceil(8) || validity.is_empty()
808 }
809 NativeColumn::Bytes {
810 offsets,
811 values,
812 validity,
813 } => {
814 let n = offsets.len().saturating_sub(1);
815 (validity.len() == n.div_ceil(8) || validity.is_empty())
816 && offsets
817 .last()
818 .map(|&last| (last as usize) <= values.len())
819 .unwrap_or(true)
820 }
821 }
822 }
823
824 pub fn null_count(&self, n: usize) -> usize {
827 if n == 0 {
828 return 0;
829 }
830 let validity = match self {
831 NativeColumn::Int64 { validity, .. }
832 | NativeColumn::Float64 { validity, .. }
833 | NativeColumn::Bool { validity, .. }
834 | NativeColumn::Bytes { validity, .. } => validity,
835 };
836 if validity.is_empty() {
837 return 0;
838 }
839 let full = n / 8;
842 let mut set = validity[..full.min(validity.len())]
843 .iter()
844 .map(|b| b.count_ones() as usize)
845 .sum::<usize>();
846 let tail_bits = n % 8;
847 if tail_bits > 0 {
848 if let Some(&b) = validity.get(full) {
849 set += (b & ((1u8 << tail_bits) - 1)).count_ones() as usize;
850 }
851 }
852 n - set
853 }
854
855 pub fn approx_bytes(&self) -> u64 {
857 match self {
858 NativeColumn::Int64 { data, validity } => {
859 (data.len() as u64) * 8 + validity.len() as u64
860 }
861 NativeColumn::Float64 { data, validity } => {
862 (data.len() as u64) * 8 + validity.len() as u64
863 }
864 NativeColumn::Bool { data, validity } => data.len() as u64 + validity.len() as u64,
865 NativeColumn::Bytes {
866 offsets,
867 values,
868 validity,
869 } => values.len() as u64 + (offsets.len() as u64) * 4 + validity.len() as u64,
870 }
871 }
872
873 pub fn int64_sequence(start: i64, n: usize) -> Self {
875 NativeColumn::Int64 {
876 data: (0..n).map(|i| start + i as i64).collect(),
877 validity: full_validity(n),
878 }
879 }
880
881 pub fn int64_constant(value: i64, n: usize) -> Self {
883 NativeColumn::Int64 {
884 data: vec![value; n],
885 validity: full_validity(n),
886 }
887 }
888
889 pub fn bool_constant(value: bool, n: usize) -> Self {
891 NativeColumn::Bool {
892 data: vec![if value { 1 } else { 0 }; n],
893 validity: full_validity(n),
894 }
895 }
896
897 pub fn gather(&self, indices: &[usize]) -> NativeColumn {
900 let bit = |v: &[u8], i: usize| (v.get(i / 8).copied().unwrap_or(0) >> (i % 8)) & 1 == 1;
901 match self {
902 NativeColumn::Int64 { data, validity } => NativeColumn::Int64 {
903 data: indices.iter().map(|&i| data[i]).collect(),
904 validity: validity_bitmap_from(indices.iter().map(|&i| bit(validity, i))),
905 },
906 NativeColumn::Float64 {
907 data: fdata,
908 validity: fval,
909 } => NativeColumn::Float64 {
910 data: indices.iter().map(|&i| fdata[i]).collect(),
911 validity: validity_bitmap_from(indices.iter().map(|&i| bit(fval, i))),
912 },
913 NativeColumn::Bool {
914 data: bdata,
915 validity: bval,
916 } => NativeColumn::Bool {
917 data: indices.iter().map(|&i| bdata[i]).collect(),
918 validity: validity_bitmap_from(indices.iter().map(|&i| bit(bval, i))),
919 },
920 NativeColumn::Bytes {
921 offsets,
922 values,
923 validity,
924 } => {
925 let mut out_offsets = Vec::with_capacity(indices.len() + 1);
926 let mut out_values = Vec::new();
927 out_offsets.push(0);
928 for &i in indices {
929 let lo = offsets[i] as usize;
930 let hi = offsets[i + 1] as usize;
931 out_values.extend_from_slice(&values[lo..hi]);
932 out_offsets.push(out_values.len() as u32);
933 }
934 NativeColumn::Bytes {
935 offsets: out_offsets,
936 values: out_values,
937 validity: validity_bitmap_from(indices.iter().map(|&i| bit(validity, i))),
938 }
939 }
940 }
941 }
942
943 pub fn value_at(&self, idx: usize) -> Option<Value> {
948 match self {
949 NativeColumn::Int64 { data, validity } => {
950 if !validity_bit(validity, idx) {
951 return None;
952 }
953 data.get(idx).copied().map(Value::Int64)
954 }
955 NativeColumn::Float64 { data, validity } => {
956 if !validity_bit(validity, idx) {
957 return None;
958 }
959 data.get(idx).copied().map(Value::Float64)
960 }
961 NativeColumn::Bool { data, validity } => {
962 if !validity_bit(validity, idx) {
963 return None;
964 }
965 data.get(idx).copied().map(|b| Value::Bool(b != 0))
966 }
967 NativeColumn::Bytes {
968 offsets,
969 values,
970 validity,
971 } => {
972 if !validity_bit(validity, idx) {
973 return None;
974 }
975 if idx + 1 >= offsets.len() {
976 return None;
977 }
978 let lo = offsets[idx] as usize;
979 let hi = offsets[idx + 1] as usize;
980 Some(Value::Bytes(values[lo..hi].to_vec()))
981 }
982 }
983 }
984
985 pub fn slice_range(&self, start: usize, end: usize) -> NativeColumn {
988 let mk_validity = |v: &[u8]| -> Vec<bool> {
989 (0..(end - start))
990 .map(|i| validity_bit(v, start + i))
991 .collect()
992 };
993 match self {
994 NativeColumn::Int64 { data, validity } => NativeColumn::Int64 {
995 data: data[start..end].to_vec(),
996 validity: validity_bitmap_from(mk_validity(validity)),
997 },
998 NativeColumn::Float64 { data, validity } => NativeColumn::Float64 {
999 data: data[start..end].to_vec(),
1000 validity: validity_bitmap_from(mk_validity(validity)),
1001 },
1002 NativeColumn::Bool { data, validity } => NativeColumn::Bool {
1003 data: data[start..end].to_vec(),
1004 validity: validity_bitmap_from(mk_validity(validity)),
1005 },
1006 NativeColumn::Bytes {
1007 offsets,
1008 values,
1009 validity,
1010 } => {
1011 let lo = offsets[start] as usize;
1012 let hi = offsets[end] as usize;
1013 let new_offsets: Vec<u32> = offsets[start..=end]
1014 .iter()
1015 .map(|o| *o - offsets[start])
1016 .collect();
1017 NativeColumn::Bytes {
1018 offsets: new_offsets,
1019 values: values[lo..hi].to_vec(),
1020 validity: validity_bitmap_from(mk_validity(validity)),
1021 }
1022 }
1023 }
1024 }
1025
1026 pub fn concat(parts: &[NativeColumn]) -> NativeColumn {
1029 match parts.first() {
1030 Some(NativeColumn::Int64 { .. }) => {
1031 let mut data = Vec::new();
1032 let mut non_null: Vec<bool> = Vec::new();
1033 for p in parts {
1034 if let NativeColumn::Int64 { data: d, validity } = p {
1035 data.extend_from_slice(d);
1036 non_null.extend((0..d.len()).map(|i| validity_bit(validity, i)));
1037 }
1038 }
1039 NativeColumn::Int64 {
1040 data,
1041 validity: validity_bitmap_from(non_null),
1042 }
1043 }
1044 Some(NativeColumn::Float64 { .. }) => {
1045 let mut data = Vec::new();
1046 let mut non_null: Vec<bool> = Vec::new();
1047 for p in parts {
1048 if let NativeColumn::Float64 { data: d, validity } = p {
1049 data.extend_from_slice(d);
1050 non_null.extend((0..d.len()).map(|i| validity_bit(validity, i)));
1051 }
1052 }
1053 NativeColumn::Float64 {
1054 data,
1055 validity: validity_bitmap_from(non_null),
1056 }
1057 }
1058 Some(NativeColumn::Bool { .. }) => {
1059 let mut data = Vec::new();
1060 let mut non_null: Vec<bool> = Vec::new();
1061 for p in parts {
1062 if let NativeColumn::Bool { data: d, validity } = p {
1063 data.extend_from_slice(d);
1064 non_null.extend((0..d.len()).map(|i| validity_bit(validity, i)));
1065 }
1066 }
1067 NativeColumn::Bool {
1068 data,
1069 validity: validity_bitmap_from(non_null),
1070 }
1071 }
1072 Some(NativeColumn::Bytes { .. }) => {
1073 let mut offsets: Vec<u32> = vec![0];
1074 let mut values = Vec::new();
1075 let mut non_null: Vec<bool> = Vec::new();
1076 for p in parts {
1077 if let NativeColumn::Bytes {
1078 offsets: off,
1079 values: val,
1080 validity,
1081 } = p
1082 {
1083 for w in off.windows(2) {
1084 values.extend_from_slice(&val[w[0] as usize..w[1] as usize]);
1085 offsets.push(values.len() as u32);
1086 }
1087 non_null.extend((0..off.len() - 1).map(|i| validity_bit(validity, i)));
1088 }
1089 }
1090 NativeColumn::Bytes {
1091 offsets,
1092 values,
1093 validity: validity_bitmap_from(non_null),
1094 }
1095 }
1096 None => NativeColumn::Bytes {
1097 offsets: vec![0],
1098 values: Vec::new(),
1099 validity: Vec::new(),
1100 },
1101 }
1102 }
1103}
1104
1105fn full_validity(n: usize) -> Vec<u8> {
1106 validity_bitmap_from(std::iter::repeat(true).take(n))
1107}
1108
1109pub fn null_native(ty: TypeId, n: usize) -> NativeColumn {
1112 let validity = vec![0u8; n.div_ceil(8)];
1113 match ty {
1114 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date64 | TypeId::Time64 => {
1115 NativeColumn::Int64 {
1116 data: vec![0; n],
1117 validity,
1118 }
1119 }
1120 TypeId::Float64 => NativeColumn::Float64 {
1121 data: vec![0.0; n],
1122 validity,
1123 },
1124 TypeId::Bool => NativeColumn::Bool {
1125 data: vec![0; n],
1126 validity,
1127 },
1128 _ => NativeColumn::Bytes {
1129 offsets: vec![0u32; n + 1],
1130 values: Vec::new(),
1131 validity,
1132 },
1133 }
1134}
1135
1136#[inline]
1138pub fn validity_bit(validity: &[u8], i: usize) -> bool {
1139 (validity.get(i / 8).copied().unwrap_or(0) >> (i % 8)) & 1 == 1
1140}
1141
1142pub fn all_non_null(validity: &[u8], n: usize) -> bool {
1145 if n == 0 {
1146 return true;
1147 }
1148 let full = n / 8;
1149 if !validity[..full].iter().all(|&b| b == 0xFF) {
1150 return false;
1151 }
1152 if n % 8 != 0 {
1153 let mask = (1u8 << (n % 8)) - 1;
1154 (validity.get(full).copied().unwrap_or(0) & mask) == mask
1155 } else {
1156 true
1157 }
1158}
1159
1160pub fn native_min_max(ty: TypeId, col: &NativeColumn) -> (Option<Vec<u8>>, Option<Vec<u8>>, u64) {
1165 let _ = ty;
1166 match col {
1167 NativeColumn::Int64 { data, validity } => {
1168 let (mut mn, mut mx, mut nulls) = (None::<i64>, None::<i64>, 0u64);
1169 for (i, v) in data.iter().enumerate() {
1170 if !validity_bit(validity, i) {
1171 nulls += 1;
1172 continue;
1173 }
1174 mn = Some(mn.map_or(*v, |m| m.min(*v)));
1175 mx = Some(mx.map_or(*v, |m| m.max(*v)));
1176 }
1177 (
1178 mn.map(|v| v.to_be_bytes().to_vec()),
1179 mx.map(|v| v.to_be_bytes().to_vec()),
1180 nulls,
1181 )
1182 }
1183 NativeColumn::Float64 { data, validity } => {
1184 let (mut mn, mut mx, mut nulls) = (None::<f64>, None::<f64>, 0u64);
1185 for (i, v) in data.iter().enumerate() {
1186 if !validity_bit(validity, i) || v.is_nan() {
1187 nulls += 1;
1188 continue;
1189 }
1190 mn = Some(mn.map_or(*v, |m| m.min(*v)));
1191 mx = Some(mx.map_or(*v, |m| m.max(*v)));
1192 }
1193 (
1194 mn.map(|v| v.to_bits().to_be_bytes().to_vec()),
1195 mx.map(|v| v.to_bits().to_be_bytes().to_vec()),
1196 nulls,
1197 )
1198 }
1199 NativeColumn::Bool { data, validity } => {
1200 let (mut any_t, mut any_f, mut nulls) = (false, false, 0u64);
1201 for (i, v) in data.iter().enumerate() {
1202 if !validity_bit(validity, i) {
1203 nulls += 1;
1204 continue;
1205 }
1206 if *v != 0 {
1207 any_t = true;
1208 } else {
1209 any_f = true;
1210 }
1211 }
1212 let min = if any_f || any_t {
1213 Some(vec![if any_f { 0 } else { 1 }])
1214 } else {
1215 None
1216 };
1217 let max = if any_t || any_f {
1218 Some(vec![if any_t { 1 } else { 0 }])
1219 } else {
1220 None
1221 };
1222 (min, max, nulls)
1223 }
1224 NativeColumn::Bytes {
1225 offsets,
1226 values,
1227 validity,
1228 } => {
1229 let mut mn: Option<&[u8]> = None;
1230 let mut mx: Option<&[u8]> = None;
1231 let mut nulls = 0u64;
1232 for i in 0..offsets.len().saturating_sub(1) {
1233 if !validity_bit(validity, i) {
1234 nulls += 1;
1235 continue;
1236 }
1237 let s = &values[offsets[i] as usize..offsets[i + 1] as usize];
1238 mn = Some(match mn {
1239 None => s,
1240 Some(m) if s < m => s,
1241 Some(m) => m,
1242 });
1243 mx = Some(match mx {
1244 None => s,
1245 Some(m) if s > m => s,
1246 Some(m) => m,
1247 });
1248 }
1249 (mn.map(|s| s.to_vec()), mx.map(|s| s.to_vec()), nulls)
1250 }
1251 }
1252}
1253
1254pub fn page_stat_for(
1258 ty: TypeId,
1259 col: &NativeColumn,
1260 first_row_id: u64,
1261 last_row_id: u64,
1262) -> crate::page::PageStat {
1263 let (min, max, null_count) = native_min_max(ty, col);
1264 crate::page::PageStat {
1265 first_row_id,
1266 last_row_id,
1267 null_count,
1268 row_count: col.len() as u32,
1269 min,
1270 max,
1271 offset: 0,
1272 compressed_len: 0,
1273 uncompressed_len: 0,
1274 }
1275}
1276
1277pub fn encode_key_native(_ty: TypeId, col: &NativeColumn, i: usize) -> Option<Vec<u8>> {
1284 match col {
1285 NativeColumn::Int64 { data, validity } if validity_bit(validity, i) => {
1286 Some(data[i].to_be_bytes().to_vec())
1287 }
1288 NativeColumn::Float64 { data, validity } if validity_bit(validity, i) => {
1289 Some(data[i].to_bits().to_be_bytes().to_vec())
1290 }
1291 NativeColumn::Bool { data, validity } if validity_bit(validity, i) => Some(vec![data[i]]),
1292 NativeColumn::Bytes {
1293 offsets,
1294 values,
1295 validity,
1296 } if validity_bit(validity, i) => {
1297 let lo = offsets[i] as usize;
1298 let hi = offsets[i + 1] as usize;
1299 Some(values[lo..hi].to_vec())
1300 }
1301 _ => None,
1302 }
1303}
1304
1305pub fn native_bytes_at(col: &NativeColumn, i: usize) -> Option<&[u8]> {
1309 match col {
1310 NativeColumn::Bytes {
1311 offsets,
1312 values,
1313 validity,
1314 } if validity_bit(validity, i) => {
1315 let lo = offsets[i] as usize;
1316 let hi = offsets[i + 1] as usize;
1317 Some(&values[lo..hi])
1318 }
1319 _ => None,
1320 }
1321}
1322
1323pub fn values_to_native(ty: TypeId, values: &[Value]) -> NativeColumn {
1326 let n = values.len();
1327 let mut non_null = vec![false; n];
1328 match ty {
1329 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date64 | TypeId::Time64 => {
1330 let mut data = Vec::with_capacity(n);
1331 for (i, v) in values.iter().enumerate() {
1332 match v {
1333 Value::Int64(x) => {
1334 non_null[i] = true;
1335 data.push(*x);
1336 }
1337 _ => data.push(0),
1338 }
1339 }
1340 NativeColumn::Int64 {
1341 data,
1342 validity: validity_bitmap_from(non_null),
1343 }
1344 }
1345 TypeId::Float64 => {
1346 let mut data = Vec::with_capacity(n);
1347 for (i, v) in values.iter().enumerate() {
1348 match v {
1349 Value::Float64(x) => {
1350 non_null[i] = true;
1351 data.push(*x);
1352 }
1353 _ => data.push(0.0),
1354 }
1355 }
1356 NativeColumn::Float64 {
1357 data,
1358 validity: validity_bitmap_from(non_null),
1359 }
1360 }
1361 TypeId::Bool => {
1362 let mut data = Vec::with_capacity(n);
1363 for (i, v) in values.iter().enumerate() {
1364 match v {
1365 Value::Bool(x) => {
1366 non_null[i] = true;
1367 data.push(if *x { 1 } else { 0 });
1368 }
1369 _ => data.push(0),
1370 }
1371 }
1372 NativeColumn::Bool {
1373 data,
1374 validity: validity_bitmap_from(non_null),
1375 }
1376 }
1377 _ => {
1378 let mut offsets = Vec::with_capacity(n + 1);
1379 let mut vals = Vec::new();
1380 offsets.push(0u32);
1381 for (i, v) in values.iter().enumerate() {
1382 if let Value::Bytes(b) = v {
1383 non_null[i] = true;
1384 vals.extend_from_slice(b);
1385 }
1386 offsets.push(vals.len() as u32);
1387 }
1388 NativeColumn::Bytes {
1389 offsets,
1390 values: vals,
1391 validity: validity_bitmap_from(non_null),
1392 }
1393 }
1394 }
1395}
1396
1397pub fn rows_to_native(ty: TypeId, rows: &[Vec<(u16, Value)>], column_id: u16) -> NativeColumn {
1406 let n = rows.len();
1407 let mut non_null = vec![false; n];
1408 fn at(row: &[(u16, Value)], column_id: u16) -> Option<&Value> {
1409 row.iter().find(|(id, _)| *id == column_id).map(|(_, v)| v)
1410 }
1411 match ty {
1412 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date64 | TypeId::Time64 => {
1413 let mut data = Vec::with_capacity(n);
1414 for (i, row) in rows.iter().enumerate() {
1415 match at(row, column_id) {
1416 Some(Value::Int64(x)) => {
1417 non_null[i] = true;
1418 data.push(*x);
1419 }
1420 _ => data.push(0),
1421 }
1422 }
1423 NativeColumn::Int64 {
1424 data,
1425 validity: validity_bitmap_from(non_null),
1426 }
1427 }
1428 TypeId::Float64 => {
1429 let mut data = Vec::with_capacity(n);
1430 for (i, row) in rows.iter().enumerate() {
1431 match at(row, column_id) {
1432 Some(Value::Float64(x)) => {
1433 non_null[i] = true;
1434 data.push(*x);
1435 }
1436 _ => data.push(0.0),
1437 }
1438 }
1439 NativeColumn::Float64 {
1440 data,
1441 validity: validity_bitmap_from(non_null),
1442 }
1443 }
1444 TypeId::Bool => {
1445 let mut data = Vec::with_capacity(n);
1446 for (i, row) in rows.iter().enumerate() {
1447 match at(row, column_id) {
1448 Some(Value::Bool(x)) => {
1449 non_null[i] = true;
1450 data.push(if *x { 1 } else { 0 });
1451 }
1452 _ => data.push(0),
1453 }
1454 }
1455 NativeColumn::Bool {
1456 data,
1457 validity: validity_bitmap_from(non_null),
1458 }
1459 }
1460 _ => {
1461 let mut offsets = Vec::with_capacity(n + 1);
1462 let mut vals = Vec::new();
1463 offsets.push(0u32);
1464 for (i, row) in rows.iter().enumerate() {
1465 if let Some(Value::Bytes(b)) = at(row, column_id) {
1466 non_null[i] = true;
1467 vals.extend_from_slice(b);
1468 }
1469 offsets.push(vals.len() as u32);
1470 }
1471 NativeColumn::Bytes {
1472 offsets,
1473 values: vals,
1474 validity: validity_bitmap_from(non_null),
1475 }
1476 }
1477 }
1478}
1479
1480fn validity_bitmap_from(non_null: impl IntoIterator<Item = bool>) -> Vec<u8> {
1481 let bits: Vec<bool> = non_null.into_iter().collect();
1482 let n = bits.len();
1483 let mut out = vec![0u8; n.div_ceil(8)];
1484 for (i, &b) in bits.iter().enumerate() {
1485 if b {
1486 out[i / 8] |= 1 << (i % 8);
1487 }
1488 }
1489 out
1490}
1491
1492pub fn encode_page_native(
1499 ty: TypeId,
1500 col: &NativeColumn,
1501 encoding: Encoding,
1502 compress: Compress,
1503 le: bool,
1504) -> Result<Vec<u8>> {
1505 let raw = matches!(compress, Compress::Plain) || matches!(encoding, Encoding::Plain);
1506 Ok(match (ty.clone(), col) {
1507 (TypeId::Int64 | TypeId::TimestampNanos, NativeColumn::Int64 { data, validity }) => {
1508 if matches!(encoding, Encoding::Delta) && !raw {
1509 let mut payload = Vec::with_capacity(4 + validity.len() + data.len() * 8);
1510 payload.extend_from_slice(&(validity.len() as u32).to_be_bytes());
1511 payload.extend_from_slice(validity);
1512 let mut deltas = Vec::with_capacity(data.len());
1513 let mut prev = 0i64;
1514 for v in data {
1515 deltas.push(v - prev);
1516 prev = *v;
1517 }
1518 if le {
1519 append_i64_le(&mut payload, &deltas);
1520 } else {
1521 append_i64_be(&mut payload, &deltas);
1522 }
1523 compress_delta_payload(&payload, compress, le)?
1524 } else {
1525 native_plain_page(validity, compress, raw, le, |p| {
1526 if le {
1527 append_i64_le(p, data);
1528 } else {
1529 append_i64_be(p, data);
1530 }
1531 })
1532 }
1533 }
1534 (
1535 TypeId::Float64,
1536 NativeColumn::Float64 {
1537 data: fdata,
1538 validity,
1539 },
1540 ) => native_plain_page(validity, compress, raw, le, |p| {
1541 let bits: &[u64] = bytemuck::cast_slice::<f64, u64>(fdata);
1542 if le {
1543 append_u64_le(p, bits);
1544 } else {
1545 append_u64_be(p, bits);
1546 }
1547 }),
1548 (
1549 TypeId::Bool,
1550 NativeColumn::Bool {
1551 data: bdata,
1552 validity,
1553 },
1554 ) => native_plain_page(validity, compress, raw, le, |p| p.extend_from_slice(bdata)),
1555 (
1556 TypeId::Bytes,
1557 NativeColumn::Bytes {
1558 offsets,
1559 values,
1560 validity,
1561 },
1562 ) => {
1563 if matches!(encoding, Encoding::Dictionary) && !raw {
1564 let dict = dict_encode_bytes_native(offsets, values, validity);
1565 compress_dict_payload(&dict, compress, le)?
1566 } else {
1567 native_plain_page(validity, compress, raw, le, |p| {
1568 let offs: Vec<u64> = offsets.iter().map(|o| *o as u64).collect();
1569 if le {
1570 append_u64_le(p, &offs);
1571 } else {
1572 append_u64_be(p, &offs);
1573 }
1574 p.extend_from_slice(values);
1575 })
1576 }
1577 }
1578 _ => {
1579 return Err(MongrelError::InvalidArgument(format!(
1580 "encode_page_native: unsupported (ty={ty:?})"
1581 )))
1582 }
1583 })
1584}
1585
1586fn compress_delta_payload(payload: &[u8], compress: Compress, le: bool) -> Result<Vec<u8>> {
1590 Ok(match compress {
1591 Compress::Plain => {
1592 let mut out = vec![algo_with_le(ALGO_PLAIN, le)];
1593 out.extend_from_slice(payload);
1594 out
1595 }
1596 Compress::Zstd(level) => {
1597 let mut out = vec![algo_with_le(ALGO_ZSTD_DELTA, le)];
1598 out.extend(zstd_compress_level(payload, level)?);
1599 out
1600 }
1601 Compress::Lz4 => {
1602 let mut out = vec![algo_with_le(ALGO_LZ4_DELTA, le)];
1603 out.extend(lz4_compress(payload));
1604 out
1605 }
1606 })
1607}
1608
1609fn compress_dict_payload(payload: &[u8], compress: Compress, _le: bool) -> Result<Vec<u8>> {
1614 Ok(match compress {
1615 Compress::Plain => {
1616 let mut out = vec![ALGO_PLAIN];
1617 out.extend_from_slice(payload);
1618 out
1619 }
1620 Compress::Zstd(level) => {
1621 let mut out = vec![ALGO_ZSTD_DICT];
1622 out.extend(zstd_compress_level(payload, level)?);
1623 out
1624 }
1625 Compress::Lz4 => {
1626 let mut out = vec![ALGO_LZ4_DICT];
1627 out.extend(lz4_compress(payload));
1628 out
1629 }
1630 })
1631}
1632
1633fn native_plain_page(
1638 validity: &[u8],
1639 compress: Compress,
1640 raw: bool,
1641 le: bool,
1642 fill_payload: impl FnOnce(&mut Vec<u8>),
1643) -> Vec<u8> {
1644 let mut payload = Vec::new();
1645 payload.extend_from_slice(&(validity.len() as u32).to_be_bytes());
1646 payload.extend_from_slice(validity);
1647 fill_payload(&mut payload);
1648 if raw {
1649 let mut out = vec![algo_with_le(ALGO_PLAIN, le)];
1650 out.extend(payload);
1651 out
1652 } else {
1653 match compress {
1654 Compress::Zstd(level) => {
1655 let mut out = vec![algo_with_le(ALGO_ZSTD_PLAIN, le)];
1656 out.extend(zstd_compress_level(&payload, level).expect("zstd compress"));
1657 out
1658 }
1659 Compress::Lz4 => {
1660 let mut out = vec![algo_with_le(ALGO_LZ4_PLAIN, le)];
1661 out.extend(lz4_compress(&payload));
1662 out
1663 }
1664 Compress::Plain => {
1665 let mut out = vec![algo_with_le(ALGO_PLAIN, le)];
1666 out.extend(payload);
1667 out
1668 }
1669 }
1670 }
1671}
1672
1673pub fn decode_page_native(ty: TypeId, page: &[u8], n: usize) -> Result<NativeColumn> {
1678 use std::borrow::Cow;
1679 if page.is_empty() {
1680 return Err(MongrelError::InvalidArgument("empty page".into()));
1681 }
1682 let algo = page[0];
1683 let body = &page[1..];
1684 let le = algo & ALGO_LE_FLAG != 0;
1688 let base = algo & !ALGO_LE_FLAG;
1689 let plain_algo = matches!(base, ALGO_PLAIN | ALGO_ZSTD_PLAIN | ALGO_LZ4_PLAIN);
1690 let delta_algo = matches!(base, ALGO_ZSTD_DELTA | ALGO_LZ4_DELTA);
1691 let dict_algo = matches!(base, ALGO_ZSTD_DICT | ALGO_LZ4_DICT);
1692 if !plain_algo && !delta_algo && !dict_algo {
1693 return Err(MongrelError::InvalidArgument(format!(
1694 "decode_page_native: unsupported algo {algo} for ty {ty:?}"
1695 )));
1696 }
1697 let raw: Cow<[u8]> = match base {
1699 ALGO_PLAIN => Cow::Borrowed(body),
1700 ALGO_ZSTD_PLAIN | ALGO_ZSTD_DELTA | ALGO_ZSTD_DICT => Cow::Owned(zstd_decompress(
1701 body,
1702 max_decompressed_bytes(ty.clone(), n, base),
1703 )?),
1704 ALGO_LZ4_PLAIN | ALGO_LZ4_DELTA | ALGO_LZ4_DICT => Cow::Owned(lz4_decompress(
1705 body,
1706 max_decompressed_bytes(ty.clone(), n, base),
1707 )?),
1708 _ => unreachable!(),
1709 };
1710
1711 if dict_algo {
1714 return if matches!(ty, TypeId::Bytes) {
1715 dict_decode_bytes_native(&raw, n)
1716 } else {
1717 Err(MongrelError::InvalidArgument(format!(
1718 "decode_page_native: dict algo {algo} only valid for Bytes, got {ty:?}"
1719 )))
1720 };
1721 }
1722
1723 let take_i64 = |p: &[u8]| -> Result<Vec<i64>> {
1725 if le {
1726 take_i64_le(p, n)
1727 } else {
1728 take_i64_be(p, n)
1729 }
1730 };
1731 if delta_algo {
1733 if !matches!(ty, TypeId::Int64 | TypeId::TimestampNanos) {
1734 return Err(MongrelError::InvalidArgument(format!(
1735 "decode_page_native: delta algo {algo} only valid for Int64, got {ty:?}"
1736 )));
1737 }
1738 let (validity, p) = split_validity(&raw)?;
1739 let deltas = take_i64(p)?;
1740 let data = delta_prefix_sum_i64(&deltas);
1741 return Ok(NativeColumn::Int64 { data, validity });
1742 }
1743
1744 match ty {
1746 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date64 | TypeId::Time64 => {
1747 let (validity, p) = split_validity(&raw)?;
1748 Ok(NativeColumn::Int64 {
1749 data: take_i64(p)?,
1750 validity,
1751 })
1752 }
1753 TypeId::Float64 => {
1754 let (validity, p) = split_validity(&raw)?;
1755 let bits = if le {
1756 take_u64_le(p, n)?
1757 } else {
1758 take_u64_be(p, n)?
1759 };
1760 let data: Vec<f64> = bits.into_iter().map(f64::from_bits).collect();
1761 Ok(NativeColumn::Float64 { data, validity })
1762 }
1763 TypeId::Bool => {
1764 let (validity, p) = split_validity(&raw)?;
1765 if p.len() < n {
1766 return Err(MongrelError::InvalidArgument(
1767 "bool payload truncated".into(),
1768 ));
1769 }
1770 Ok(NativeColumn::Bool {
1771 data: p[..n].to_vec(),
1772 validity,
1773 })
1774 }
1775 TypeId::Bytes => decode_bytes_plain_payload(&raw, n, le),
1776 _ => Err(MongrelError::InvalidArgument(format!(
1777 "decode_page_native: unsupported ty {ty:?}"
1778 ))),
1779 }
1780}
1781
1782fn split_validity(raw: &[u8]) -> Result<(Vec<u8>, &[u8])> {
1783 if raw.len() < 4 {
1784 return Err(MongrelError::InvalidArgument("page validity header".into()));
1785 }
1786 let vlen = u32::from_be_bytes([raw[0], raw[1], raw[2], raw[3]]) as usize;
1787 if 4 + vlen > raw.len() {
1788 return Err(MongrelError::InvalidArgument("page validity range".into()));
1789 }
1790 Ok((raw[4..4 + vlen].to_vec(), &raw[4 + vlen..]))
1791}
1792
1793fn decode_bytes_plain_payload(raw: &[u8], n: usize, le: bool) -> Result<NativeColumn> {
1796 let (validity, p) = split_validity(raw)?;
1797 let table = (n + 1) * 8;
1798 if p.len() < table {
1799 return Err(MongrelError::InvalidArgument(
1800 "bytes offsets truncated".into(),
1801 ));
1802 }
1803 let offsets_be: Vec<u64> = if le {
1804 take_u64_le(p, n + 1)?
1805 } else {
1806 take_u64_be(p, n + 1)?
1807 };
1808 let offsets: Vec<u32> = offsets_be.into_iter().map(|o| o as u32).collect();
1809 let values = p[table..].to_vec();
1810 Ok(NativeColumn::Bytes {
1811 offsets,
1812 values,
1813 validity,
1814 })
1815}
1816
1817fn take_i64_be(p: &[u8], n: usize) -> Result<Vec<i64>> {
1818 if p.len() < n * 8 {
1819 return Err(MongrelError::InvalidArgument(
1820 "int64 payload truncated".into(),
1821 ));
1822 }
1823 Ok(take_u64_be(p, n)?.into_iter().map(|u| u as i64).collect())
1824}
1825
1826fn delta_prefix_sum_i64(deltas: &[i64]) -> Vec<i64> {
1834 let mut out = vec![0i64; deltas.len()];
1835 prefix_sum_i64_into(deltas, &mut out);
1836 out
1837}
1838
1839fn prefix_sum_i64_into(deltas: &[i64], out: &mut [i64]) {
1840 debug_assert_eq!(deltas.len(), out.len());
1841 #[cfg(target_arch = "x86_64")]
1842 {
1843 if is_x86_feature_detected!("avx2") && deltas.len() >= 4 {
1844 unsafe {
1847 prefix_sum_avx2(deltas, out);
1848 }
1849 return;
1850 }
1851 }
1852 prefix_sum_scalar(deltas, out);
1853}
1854
1855fn prefix_sum_scalar(deltas: &[i64], out: &mut [i64]) {
1856 let mut acc = 0i64;
1857 for (i, &d) in deltas.iter().enumerate() {
1858 acc = acc.wrapping_add(d);
1859 out[i] = acc;
1860 }
1861}
1862
1863#[cfg(target_arch = "x86_64")]
1864#[target_feature(enable = "avx2")]
1865unsafe fn prefix_sum_avx2(deltas: &[i64], out: &mut [i64]) {
1866 use std::arch::x86_64::*;
1867 let n = deltas.len();
1868 let mut running = _mm256_setzero_si256(); let mut i = 0usize;
1870 while i + 4 <= n {
1871 let mut x = _mm256_loadu_si256(deltas.as_ptr().add(i) as *const __m256i);
1872 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);
1881 running = _mm256_permute4x64_epi64(x, 0xFF); i += 4;
1883 }
1884 let mut acc = if i == 0 { 0 } else { out[i - 1] };
1887 while i < n {
1888 acc = acc.wrapping_add(deltas[i]);
1889 out[i] = acc;
1890 i += 1;
1891 }
1892}
1893
1894fn append_u64_be(out: &mut Vec<u8>, data: &[u64]) {
1896 if cfg!(target_endian = "little") {
1897 let swapped: Vec<u64> = data.iter().map(|v| v.swap_bytes()).collect();
1898 out.extend_from_slice(bytemuck::cast_slice::<u64, u8>(&swapped));
1899 } else {
1900 out.extend_from_slice(bytemuck::cast_slice::<u64, u8>(data));
1901 }
1902}
1903
1904fn append_i64_be(out: &mut Vec<u8>, data: &[i64]) {
1906 if cfg!(target_endian = "little") {
1907 let swapped: Vec<i64> = data.iter().map(|v| v.swap_bytes()).collect();
1908 out.extend_from_slice(bytemuck::cast_slice::<i64, u8>(&swapped));
1909 } else {
1910 out.extend_from_slice(bytemuck::cast_slice::<i64, u8>(data));
1911 }
1912}
1913
1914fn take_u64_be(p: &[u8], n: usize) -> Result<Vec<u64>> {
1918 if p.len() < n * 8 {
1919 return Err(MongrelError::InvalidArgument(
1920 "u64 payload truncated".into(),
1921 ));
1922 }
1923 let bytes = &p[..n * 8];
1924 if let Ok(native) = bytemuck::try_cast_slice::<u8, u64>(bytes) {
1925 Ok(native.iter().map(|v| v.swap_bytes()).collect())
1926 } else {
1927 let mut out = Vec::with_capacity(n);
1928 for chunk in bytes.chunks_exact(8) {
1929 out.push(u64::from_be_bytes(chunk.try_into().unwrap()));
1930 }
1931 Ok(out)
1932 }
1933}
1934
1935fn append_u64_le(out: &mut Vec<u8>, data: &[u64]) {
1939 if cfg!(target_endian = "little") {
1940 out.extend_from_slice(bytemuck::cast_slice::<u64, u8>(data));
1941 } else {
1942 let swapped: Vec<u64> = data.iter().map(|v| v.swap_bytes()).collect();
1943 out.extend_from_slice(bytemuck::cast_slice::<u64, u8>(&swapped));
1944 }
1945}
1946
1947fn append_i64_le(out: &mut Vec<u8>, data: &[i64]) {
1949 if cfg!(target_endian = "little") {
1950 out.extend_from_slice(bytemuck::cast_slice::<i64, u8>(data));
1951 } else {
1952 let swapped: Vec<i64> = data.iter().map(|v| v.swap_bytes()).collect();
1953 out.extend_from_slice(bytemuck::cast_slice::<i64, u8>(&swapped));
1954 }
1955}
1956
1957fn take_u64_le(p: &[u8], n: usize) -> Result<Vec<u64>> {
1963 if p.len() < n * 8 {
1964 return Err(MongrelError::InvalidArgument(
1965 "u64 payload truncated".into(),
1966 ));
1967 }
1968 let bytes = &p[..n * 8];
1969 if cfg!(target_endian = "little") {
1970 if let Ok(native) = bytemuck::try_cast_slice::<u8, u64>(bytes) {
1971 return Ok(native.to_vec());
1973 }
1974 Ok(bytes
1975 .chunks_exact(8)
1976 .map(|c| u64::from_le_bytes(c.try_into().unwrap()))
1977 .collect())
1978 } else {
1979 Ok(bytes
1980 .chunks_exact(8)
1981 .map(|c| u64::from_le_bytes(c.try_into().unwrap()))
1982 .collect())
1983 }
1984}
1985
1986fn take_i64_le(p: &[u8], n: usize) -> Result<Vec<i64>> {
1988 if p.len() < n * 8 {
1989 return Err(MongrelError::InvalidArgument(
1990 "int64 payload truncated".into(),
1991 ));
1992 }
1993 Ok(take_u64_le(p, n)?.into_iter().map(|u| u as i64).collect())
1994}
1995
1996fn dict_encode_bytes_native(offsets: &[u32], values: &[u8], validity: &[u8]) -> Vec<u8> {
1997 let n = offsets.len() - 1;
1998 let mut table: Vec<Vec<u8>> = Vec::new();
1999 let mut index_of: std::collections::HashMap<&[u8], u32> = std::collections::HashMap::new();
2000 let mut indices: Vec<u32> = Vec::with_capacity(n);
2001 for i in 0..n {
2002 let lo = offsets[i] as usize;
2003 let hi = offsets[i + 1] as usize;
2004 let slice = &values[lo..hi];
2005 let idx = if let Some(&idx) = index_of.get(slice) {
2006 idx
2007 } else {
2008 let idx = table.len() as u32;
2009 index_of.insert(slice, idx);
2010 table.push(slice.to_vec());
2011 idx
2012 };
2013 indices.push(idx);
2014 }
2015 let mut out = Vec::new();
2016 out.extend_from_slice(&(validity.len() as u32).to_be_bytes());
2017 out.extend_from_slice(validity);
2018 out.extend_from_slice(&(indices.len() as u32).to_be_bytes());
2019 for i in &indices {
2020 out.extend_from_slice(&i.to_be_bytes());
2021 }
2022 out.extend_from_slice(&(table.len() as u32).to_be_bytes());
2023 for entry in &table {
2024 out.extend_from_slice(&(entry.len() as u32).to_be_bytes());
2025 out.extend_from_slice(entry);
2026 }
2027 out
2028}
2029
2030fn dict_decode_bytes_native(data: &[u8], n: usize) -> Result<NativeColumn> {
2031 let mut cur = 0usize;
2032 let vlen = read_u32_be(data, &mut cur)? as usize;
2033 let validity = checked_slice(data, &mut cur, vlen)?.to_vec();
2034 let index_count = read_u32_be(data, &mut cur)? as usize;
2035 if index_count < n {
2036 return Err(MongrelError::InvalidArgument("dict index_count < n".into()));
2037 }
2038 let mut indices = Vec::with_capacity(index_count.min(n));
2039 for _ in 0..index_count {
2040 indices.push(read_u32_be(data, &mut cur)?);
2041 }
2042 let table_count = read_u32_be(data, &mut cur)? as usize;
2043 let mut table: Vec<(usize, usize)> = Vec::with_capacity(table_count); let mut values = Vec::new();
2045 for _ in 0..table_count {
2046 let len = read_u32_be(data, &mut cur)? as usize;
2047 let chunk = checked_slice(data, &mut cur, len)?;
2048 let start = values.len();
2049 values.extend_from_slice(chunk);
2050 table.push((start, len));
2051 }
2052 let mut merged = Vec::new();
2053 let mut offs = vec![0u32];
2054 for (i, &idx) in indices.iter().enumerate().take(n) {
2055 let non_null = (validity.get(i / 8).copied().unwrap_or(0) >> (i % 8)) & 1 == 1;
2059 if non_null {
2060 let (start, len) = table
2061 .get(idx as usize)
2062 .copied()
2063 .ok_or_else(|| MongrelError::InvalidArgument("dict index out of range".into()))?;
2064 merged.extend_from_slice(&values[start..start + len]);
2065 }
2066 offs.push(merged.len() as u32);
2067 }
2068 Ok(NativeColumn::Bytes {
2069 offsets: offs,
2070 values: merged,
2071 validity,
2072 })
2073}
2074
2075#[cfg(test)]
2076mod native_tests {
2077 use super::*;
2078
2079 #[test]
2080 fn native_int64_plain_round_trip() {
2081 let col = NativeColumn::Int64 {
2082 data: (0..1000).collect(),
2083 validity: full_validity(1000),
2084 };
2085 let page = encode_page_native(
2086 TypeId::Int64,
2087 &col,
2088 Encoding::Zstd,
2089 Compress::Zstd(3),
2090 false,
2091 )
2092 .unwrap();
2093 let back = decode_page_native(TypeId::Int64, &page, 1000).unwrap();
2094 match back {
2095 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2096 _ => panic!(),
2097 }
2098 }
2099
2100 #[test]
2101 fn native_int64_delta_crushes_sequential() {
2102 let col = NativeColumn::int64_sequence(0, 100_000);
2103 let plain = encode_page_native(
2104 TypeId::Int64,
2105 &col,
2106 Encoding::Zstd,
2107 Compress::Zstd(3),
2108 false,
2109 )
2110 .unwrap();
2111 let delta = encode_page_native(
2112 TypeId::Int64,
2113 &col,
2114 Encoding::Delta,
2115 Compress::Zstd(3),
2116 false,
2117 )
2118 .unwrap();
2119 assert!(
2120 delta.len() < plain.len() / 5,
2121 "delta must crush sequential ints"
2122 );
2123 let back = decode_page_native(TypeId::Int64, &delta, 100_000).unwrap();
2124 match back {
2125 NativeColumn::Int64 { data, .. } => assert_eq!(data.len(), 100_000),
2126 _ => panic!(),
2127 }
2128 }
2129
2130 #[test]
2131 fn native_bytes_dict_round_trip() {
2132 let n = 500;
2133 let mut offsets = vec![0u32];
2134 let mut values = Vec::new();
2135 for i in 0..n {
2136 let s = ["red", "green", "blue"][i % 3];
2137 values.extend_from_slice(s.as_bytes());
2138 offsets.push(values.len() as u32);
2139 }
2140 let col = NativeColumn::Bytes {
2141 offsets,
2142 values,
2143 validity: full_validity(n),
2144 };
2145 let page = encode_page_native(
2146 TypeId::Bytes,
2147 &col,
2148 Encoding::Dictionary,
2149 Compress::Zstd(3),
2150 false,
2151 )
2152 .unwrap();
2153 assert!(page.len() < 100, "dict page tiny, got {}", page.len());
2154 let back = decode_page_native(TypeId::Bytes, &page, n).unwrap();
2155 assert_eq!(back.len(), n);
2156 }
2157
2158 #[test]
2159 fn native_gather_picks_indices() {
2160 let col = NativeColumn::Int64 {
2161 data: vec![10, 20, 30, 40],
2162 validity: full_validity(4),
2163 };
2164 let g = col.gather(&[0, 2, 3]);
2165 match g {
2166 NativeColumn::Int64 { data, .. } => assert_eq!(data, vec![10, 30, 40]),
2167 _ => panic!(),
2168 }
2169 }
2170
2171 #[test]
2174 fn native_plain_no_zstd_round_trips_all_types() {
2175 let i = NativeColumn::Int64 {
2176 data: (0..1000).collect(),
2177 validity: full_validity(1000),
2178 };
2179 let p =
2180 encode_page_native(TypeId::Int64, &i, Encoding::Plain, Compress::Plain, false).unwrap();
2181 assert_eq!(p[0], ALGO_PLAIN, "Int64 plain must be ALGO_PLAIN");
2182 match decode_page_native(TypeId::Int64, &p, 1000).unwrap() {
2183 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2184 _ => panic!(),
2185 }
2186
2187 let f = NativeColumn::Float64 {
2188 data: (0..500).map(|x| x as f64 * 1.5).collect(),
2189 validity: full_validity(500),
2190 };
2191 let p = encode_page_native(TypeId::Float64, &f, Encoding::Plain, Compress::Plain, false)
2192 .unwrap();
2193 assert_eq!(p[0], ALGO_PLAIN);
2194 match decode_page_native(TypeId::Float64, &p, 500).unwrap() {
2195 NativeColumn::Float64 { data, .. } => {
2196 assert_eq!(data, (0..500).map(|x| x as f64 * 1.5).collect::<Vec<_>>())
2197 }
2198 _ => panic!(),
2199 }
2200
2201 let b = NativeColumn::Bool {
2202 data: (0..64).map(|i| (i % 2) as u8).collect(),
2203 validity: full_validity(64),
2204 };
2205 let p =
2206 encode_page_native(TypeId::Bool, &b, Encoding::Plain, Compress::Plain, false).unwrap();
2207 assert_eq!(p[0], ALGO_PLAIN);
2208 match decode_page_native(TypeId::Bool, &p, 64).unwrap() {
2209 NativeColumn::Bool { data, .. } => {
2210 assert_eq!(data, (0..64).map(|i| (i % 2) as u8).collect::<Vec<_>>())
2211 }
2212 _ => panic!(),
2213 }
2214
2215 let mut offsets = vec![0u32];
2216 let mut values = Vec::new();
2217 for i in 0..200u32 {
2218 values.extend_from_slice(format!("v{i}").as_bytes());
2219 offsets.push(values.len() as u32);
2220 }
2221 let s = NativeColumn::Bytes {
2222 offsets,
2223 values,
2224 validity: full_validity(200),
2225 };
2226 let p =
2227 encode_page_native(TypeId::Bytes, &s, Encoding::Plain, Compress::Plain, false).unwrap();
2228 assert_eq!(p[0], ALGO_PLAIN);
2229 match decode_page_native(TypeId::Bytes, &p, 200).unwrap() {
2230 NativeColumn::Bytes {
2231 offsets: o,
2232 values: v,
2233 ..
2234 } => {
2235 assert_eq!(o.len(), 201);
2236 for i in 0..200 {
2237 let lo = o[i] as usize;
2238 let hi = o[i + 1] as usize;
2239 assert_eq!(&v[lo..hi], format!("v{i}").as_bytes());
2240 }
2241 }
2242 _ => panic!(),
2243 }
2244 }
2245
2246 #[test]
2249 fn lz4_pages_round_trip_all_types() {
2250 let i = NativeColumn::int64_sequence(0, 1000);
2252 let p =
2253 encode_page_native(TypeId::Int64, &i, Encoding::Delta, Compress::Lz4, false).unwrap();
2254 assert_eq!(p[0], ALGO_LZ4_DELTA);
2255 match decode_page_native(TypeId::Int64, &p, 1000).unwrap() {
2256 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2257 _ => panic!(),
2258 }
2259 let p =
2261 encode_page_native(TypeId::Int64, &i, Encoding::Zstd, Compress::Lz4, false).unwrap();
2262 assert_eq!(p[0], ALGO_LZ4_PLAIN);
2263 match decode_page_native(TypeId::Int64, &p, 1000).unwrap() {
2264 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2265 _ => panic!(),
2266 }
2267 let f = NativeColumn::Float64 {
2269 data: (0..500).map(|x| x as f64 * 1.5).collect(),
2270 validity: full_validity(500),
2271 };
2272 let p =
2273 encode_page_native(TypeId::Float64, &f, Encoding::Zstd, Compress::Lz4, false).unwrap();
2274 assert_eq!(p[0], ALGO_LZ4_PLAIN);
2275 match decode_page_native(TypeId::Float64, &p, 500).unwrap() {
2276 NativeColumn::Float64 { data, .. } => {
2277 assert_eq!(data, (0..500).map(|x| x as f64 * 1.5).collect::<Vec<_>>())
2278 }
2279 _ => panic!(),
2280 }
2281 let mut offsets = vec![0u32];
2283 let mut values = Vec::new();
2284 for i in 0..300u32 {
2285 values.extend_from_slice(["red", "green", "blue"][(i % 3) as usize].as_bytes());
2286 offsets.push(values.len() as u32);
2287 }
2288 let s = NativeColumn::Bytes {
2289 offsets,
2290 values,
2291 validity: full_validity(300),
2292 };
2293 let p = encode_page_native(
2294 TypeId::Bytes,
2295 &s,
2296 Encoding::Dictionary,
2297 Compress::Lz4,
2298 false,
2299 )
2300 .unwrap();
2301 assert_eq!(p[0], ALGO_LZ4_DICT);
2302 match decode_page_native(TypeId::Bytes, &p, 300).unwrap() {
2303 NativeColumn::Bytes { offsets: o, .. } => assert_eq!(o.len(), 301),
2304 _ => panic!(),
2305 }
2306 }
2307
2308 #[test]
2312 fn le_pages_round_trip_all_types() {
2313 let assert_le = |page: &[u8]| assert_ne!(page[0] & ALGO_LE_FLAG, 0, "LE flag must be set");
2314
2315 let i = NativeColumn::Int64 {
2317 data: (0..1000).collect(),
2318 validity: full_validity(1000),
2319 };
2320 let p =
2321 encode_page_native(TypeId::Int64, &i, Encoding::Plain, Compress::Plain, true).unwrap();
2322 assert_eq!(p[0], ALGO_LE_FLAG, "raw LE Int64 algo = flag only");
2323 match decode_page_native(TypeId::Int64, &p, 1000).unwrap() {
2324 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2325 _ => panic!(),
2326 }
2327
2328 let p =
2330 encode_page_native(TypeId::Int64, &i, Encoding::Zstd, Compress::Zstd(3), true).unwrap();
2331 assert_le(&p);
2332 match decode_page_native(TypeId::Int64, &p, 1000).unwrap() {
2333 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2334 _ => panic!(),
2335 }
2336
2337 let seq = NativeColumn::int64_sequence(0, 1000);
2339 let p =
2340 encode_page_native(TypeId::Int64, &seq, Encoding::Delta, Compress::Lz4, true).unwrap();
2341 assert_eq!(p[0], ALGO_LZ4_DELTA | ALGO_LE_FLAG);
2342 match decode_page_native(TypeId::Int64, &p, 1000).unwrap() {
2343 NativeColumn::Int64 { data, .. } => assert_eq!(data, (0..1000).collect::<Vec<_>>()),
2344 _ => panic!(),
2345 }
2346
2347 let f = NativeColumn::Float64 {
2349 data: (0..500).map(|x| x as f64 * 1.5).collect(),
2350 validity: full_validity(500),
2351 };
2352 let p = encode_page_native(TypeId::Float64, &f, Encoding::Zstd, Compress::Zstd(3), true)
2353 .unwrap();
2354 assert_le(&p);
2355 match decode_page_native(TypeId::Float64, &p, 500).unwrap() {
2356 NativeColumn::Float64 { data, .. } => {
2357 assert_eq!(data, (0..500).map(|x| x as f64 * 1.5).collect::<Vec<_>>())
2358 }
2359 _ => panic!(),
2360 }
2361 }
2362
2363 #[test]
2367 fn malicious_lz4_size_prefix_is_rejected() {
2368 let mut evil = vec![ALGO_LZ4_PLAIN];
2370 evil.extend_from_slice(&0xFFFF_FFFFu32.to_le_bytes());
2371 evil.extend_from_slice(b"junk");
2372 let err =
2373 decode_page_native(TypeId::Int64, &evil, 1000).expect_err("must reject oversize lz4");
2374 let msg = format!("{err}");
2375 assert!(msg.contains("exceeds page limit"), "got: {msg}");
2376
2377 let err = decode_page(TypeId::Int64, &evil, 1000).expect_err("must reject oversize lz4");
2379 assert!(format!("{err}").contains("exceeds page limit"));
2380 }
2381
2382 #[test]
2386 fn truncated_dict_payload_does_not_panic() {
2387 let mut body = vec![ALGO_ZSTD_DICT];
2390 body.extend_from_slice(&zstd_compress(&[0u8; 2]).unwrap()[..]);
2391 decode_page(TypeId::Bytes, &body, 4).expect_err("value dict trunc must Err");
2392
2393 let truncated: &[u8] = &[
2395 0x00, 0x00, 0x00, 0x10, ];
2397 dict_decode_bytes_native(truncated, 2).expect_err("native dict trunc must Err");
2398 dict_decode_bytes(truncated, 2).expect_err("value dict trunc must Err");
2399
2400 let mut bad = vec![0u8; 0];
2403 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");
2409 dict_decode_bytes_native(&bad, 1).expect_err("oob index must Err");
2410 }
2411
2412 #[test]
2416 fn delta_prefix_sum_matches_scalar_all_lengths() {
2417 let mut deltas: Vec<i64> = Vec::with_capacity(2000);
2419 let mut s: u64 = 0x9E37_79B9_7F4A_7C15;
2420 for _ in 0..2000 {
2421 s = s
2422 .wrapping_mul(6364136223846793005)
2423 .wrapping_add(1442695040888963407);
2424 deltas.push((s as i32) as i64); }
2426 for &len in &[
2427 0usize, 1, 2, 3, 4, 5, 6, 7, 8, 9, 15, 16, 17, 31, 32, 33, 63, 64, 65, 127, 128, 129,
2428 1000,
2429 ] {
2430 let input = &deltas[..len];
2431 let simd = delta_prefix_sum_i64(input);
2432 let mut ref_out = vec![0i64; len];
2433 prefix_sum_scalar(input, &mut ref_out);
2434 assert_eq!(simd, ref_out, "len {len}: SIMD prefix sum diverged");
2435 }
2436
2437 let mut ids = vec![0i64; 333];
2439 for x in ids.iter_mut().skip(1) {
2440 *x = 1;
2441 }
2442 let got = delta_prefix_sum_i64(&ids);
2443 assert_eq!(got, (0..333).map(|x| x as i64).collect::<Vec<_>>());
2444 }
2445
2446 #[test]
2449 fn delta_page_decodes_with_null_via_vectorized_path() {
2450 let mut validity = full_validity(17);
2451 validity[9 / 8] &= !(1 << (9 % 8)); let col = NativeColumn::Int64 {
2453 data: (0..17).collect(),
2454 validity: validity.clone(),
2455 };
2456 for (name, comp) in [("zstd", Compress::Zstd(3)), ("lz4", Compress::Lz4)] {
2457 let page =
2458 encode_page_native(TypeId::Int64, &col, Encoding::Delta, comp, false).unwrap();
2459 assert!(matches!(page[0], ALGO_ZSTD_DELTA | ALGO_LZ4_DELTA));
2461 let back = decode_page_native(TypeId::Int64, &page, 17).unwrap();
2462 match back {
2463 NativeColumn::Int64 { data, validity: v } => {
2464 assert_eq!(data, (0..17).collect::<Vec<_>>(), "{name} data");
2465 assert_eq!(v, validity, "{name} validity");
2466 }
2467 _ => panic!(),
2468 }
2469 }
2470 }
2471}