1use std::io::{self, Write};
13
14use byteorder::{LittleEndian, WriteBytesExt};
15
16use super::{bitpack_read, bitpack_write, bits_needed_u64};
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22#[repr(u8)]
23pub enum CodecType {
24 Constant = 0,
25 Bitpacked = 1,
26 Linear = 2,
27 BlockwiseLinear = 3,
28}
29
30impl CodecType {
31 pub fn from_u8(v: u8) -> Option<Self> {
32 match v {
33 0 => Some(Self::Constant),
34 1 => Some(Self::Bitpacked),
35 2 => Some(Self::Linear),
36 3 => Some(Self::BlockwiseLinear),
37 _ => None,
38 }
39 }
40}
41
42pub const BLOCKWISE_LINEAR_BLOCK_SIZE: usize = 512;
44
45pub fn validate_auto(data: &[u8], expected_values: usize) -> io::Result<()> {
49 let (&codec_id, rest) = data.split_first().ok_or_else(|| {
50 io::Error::new(io::ErrorKind::UnexpectedEof, "fast field codec is missing")
51 })?;
52
53 let packed_len = |count: usize, bpv: u8| -> io::Result<usize> {
54 if bpv > 64 {
55 return Err(io::Error::new(
56 io::ErrorKind::InvalidData,
57 "fast field bit width exceeds 64",
58 ));
59 }
60 count
61 .checked_mul(bpv as usize)
62 .and_then(|bits| bits.checked_add(7))
63 .map(|bits| bits / 8)
64 .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "fast field size overflow"))
65 };
66
67 match CodecType::from_u8(codec_id) {
68 Some(CodecType::Constant) => {
69 if rest.len() != 8 {
70 return Err(io::Error::new(
71 io::ErrorKind::InvalidData,
72 "invalid constant fast field length",
73 ));
74 }
75 }
76 Some(CodecType::Bitpacked) => {
77 if rest.len() < 9 {
78 return Err(io::Error::new(
79 io::ErrorKind::UnexpectedEof,
80 "bitpacked fast field header is truncated",
81 ));
82 }
83 let expected_len = 9usize
84 .checked_add(packed_len(expected_values, rest[8])?)
85 .ok_or_else(|| {
86 io::Error::new(io::ErrorKind::InvalidData, "fast field size overflow")
87 })?;
88 if rest.len() != expected_len {
89 return Err(io::Error::new(
90 io::ErrorKind::InvalidData,
91 "bitpacked fast field length is inconsistent",
92 ));
93 }
94 }
95 Some(CodecType::Linear) => {
96 if rest.len() < 29 {
97 return Err(io::Error::new(
98 io::ErrorKind::UnexpectedEof,
99 "linear fast field header is truncated",
100 ));
101 }
102 let count = u32::from_le_bytes(rest[16..20].try_into().unwrap()) as usize;
103 if count != expected_values || count < 2 {
104 return Err(io::Error::new(
105 io::ErrorKind::InvalidData,
106 "linear fast field value count is inconsistent",
107 ));
108 }
109 let expected_len = 29usize
110 .checked_add(packed_len(count, rest[28])?)
111 .ok_or_else(|| {
112 io::Error::new(io::ErrorKind::InvalidData, "fast field size overflow")
113 })?;
114 if rest.len() != expected_len {
115 return Err(io::Error::new(
116 io::ErrorKind::InvalidData,
117 "linear fast field length is inconsistent",
118 ));
119 }
120 }
121 Some(CodecType::BlockwiseLinear) => {
122 if rest.len() < 8 {
123 return Err(io::Error::new(
124 io::ErrorKind::UnexpectedEof,
125 "blockwise fast field header is truncated",
126 ));
127 }
128 let count = u32::from_le_bytes(rest[0..4].try_into().unwrap()) as usize;
129 let num_blocks = u32::from_le_bytes(rest[4..8].try_into().unwrap()) as usize;
130 if count != expected_values || num_blocks != count.div_ceil(BLOCKWISE_LINEAR_BLOCK_SIZE)
131 {
132 return Err(io::Error::new(
133 io::ErrorKind::InvalidData,
134 "blockwise fast field counts are inconsistent",
135 ));
136 }
137
138 let mut pos = 8usize;
139 for block_idx in 0..num_blocks {
140 let header_end = pos.checked_add(29).ok_or_else(|| {
141 io::Error::new(io::ErrorKind::InvalidData, "fast field offset overflow")
142 })?;
143 if header_end > rest.len() {
144 return Err(io::Error::new(
145 io::ErrorKind::UnexpectedEof,
146 "blockwise fast field block header is truncated",
147 ));
148 }
149 let bpv = rest[pos + 24];
150 let declared =
151 u32::from_le_bytes(rest[pos + 25..header_end].try_into().unwrap()) as usize;
152 let block_start = block_idx * BLOCKWISE_LINEAR_BLOCK_SIZE;
153 let block_count = (count - block_start).min(BLOCKWISE_LINEAR_BLOCK_SIZE);
154 let expected = packed_len(block_count, bpv)?;
155 if declared != expected {
156 return Err(io::Error::new(
157 io::ErrorKind::InvalidData,
158 "blockwise fast field packed length is inconsistent",
159 ));
160 }
161 pos = header_end.checked_add(declared).ok_or_else(|| {
162 io::Error::new(io::ErrorKind::InvalidData, "fast field offset overflow")
163 })?;
164 if pos > rest.len() {
165 return Err(io::Error::new(
166 io::ErrorKind::UnexpectedEof,
167 "blockwise fast field data is truncated",
168 ));
169 }
170 }
171 if pos != rest.len() {
172 return Err(io::Error::new(
173 io::ErrorKind::InvalidData,
174 "blockwise fast field contains trailing data",
175 ));
176 }
177 }
178 None => {
179 return Err(io::Error::new(
180 io::ErrorKind::InvalidData,
181 "unknown fast field codec",
182 ));
183 }
184 }
185 Ok(())
186}
187
188pub trait CodecEstimator {
195 fn collect(&mut self, value: u64);
196 fn finalize(&mut self) {}
197 fn estimate(&self) -> Option<u64>;
198 fn serialize(&self, values: &[u64], writer: &mut dyn Write) -> io::Result<u64>;
199}
200
201#[derive(Default)]
205pub struct ConstantEstimator {
206 first: Option<u64>,
207 all_same: bool,
208}
209
210impl CodecEstimator for ConstantEstimator {
211 fn collect(&mut self, value: u64) {
212 match self.first {
213 None => {
214 self.first = Some(value);
215 self.all_same = true;
216 }
217 Some(f) => {
218 if value != f {
219 self.all_same = false;
220 }
221 }
222 }
223 }
224
225 fn estimate(&self) -> Option<u64> {
226 if self.all_same {
227 Some(9)
229 } else {
230 None
231 }
232 }
233
234 fn serialize(&self, values: &[u64], writer: &mut dyn Write) -> io::Result<u64> {
235 let val = if values.is_empty() { 0 } else { values[0] };
236 writer.write_u8(CodecType::Constant as u8)?;
237 writer.write_u64::<LittleEndian>(val)?;
238 Ok(9)
239 }
240}
241
242#[derive(Default)]
246pub struct BitpackedEstimator {
247 min: u64,
248 max: u64,
249 count: usize,
250 initialized: bool,
251}
252
253impl CodecEstimator for BitpackedEstimator {
254 fn collect(&mut self, value: u64) {
255 if !self.initialized {
256 self.min = value;
257 self.max = value;
258 self.initialized = true;
259 } else {
260 self.min = self.min.min(value);
261 self.max = self.max.max(value);
262 }
263 self.count += 1;
264 }
265
266 fn estimate(&self) -> Option<u64> {
267 if self.count == 0 {
268 return Some(0);
269 }
270 let range = self.max - self.min;
271 let bpv = bits_needed_u64(range) as u64;
272 let data_bits = self.count as u64 * bpv;
274 let data_bytes = data_bits.div_ceil(8);
275 Some(1 + 8 + 1 + data_bytes)
276 }
277
278 fn serialize(&self, values: &[u64], writer: &mut dyn Write) -> io::Result<u64> {
279 let (min_value, bpv) = if values.is_empty() {
280 (0u64, 0u8)
281 } else {
282 let min_val = values.iter().copied().min().unwrap();
283 let max_val = values.iter().copied().max().unwrap();
284 (min_val, bits_needed_u64(max_val - min_val))
285 };
286
287 writer.write_u8(CodecType::Bitpacked as u8)?;
288 writer.write_u64::<LittleEndian>(min_value)?;
289 writer.write_u8(bpv)?;
290 let mut bytes_written = 10u64; if bpv > 0 && !values.is_empty() {
293 let shifted: Vec<u64> = values.iter().map(|&v| v - min_value).collect();
294 let mut packed = Vec::new();
295 bitpack_write(&shifted, bpv, &mut packed);
296 writer.write_all(&packed)?;
297 bytes_written += packed.len() as u64;
298 }
299 Ok(bytes_written)
300 }
301}
302
303#[inline]
307pub fn bitpacked_read(data: &[u8], index: usize) -> u64 {
308 let min_value = u64::from_le_bytes(data[0..8].try_into().unwrap());
309 let bpv = data[8];
310 if bpv == 0 {
311 return min_value;
312 }
313 let packed = &data[9..];
314 bitpack_read(packed, bpv, index).wrapping_add(min_value)
315}
316
317#[derive(Default)]
331pub struct LinearEstimator {
332 count: usize,
333 first: u64,
334 last: u64,
335 min_val: u64,
336 max_val: u64,
337 min_residual: i64,
338 max_residual: i64,
339 values_collected: bool,
340 overflow: bool,
342}
343
344impl CodecEstimator for LinearEstimator {
345 fn collect(&mut self, value: u64) {
346 if !self.values_collected {
347 self.first = value;
348 self.min_val = value;
349 self.max_val = value;
350 self.values_collected = true;
351 } else {
352 self.min_val = self.min_val.min(value);
353 self.max_val = self.max_val.max(value);
354 }
355 self.last = value;
356 self.count += 1;
357 }
358
359 fn finalize(&mut self) {
360 if self.count < 2 {
361 return;
362 }
363 let pred_min = self.first.min(self.last) as i128;
371 let pred_max = self.first.max(self.last) as i128;
372 let min_res = self.min_val as i128 - pred_max;
373 let max_res = self.max_val as i128 - pred_min;
374 if min_res < i64::MIN as i128 || max_res > i64::MAX as i128 {
377 self.overflow = true;
378 return;
379 }
380 self.min_residual = min_res as i64;
381 self.max_residual = max_res as i64;
382 }
383
384 fn estimate(&self) -> Option<u64> {
385 if self.count < 2 || self.overflow {
386 return None;
387 }
388 let range = (self.max_residual as i128 - self.min_residual as i128) as u64;
390 let bpv = bits_needed_u64(range) as u64;
391 let data_bits = self.count as u64 * bpv;
392 let data_bytes = data_bits.div_ceil(8);
393 Some(1 + 8 + 8 + 4 + 8 + 1 + data_bytes)
395 }
396
397 fn serialize(&self, values: &[u64], writer: &mut dyn Write) -> io::Result<u64> {
398 let n = values.len();
399 if n < 2 {
400 return Err(io::Error::new(
401 io::ErrorKind::InvalidInput,
402 "linear needs ≥ 2 values",
403 ));
404 }
405 let first = values[0];
406 let last = values[n - 1];
407
408 let mut min_residual = i128::MAX;
410 for (i, &val) in values.iter().enumerate() {
411 let predicted = interpolate(first, last, n, i);
412 let residual = val as i128 - predicted as i128;
413 min_residual = min_residual.min(residual);
414 }
415
416 if min_residual < i64::MIN as i128 || min_residual > i64::MAX as i128 {
418 return Err(io::Error::new(
419 io::ErrorKind::InvalidInput,
420 "linear codec: residual offset exceeds i64 range",
421 ));
422 }
423 let min_residual_i64 = min_residual as i64;
424
425 let shifted: Vec<u64> = values
427 .iter()
428 .enumerate()
429 .map(|(i, &val)| {
430 let predicted = interpolate(first, last, n, i);
431 let residual = val as i128 - predicted as i128;
432 (residual - min_residual) as u64
433 })
434 .collect();
435 let max_shifted = shifted.iter().copied().max().unwrap_or(0);
436 let bpv = bits_needed_u64(max_shifted);
437 writer.write_u8(CodecType::Linear as u8)?;
438 writer.write_u64::<LittleEndian>(first)?;
439 writer.write_u64::<LittleEndian>(last)?;
440 writer.write_u32::<LittleEndian>(n as u32)?;
441 writer.write_i64::<LittleEndian>(min_residual_i64)?;
442 writer.write_u8(bpv)?;
443 let mut bytes_written = 30u64; if bpv > 0 {
446 let mut packed = Vec::new();
447 bitpack_write(&shifted, bpv, &mut packed);
448 writer.write_all(&packed)?;
449 bytes_written += packed.len() as u64;
450 }
451
452 Ok(bytes_written)
453 }
454}
455
456#[inline]
458fn interpolate(first: u64, last: u64, n: usize, i: usize) -> u64 {
459 if n <= 1 {
460 return first;
461 }
462 let first = first as i128;
464 let last = last as i128;
465 let n = n as i128;
466 let i = i as i128;
467 let result = first + (last - first) * i / (n - 1);
468 result as u64
469}
470
471#[inline]
475pub fn linear_read(data: &[u8], index: usize) -> u64 {
476 let first = u64::from_le_bytes(data[0..8].try_into().unwrap());
477 let last = u64::from_le_bytes(data[8..16].try_into().unwrap());
478 let n = u32::from_le_bytes(data[16..20].try_into().unwrap()) as usize;
479 let offset = i64::from_le_bytes(data[20..28].try_into().unwrap());
480 let bpv = data[28];
481 let predicted = interpolate(first, last, n, index);
482 let residual = if bpv == 0 {
483 0u64
484 } else {
485 bitpack_read(&data[29..], bpv, index)
486 };
487 (predicted as i128 + offset as i128 + residual as i128) as u64
489}
490
491#[derive(Clone, Copy)]
498struct BlockwiseLinearBlockEstimate {
499 min_residual: i64,
500 bits_per_value: u8,
501 serialized_size: u64,
502}
503
504#[derive(Default)]
505pub struct BlockwiseLinearEstimator {
506 count: usize,
507 completed_size: u64,
508 current_block: Vec<u64>,
509 completed_blocks: Vec<BlockwiseLinearBlockEstimate>,
510 tail_estimate: Option<BlockwiseLinearBlockEstimate>,
511 overflow: bool,
512}
513
514impl BlockwiseLinearEstimator {
515 fn collect_values(&mut self, values: &[u64]) {
521 debug_assert_eq!(self.count, 0);
522 debug_assert!(self.current_block.is_empty());
523
524 self.count = values.len();
525 let mut blocks = values.chunks_exact(BLOCKWISE_LINEAR_BLOCK_SIZE);
526 for block in &mut blocks {
527 match estimate_blockwise_linear_block(block) {
528 Some(estimate) => {
529 self.completed_size += estimate.serialized_size;
530 self.completed_blocks.push(estimate);
531 }
532 None => {
533 self.overflow = true;
534 return;
535 }
536 }
537 }
538 self.current_block.extend_from_slice(blocks.remainder());
539 }
540}
541
542fn estimate_blockwise_linear_block(block: &[u64]) -> Option<BlockwiseLinearBlockEstimate> {
548 debug_assert!(!block.is_empty());
549 let block_len = block.len();
550 if block_len < 2 {
551 return Some(BlockwiseLinearBlockEstimate {
552 min_residual: 0,
553 bits_per_value: 0,
554 serialized_size: 29,
555 });
556 }
557
558 let first = block[0];
559 let last = block[block_len - 1];
560 let mut min_res = i128::MAX;
561 let mut max_res = i128::MIN;
562 for (i, &val) in block.iter().enumerate() {
563 let pred = interpolate(first, last, block_len, i);
564 let res = val as i128 - pred as i128;
565 min_res = min_res.min(res);
566 max_res = max_res.max(res);
567 }
568
569 if min_res < i64::MIN as i128 || max_res > i64::MAX as i128 {
572 return None;
573 }
574
575 let bits_per_value = bits_needed_u64((max_res - min_res) as u64);
576 let data_bytes = (block_len as u64 * u64::from(bits_per_value)).div_ceil(8);
577 Some(BlockwiseLinearBlockEstimate {
578 min_residual: min_res as i64,
579 bits_per_value,
580 serialized_size: 29 + data_bytes,
581 })
582}
583
584impl CodecEstimator for BlockwiseLinearEstimator {
585 fn collect(&mut self, value: u64) {
586 self.count += 1;
587 self.tail_estimate = None;
588 if self.overflow {
589 return;
590 }
591
592 self.current_block.push(value);
593 if self.current_block.len() == BLOCKWISE_LINEAR_BLOCK_SIZE {
594 match estimate_blockwise_linear_block(&self.current_block) {
595 Some(estimate) => {
596 self.completed_size += estimate.serialized_size;
597 self.completed_blocks.push(estimate);
598 }
599 None => self.overflow = true,
600 }
601 self.current_block.clear();
602 }
603 }
604
605 fn finalize(&mut self) {
606 self.tail_estimate = if self.current_block.is_empty() || self.overflow {
607 None
608 } else {
609 estimate_blockwise_linear_block(&self.current_block)
610 };
611 if !self.current_block.is_empty() && self.tail_estimate.is_none() {
612 self.overflow = true;
613 }
614 }
615
616 fn estimate(&self) -> Option<u64> {
617 if self.count < 2 * BLOCKWISE_LINEAR_BLOCK_SIZE || self.overflow {
618 return None;
620 }
621
622 let mut total = 9 + self.completed_size;
624 if !self.current_block.is_empty() {
625 total += self
626 .tail_estimate
627 .or_else(|| estimate_blockwise_linear_block(&self.current_block))?
628 .serialized_size;
629 }
630 Some(total)
631 }
632
633 fn serialize(&self, values: &[u64], writer: &mut dyn Write) -> io::Result<u64> {
634 let n = values.len();
635 let num_blocks = n.div_ceil(BLOCKWISE_LINEAR_BLOCK_SIZE);
636
637 writer.write_u8(CodecType::BlockwiseLinear as u8)?;
638 writer.write_u32::<LittleEndian>(n as u32)?;
639 writer.write_u32::<LittleEndian>(num_blocks as u32)?;
640 let mut bytes_written = 9u64;
641
642 let mut shifted = Vec::new();
645 let mut packed = Vec::new();
646 let estimates_match = self.count == n
647 && self.completed_blocks.len() == n / BLOCKWISE_LINEAR_BLOCK_SIZE
648 && (n.is_multiple_of(BLOCKWISE_LINEAR_BLOCK_SIZE) || self.tail_estimate.is_some());
649
650 for b in 0..num_blocks {
651 let start = b * BLOCKWISE_LINEAR_BLOCK_SIZE;
652 let end = (start + BLOCKWISE_LINEAR_BLOCK_SIZE).min(n);
653 let block = &values[start..end];
654 let block_len = block.len();
655
656 let first = block[0];
657 let last = if block_len > 1 {
658 block[block_len - 1]
659 } else {
660 first
661 };
662
663 let estimate = if estimates_match {
664 self.completed_blocks.get(b).copied().or(self.tail_estimate)
665 } else {
666 None
667 };
668 let estimate = match estimate {
669 Some(estimate) => estimate,
670 None => estimate_blockwise_linear_block(block).ok_or_else(|| {
671 io::Error::new(
672 io::ErrorKind::InvalidInput,
673 "blockwise linear codec: per-block residual offset exceeds i64 range",
674 )
675 })?,
676 };
677 let min_residual = i128::from(estimate.min_residual);
678
679 shifted.clear();
680 shifted.extend(block.iter().enumerate().map(|(i, &val)| {
681 if block_len < 2 {
682 0
683 } else {
684 let pred = interpolate(first, last, block_len, i);
685 let res = val as i128 - pred as i128;
686 (res - min_residual) as u64
687 }
688 }));
689 writer.write_u64::<LittleEndian>(first)?;
690 writer.write_u64::<LittleEndian>(last)?;
691 writer.write_i64::<LittleEndian>(estimate.min_residual)?;
692 writer.write_u8(estimate.bits_per_value)?;
693
694 packed.clear();
695 if estimate.bits_per_value > 0 {
696 bitpack_write(&shifted, estimate.bits_per_value, &mut packed);
697 }
698 writer.write_u32::<LittleEndian>(packed.len() as u32)?;
699 writer.write_all(&packed)?;
700 bytes_written += 29 + packed.len() as u64;
701 }
702
703 Ok(bytes_written)
704 }
705}
706
707pub fn blockwise_linear_read(data: &[u8], index: usize) -> u64 {
711 let _num_values = u32::from_le_bytes(data[0..4].try_into().unwrap()) as usize;
712 let num_blocks = u32::from_le_bytes(data[4..8].try_into().unwrap()) as usize;
713
714 let target_block = index / BLOCKWISE_LINEAR_BLOCK_SIZE;
715 let index_in_block = index % BLOCKWISE_LINEAR_BLOCK_SIZE;
716
717 let mut pos = 8usize;
719 for b in 0..num_blocks {
720 let first = u64::from_le_bytes(data[pos..pos + 8].try_into().unwrap());
721 let last = u64::from_le_bytes(data[pos + 8..pos + 16].try_into().unwrap());
722 let offset = i64::from_le_bytes(data[pos + 16..pos + 24].try_into().unwrap());
723 let bpv = data[pos + 24];
724 let packed_len = u32::from_le_bytes(data[pos + 25..pos + 29].try_into().unwrap()) as usize;
725
726 if b == target_block {
727 let block_start = b * BLOCKWISE_LINEAR_BLOCK_SIZE;
728 let block_end = ((b + 1) * BLOCKWISE_LINEAR_BLOCK_SIZE).min(_num_values);
729 let block_len = block_end - block_start;
730
731 let predicted = interpolate(first, last, block_len, index_in_block);
732 let residual = if bpv == 0 {
733 0u64
734 } else {
735 bitpack_read(&data[pos + 29..], bpv, index_in_block)
736 };
737 return (predicted as i128 + offset as i128 + residual as i128) as u64;
738 }
739
740 pos += 29 + packed_len;
741 }
742
743 0 }
745
746pub fn blockwise_linear_read_batch(data: &[u8], start_index: usize, out: &mut [u64]) {
752 if out.is_empty() {
753 return;
754 }
755
756 let num_values = u32::from_le_bytes(data[0..4].try_into().unwrap()) as usize;
757 let num_blocks = u32::from_le_bytes(data[4..8].try_into().unwrap()) as usize;
758 let valid_len = out.len().min(num_values.saturating_sub(start_index));
759 out[valid_len..].fill(0);
760 if valid_len == 0 {
761 return;
762 }
763
764 let target_block = start_index / BLOCKWISE_LINEAR_BLOCK_SIZE;
765 let mut pos = 8usize;
766 let mut written = 0usize;
767
768 for block_idx in 0..num_blocks {
769 let packed_len = u32::from_le_bytes(data[pos + 25..pos + 29].try_into().unwrap()) as usize;
770 if block_idx < target_block {
771 pos += 29 + packed_len;
772 continue;
773 }
774
775 let first = u64::from_le_bytes(data[pos..pos + 8].try_into().unwrap());
776 let last = u64::from_le_bytes(data[pos + 8..pos + 16].try_into().unwrap());
777 let offset = i64::from_le_bytes(data[pos + 16..pos + 24].try_into().unwrap());
778 let bpv = data[pos + 24];
779 let packed = &data[pos + 29..pos + 29 + packed_len];
780
781 let block_start = block_idx * BLOCKWISE_LINEAR_BLOCK_SIZE;
782 let block_len = (num_values - block_start).min(BLOCKWISE_LINEAR_BLOCK_SIZE);
783 let index_in_block = if block_idx == target_block {
784 start_index - block_start
785 } else {
786 0
787 };
788 let take = (block_len - index_in_block).min(valid_len - written);
789
790 for (i, value) in out[written..written + take].iter_mut().enumerate() {
791 let block_index = index_in_block + i;
792 let predicted = interpolate(first, last, block_len, block_index);
793 let residual = if bpv == 0 {
794 0
795 } else {
796 bitpack_read(packed, bpv, block_index)
797 };
798 *value = (predicted as i128 + offset as i128 + residual as i128) as u64;
799 }
800
801 written += take;
802 if written == valid_len {
803 break;
804 }
805 pos += 29 + packed_len;
806 }
807}
808
809pub fn serialize_auto(values: &[u64], writer: &mut dyn Write) -> io::Result<u64> {
815 let mut constant = ConstantEstimator::default();
816 let mut bitpacked = BitpackedEstimator::default();
817 let mut linear = LinearEstimator::default();
818 let mut blockwise = BlockwiseLinearEstimator::default();
819
820 for &v in values {
822 constant.collect(v);
823 bitpacked.collect(v);
824 linear.collect(v);
825 }
826 blockwise.collect_values(values);
827
828 constant.finalize();
830 bitpacked.finalize();
831 linear.finalize();
832 blockwise.finalize();
833
834 let candidates: Vec<(&dyn CodecEstimator, &str)> = vec![
836 (&constant, "constant"),
837 (&bitpacked, "bitpacked"),
838 (&linear, "linear"),
839 (&blockwise, "blockwise_linear"),
840 ];
841
842 let (best, _name) = candidates
843 .into_iter()
844 .filter_map(|(est, name)| est.estimate().map(|size| (est, name, size)))
845 .min_by_key(|&(_, _, size)| size)
846 .map(|(est, name, _)| (est, name))
847 .unwrap_or((&bitpacked as &dyn CodecEstimator, "bitpacked"));
848
849 best.serialize(values, writer)
850}
851
852pub fn bitpacked_read_batch(data: &[u8], start_index: usize, out: &mut [u64]) {
859 let min_value = u64::from_le_bytes(data[0..8].try_into().unwrap());
860 let bpv = data[8];
861
862 if bpv == 0 {
863 out.iter_mut().for_each(|v| *v = min_value);
864 return;
865 }
866
867 let packed = &data[9..];
868
869 match bpv {
870 8 => {
872 for (i, v) in out.iter_mut().enumerate() {
873 let idx = start_index + i;
874 *v = (packed[idx] as u64).wrapping_add(min_value);
875 }
876 }
877 16 => {
878 for (i, v) in out.iter_mut().enumerate() {
879 let idx = start_index + i;
880 let byte_off = idx * 2;
881 let raw = u16::from_le_bytes([packed[byte_off], packed[byte_off + 1]]);
882 *v = (raw as u64).wrapping_add(min_value);
883 }
884 }
885 32 => {
886 for (i, v) in out.iter_mut().enumerate() {
887 let idx = start_index + i;
888 let byte_off = idx * 4;
889 let raw = u32::from_le_bytes(packed[byte_off..byte_off + 4].try_into().unwrap());
890 *v = (raw as u64).wrapping_add(min_value);
891 }
892 }
893 64 => {
894 for (i, v) in out.iter_mut().enumerate() {
895 let idx = start_index + i;
896 let byte_off = idx * 8;
897 let raw = u64::from_le_bytes(packed[byte_off..byte_off + 8].try_into().unwrap());
898 *v = raw.wrapping_add(min_value);
899 }
900 }
901 _ => {
903 for (i, v) in out.iter_mut().enumerate() {
904 *v = super::bitpack_read(packed, bpv, start_index + i).wrapping_add(min_value);
905 }
906 }
907 }
908}
909
910pub fn auto_read_batch(data: &[u8], start_index: usize, out: &mut [u64]) {
915 if data.is_empty() || out.is_empty() {
916 out.iter_mut().for_each(|v| *v = 0);
917 return;
918 }
919 let codec_id = data[0];
920 let rest = &data[1..];
921 match CodecType::from_u8(codec_id) {
922 Some(CodecType::Constant) => {
923 let val = u64::from_le_bytes(rest[0..8].try_into().unwrap());
924 out.iter_mut().for_each(|v| *v = val);
925 }
926 Some(CodecType::Bitpacked) => bitpacked_read_batch(rest, start_index, out),
927 Some(CodecType::Linear) => {
928 for (i, v) in out.iter_mut().enumerate() {
929 *v = linear_read(rest, start_index + i);
930 }
931 }
932 Some(CodecType::BlockwiseLinear) => blockwise_linear_read_batch(rest, start_index, out),
933 None => out.iter_mut().for_each(|v| *v = 0),
934 }
935}
936
937#[inline]
941pub fn auto_read(data: &[u8], index: usize) -> u64 {
942 if data.is_empty() {
943 return 0;
944 }
945 let codec_id = data[0];
946 let rest = &data[1..];
947 match CodecType::from_u8(codec_id) {
948 Some(CodecType::Constant) => {
949 u64::from_le_bytes(rest[0..8].try_into().unwrap())
951 }
952 Some(CodecType::Bitpacked) => bitpacked_read(rest, index),
953 Some(CodecType::Linear) => linear_read(rest, index),
954 Some(CodecType::BlockwiseLinear) => blockwise_linear_read(rest, index),
955 None => 0,
956 }
957}
958
959#[cfg(test)]
962mod tests {
963 use super::*;
964
965 fn roundtrip(values: &[u64]) -> Vec<u64> {
966 let mut buf = Vec::new();
967 serialize_auto(values, &mut buf).unwrap();
968 (0..values.len()).map(|i| auto_read(&buf, i)).collect()
969 }
970
971 fn blockwise_values(len: usize) -> Vec<u64> {
972 (0..len)
973 .map(|i| {
974 let block = i / BLOCKWISE_LINEAR_BLOCK_SIZE;
975 let index = i % BLOCKWISE_LINEAR_BLOCK_SIZE;
976 block as u64 * 1_000_000
977 + index as u64 * (block as u64 + 3)
978 + ((i * 17 + block * 11) % 23) as u64
979 })
980 .collect()
981 }
982
983 fn serialize_blockwise_reference(values: &[u64]) -> Vec<u8> {
987 let n = values.len();
988 let num_blocks = n.div_ceil(BLOCKWISE_LINEAR_BLOCK_SIZE);
989 let mut encoded = Vec::new();
990 encoded.push(CodecType::BlockwiseLinear as u8);
991 encoded.extend_from_slice(&(n as u32).to_le_bytes());
992 encoded.extend_from_slice(&(num_blocks as u32).to_le_bytes());
993
994 for block in values.chunks(BLOCKWISE_LINEAR_BLOCK_SIZE) {
995 let block_len = block.len();
996 let first = block[0];
997 let last = if block_len > 1 {
998 block[block_len - 1]
999 } else {
1000 first
1001 };
1002 let min_residual = if block_len < 2 {
1003 0
1004 } else {
1005 block
1006 .iter()
1007 .enumerate()
1008 .map(|(i, &value)| {
1009 value as i128 - interpolate(first, last, block_len, i) as i128
1010 })
1011 .min()
1012 .unwrap()
1013 };
1014 let shifted: Vec<u64> = block
1015 .iter()
1016 .enumerate()
1017 .map(|(i, &value)| {
1018 if block_len < 2 {
1019 0
1020 } else {
1021 let predicted = interpolate(first, last, block_len, i);
1022 (value as i128 - predicted as i128 - min_residual) as u64
1023 }
1024 })
1025 .collect();
1026 let bpv = bits_needed_u64(shifted.iter().copied().max().unwrap_or(0));
1027 let mut packed = Vec::new();
1028 bitpack_write(&shifted, bpv, &mut packed);
1029
1030 encoded.extend_from_slice(&first.to_le_bytes());
1031 encoded.extend_from_slice(&last.to_le_bytes());
1032 encoded.extend_from_slice(&(min_residual as i64).to_le_bytes());
1033 encoded.push(bpv);
1034 encoded.extend_from_slice(&(packed.len() as u32).to_le_bytes());
1035 encoded.extend_from_slice(&packed);
1036 }
1037
1038 encoded
1039 }
1040
1041 #[test]
1042 fn test_constant_codec() {
1043 let values: Vec<u64> = vec![42; 100];
1044 let mut buf = Vec::new();
1045 serialize_auto(&values, &mut buf).unwrap();
1046 assert_eq!(buf[0], CodecType::Constant as u8);
1047 assert_eq!(buf.len(), 9);
1048 assert_eq!(roundtrip(&values), values);
1049 }
1050
1051 #[test]
1052 fn test_bitpacked_codec() {
1053 let values: Vec<u64> = (0..50).map(|i| 1000 + (i % 7) * 13).collect();
1054 let result = roundtrip(&values);
1055 assert_eq!(result, values);
1056 }
1057
1058 #[test]
1059 fn test_linear_codec_sequential() {
1060 let values: Vec<u64> = (0..1000).map(|i| 100 + i * 3).collect();
1062 let mut buf = Vec::new();
1063 serialize_auto(&values, &mut buf).unwrap();
1064 assert_eq!(roundtrip(&values), values);
1066 }
1067
1068 #[test]
1069 fn test_blockwise_linear_codec() {
1070 let mut values: Vec<u64> = Vec::new();
1072 for i in 0..1500 {
1073 if i < 750 {
1074 values.push(100 + i * 2);
1075 } else {
1076 values.push(5000 + (i - 750) * 5);
1077 }
1078 }
1079 let result = roundtrip(&values);
1080 assert_eq!(result, values);
1081 }
1082
1083 #[test]
1084 fn test_blockwise_estimate_is_bounded_and_serialization_is_byte_compatible() {
1085 for len in [1024, 1025, 1536, 1537, 4097] {
1086 let values = blockwise_values(len);
1087 let reference = serialize_blockwise_reference(&values);
1088 let mut estimator = BlockwiseLinearEstimator::default();
1089 for &value in &values {
1090 estimator.collect(value);
1091 }
1092 estimator.finalize();
1093
1094 assert_eq!(estimator.estimate(), Some(reference.len() as u64));
1095 assert!(
1096 estimator.current_block.len() < BLOCKWISE_LINEAR_BLOCK_SIZE,
1097 "estimator retained more than one partial block"
1098 );
1099 assert!(
1100 estimator.current_block.capacity() <= BLOCKWISE_LINEAR_BLOCK_SIZE,
1101 "estimator scratch grew beyond one block"
1102 );
1103
1104 let mut encoded = Vec::new();
1105 estimator.serialize(&values, &mut encoded).unwrap();
1106 assert_eq!(
1107 encoded, reference,
1108 "serialized bytes changed for {len} values"
1109 );
1110 }
1111 }
1112
1113 #[test]
1114 fn test_empty() {
1115 let values: Vec<u64> = vec![];
1116 let mut buf = Vec::new();
1117 serialize_auto(&values, &mut buf).unwrap();
1118 assert!(buf.len() <= 10);
1119 }
1120
1121 #[test]
1122 fn test_validate_rejects_truncated_and_inconsistent_payloads() {
1123 assert!(validate_auto(&[], 1).is_err());
1124 assert!(validate_auto(&[CodecType::Constant as u8], 1).is_err());
1125
1126 let mut bitpacked = vec![CodecType::Bitpacked as u8];
1127 bitpacked.extend_from_slice(&0u64.to_le_bytes());
1128 bitpacked.push(65);
1129 assert!(validate_auto(&bitpacked, 1).is_err());
1130
1131 let mut valid = Vec::new();
1132 serialize_auto(&[1, 2, 3, 4], &mut valid).unwrap();
1133 assert!(validate_auto(&valid, 4).is_ok());
1134 valid.pop();
1135 assert!(validate_auto(&valid, 4).is_err());
1136 }
1137
1138 #[test]
1139 fn test_single_value() {
1140 let values = vec![999u64];
1141 assert_eq!(roundtrip(&values), values);
1142 }
1143
1144 #[test]
1145 fn test_two_values() {
1146 let values = vec![10u64, 20];
1147 assert_eq!(roundtrip(&values), values);
1148 }
1149
1150 #[test]
1151 fn test_large_range() {
1152 let values = vec![0u64, u64::MAX / 2, u64::MAX];
1153 assert_eq!(roundtrip(&values), values);
1154 }
1155
1156 #[test]
1157 fn test_timestamps_pick_linear_or_blockwise() {
1158 let mut values: Vec<u64> = Vec::new();
1160 let mut ts = 1_700_000_000u64;
1161 for _ in 0..2000 {
1162 values.push(ts);
1163 ts += 1000 + (ts % 7); }
1165 let result = roundtrip(&values);
1166 assert_eq!(result, values);
1167 }
1168
1169 fn roundtrip_batch(values: &[u64]) {
1171 let mut buf = Vec::new();
1172 serialize_auto(values, &mut buf).unwrap();
1173
1174 let mut batch_out = vec![0u64; values.len()];
1176 auto_read_batch(&buf, 0, &mut batch_out);
1177 assert_eq!(batch_out, values, "batch read mismatch");
1178
1179 if values.len() >= 10 {
1181 let start = 3;
1182 let count = values.len() - 6;
1183 let mut sub = vec![0u64; count];
1184 auto_read_batch(&buf, start, &mut sub);
1185 assert_eq!(
1186 sub,
1187 &values[start..start + count],
1188 "sub-range batch mismatch"
1189 );
1190 }
1191 }
1192
1193 #[test]
1194 fn test_batch_read_constant() {
1195 roundtrip_batch(&vec![42u64; 100]);
1196 }
1197
1198 #[test]
1199 fn test_batch_read_bitpacked_8bit() {
1200 let values: Vec<u64> = (0..200).map(|i| 1000 + (i % 200)).collect();
1202 roundtrip_batch(&values);
1203 }
1204
1205 #[test]
1206 fn test_batch_read_bitpacked_16bit() {
1207 let values: Vec<u64> = (0..200).map(|i| 50000 + i * 100).collect();
1209 roundtrip_batch(&values);
1210 }
1211
1212 #[test]
1213 fn test_batch_read_bitpacked_arbitrary() {
1214 let values: Vec<u64> = (0..100).map(|i| 999 + (i * 37) % 8000).collect();
1216 roundtrip_batch(&values);
1217 }
1218
1219 #[test]
1220 fn test_batch_read_linear() {
1221 let values: Vec<u64> = (0..500).map(|i| 100 + i * 3).collect();
1222 roundtrip_batch(&values);
1223 }
1224
1225 #[test]
1226 fn test_batch_read_blockwise() {
1227 let mut values = Vec::new();
1228 for i in 0..1500u64 {
1229 values.push(if i < 750 {
1230 100 + i * 2
1231 } else {
1232 5000 + (i - 750) * 5
1233 });
1234 }
1235 roundtrip_batch(&values);
1236 }
1237
1238 #[test]
1239 fn test_blockwise_batch_read_across_block_boundaries() {
1240 let values = blockwise_values(2053);
1241 let estimator = BlockwiseLinearEstimator::default();
1242 let mut encoded = Vec::new();
1243 estimator.serialize(&values, &mut encoded).unwrap();
1244 assert_eq!(encoded[0], CodecType::BlockwiseLinear as u8);
1245
1246 for (start, len) in [
1247 (0, values.len()),
1248 (510, 5),
1249 (511, 3),
1250 (512, 513),
1251 (777, 900),
1252 (1023, 514),
1253 (1535, 518),
1254 (2048, 5),
1255 ] {
1256 let expected = &values[start..start + len];
1257
1258 let mut direct = vec![u64::MAX; len];
1259 blockwise_linear_read_batch(&encoded[1..], start, &mut direct);
1260 assert_eq!(direct, expected, "direct batch mismatch at {start}");
1261
1262 let mut auto = vec![u64::MAX; len];
1263 auto_read_batch(&encoded, start, &mut auto);
1264 assert_eq!(auto, expected, "auto batch mismatch at {start}");
1265 }
1266 }
1267
1268 #[test]
1272 fn test_zigzag_timestamps_with_missing() {
1273 use super::super::{FAST_FIELD_MISSING, zigzag_encode};
1274
1275 let timestamps: Vec<i64> = vec![
1277 1724630400, 1724716800, 1724803200, 1700000000, 1680000000, 1724630400, ];
1284
1285 let mut values = Vec::new();
1287 for (i, &ts) in timestamps.iter().enumerate() {
1288 values.push(zigzag_encode(ts));
1289 if i % 2 == 1 {
1291 values.push(FAST_FIELD_MISSING);
1292 }
1293 }
1294
1295 let result = roundtrip(&values);
1296 assert_eq!(
1297 result, values,
1298 "zigzag timestamps + missing roundtrip failed"
1299 );
1300 }
1301
1302 #[test]
1304 fn test_codecs_individually_with_zigzag_and_missing() {
1305 use super::super::{FAST_FIELD_MISSING, zigzag_encode};
1306
1307 let values: Vec<u64> = vec![
1308 zigzag_encode(1724630400), zigzag_encode(1700000000), FAST_FIELD_MISSING,
1311 zigzag_encode(1680000000), zigzag_encode(1724716800), FAST_FIELD_MISSING,
1314 zigzag_encode(1724630400), zigzag_encode(0), ];
1317
1318 {
1320 let mut est = BitpackedEstimator::default();
1321 for &v in &values {
1322 est.collect(v);
1323 }
1324 est.finalize();
1325 if est.estimate().is_some() {
1326 let mut buf = Vec::new();
1327 est.serialize(&values, &mut buf).unwrap();
1328 for (i, &expected) in values.iter().enumerate() {
1329 let got = auto_read(&buf, i);
1330 assert_eq!(
1331 got, expected,
1332 "bitpacked: index {} expected {} got {}",
1333 i, expected, got
1334 );
1335 }
1336 }
1337 }
1338
1339 {
1341 let mut est = LinearEstimator::default();
1342 for &v in &values {
1343 est.collect(v);
1344 }
1345 est.finalize();
1346 if est.estimate().is_some() {
1347 let mut buf = Vec::new();
1348 est.serialize(&values, &mut buf).unwrap();
1349 for (i, &expected) in values.iter().enumerate() {
1350 let got = auto_read(&buf, i);
1351 assert_eq!(
1352 got, expected,
1353 "linear: index {} expected {} got {}",
1354 i, expected, got
1355 );
1356 }
1357 }
1358 }
1359
1360 let result = roundtrip(&values);
1362 assert_eq!(result, values, "auto codec roundtrip failed");
1363 }
1364
1365 #[test]
1367 fn test_specific_issued_at_roundtrip() {
1368 use super::super::{FAST_FIELD_MISSING, zigzag_encode};
1369
1370 let mut values = Vec::new();
1372 let base_ts = 1724630400i64; for i in 0..100u64 {
1374 if i % 5 == 0 {
1375 values.push(FAST_FIELD_MISSING);
1377 } else {
1378 let ts = base_ts - (i as i64 * 86400); values.push(zigzag_encode(ts));
1381 }
1382 }
1383
1384 let result = roundtrip(&values);
1385 for (i, (&expected, &got)) in values.iter().zip(result.iter()).enumerate() {
1386 assert_eq!(
1387 got,
1388 expected,
1389 "doc {}: expected {} (zigzag of {}), got {}",
1390 i,
1391 expected,
1392 if expected == FAST_FIELD_MISSING {
1393 -1 } else {
1395 super::super::zigzag_decode(expected)
1396 },
1397 got
1398 );
1399 }
1400 }
1401
1402 #[test]
1405 fn test_large_scale_timestamp_roundtrip() {
1406 use super::super::{FAST_FIELD_MISSING, zigzag_encode};
1407
1408 for num_docs in [10_000, 50_000, 100_000] {
1409 let mut values = Vec::with_capacity(num_docs);
1410 let base_ts = 1724630400i64;
1411
1412 for i in 0..num_docs {
1413 if i % 7 == 0 {
1414 values.push(FAST_FIELD_MISSING);
1415 } else {
1416 let ts = base_ts - (i as i64 * 3600) + ((i as i64 * 37) % 1000);
1418 values.push(zigzag_encode(ts));
1419 }
1420 }
1421
1422 let mut buf = Vec::new();
1424 serialize_auto(&values, &mut buf).unwrap();
1425 let codec_id = buf[0];
1426 let codec_name = match CodecType::from_u8(codec_id) {
1427 Some(CodecType::Constant) => "constant",
1428 Some(CodecType::Bitpacked) => "bitpacked",
1429 Some(CodecType::Linear) => "linear",
1430 Some(CodecType::BlockwiseLinear) => "blockwise_linear",
1431 None => "unknown",
1432 };
1433
1434 let mut failures = Vec::new();
1436 for (i, &expected) in values.iter().enumerate() {
1437 let got = auto_read(&buf, i);
1438 if got != expected {
1439 failures.push((i, expected, got));
1440 if failures.len() >= 5 {
1441 break;
1442 }
1443 }
1444 }
1445
1446 assert!(
1447 failures.is_empty(),
1448 "num_docs={}, codec={}: {} failures. First 5: {:?}",
1449 num_docs,
1450 codec_name,
1451 failures.len(),
1452 failures
1453 );
1454 }
1455 }
1456
1457 #[test]
1461 fn test_blockwise_linear_with_clustered_missing() {
1462 use super::super::{FAST_FIELD_MISSING, zigzag_encode};
1463
1464 let mut values = Vec::new();
1469
1470 for i in 0..512 {
1473 if i < 100 {
1474 values.push(FAST_FIELD_MISSING);
1475 } else {
1476 let ts = 1724630400i64 + (i as i64 * 100);
1477 values.push(zigzag_encode(ts));
1478 }
1479 }
1480
1481 for i in 512..3072 {
1483 let ts = 1724630400i64 + (i as i64 * 100);
1484 values.push(zigzag_encode(ts));
1485 }
1486
1487 let result = roundtrip(&values);
1488 let mut failures = Vec::new();
1489 for (i, (&expected, &got)) in values.iter().zip(result.iter()).enumerate() {
1490 if got != expected {
1491 failures.push((i, expected, got));
1492 }
1493 }
1494 assert!(
1495 failures.is_empty(),
1496 "blockwise linear with clustered missing: {} failures. First 5: {:?}",
1497 failures.len(),
1498 &failures[..failures.len().min(5)]
1499 );
1500 }
1501
1502 #[test]
1505 fn test_forced_codecs_with_timestamps_and_missing() {
1506 use super::super::{FAST_FIELD_MISSING, zigzag_encode};
1507
1508 let mut values = Vec::new();
1509 let base_ts = 1724630400i64;
1510 for i in 0..200 {
1511 if i % 5 == 0 {
1512 values.push(FAST_FIELD_MISSING);
1513 } else {
1514 let ts = base_ts - (i as i64 * 86400);
1515 values.push(zigzag_encode(ts));
1516 }
1517 }
1518
1519 {
1521 let est = BitpackedEstimator::default();
1522 let mut buf = Vec::new();
1523 est.serialize(&values, &mut buf).unwrap();
1524 for (i, &expected) in values.iter().enumerate() {
1525 let got = bitpacked_read(&buf[1..], i); assert_eq!(got, expected, "forced bitpacked: index {} failed", i);
1527 }
1528 }
1529
1530 {
1533 let est = LinearEstimator::default();
1534 let mut buf = Vec::new();
1535 let result = est.serialize(&values, &mut buf);
1536 assert!(
1537 result.is_err(),
1538 "linear codec should reject data with residuals exceeding i64"
1539 );
1540 }
1541
1542 {
1544 let safe_values: Vec<u64> = values
1545 .iter()
1546 .filter(|&&v| v != FAST_FIELD_MISSING)
1547 .copied()
1548 .collect();
1549 let est = LinearEstimator::default();
1550 let mut buf = Vec::new();
1551 est.serialize(&safe_values, &mut buf).unwrap();
1552 for (i, &expected) in safe_values.iter().enumerate() {
1553 let got = linear_read(&buf[1..], i);
1554 assert_eq!(got, expected, "forced linear (safe): index {} failed", i);
1555 }
1556 }
1557
1558 {
1560 let mut large_values = Vec::new();
1561 for i in 0..2000 {
1562 if i % 5 == 0 {
1563 large_values.push(FAST_FIELD_MISSING);
1564 } else {
1565 let ts = base_ts - (i as i64 * 86400);
1566 large_values.push(zigzag_encode(ts));
1567 }
1568 }
1569 let mut est = BlockwiseLinearEstimator::default();
1570 for &v in &large_values {
1571 est.collect(v);
1572 }
1573 assert!(
1574 est.estimate().is_none(),
1575 "blockwise linear should reject data with per-block residuals exceeding i64"
1576 );
1577 }
1578 }
1579}