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