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