Skip to main content

rudb_encoding/
integer.rs

1//! The single column integer encodings and the cascade over them.
2//!
3//! `spec/06-compression.md` section 6.2 lists the encoding set and section 6.3 says the ratios are
4//! in the cascade rather than in any one encoding. This module is both: the six candidate shapes
5//! for an integer column, each of which encodes its own output by calling back into the chooser, so
6//! that RLE over a dictionary over a bit packed code array is a thing that happens by construction
7//! rather than a case somebody wrote out.
8//!
9//! Everything here works on `i64`. A narrower column is widened on the way in and nothing is lost
10//! by it, because every encoding's size comes from the range of the values rather than from the
11//! declared width of the type: a `SMALLINT` column of values 100 to 130 packs to 5 bits whether it
12//! arrived as `i16` or as `i64`. The one place the widening would cost something is a raw copy, and
13//! there is no raw copy, because a bit packed unit at width 64 is exactly that and the chooser
14//! reaches it on its own when nothing else fits.
15//!
16//! ## The unit
17//!
18//! Bit packing is per 1024 values, per [`crate::bitpack`]. Everything else is per chunk, where a
19//! chunk is however many values the caller passes in and is meant to be a row group. The two
20//! granularities are the point rather than an accident. A frame of reference base that is chosen
21//! per 1024 values tracks a column that drifts, which is what a timestamp column and an
22//! autoincrementing key both do, and one base per row group would pay the whole range of the row
23//! group on every value. A dictionary, on the other hand, is worth more the larger the unit it
24//! covers, which is the argument section 6.5 takes all the way to a dictionary per table.
25//!
26//! ## The serialized form
27//!
28//! A chunk is a tag byte, a value count, and a body whose shape depends on the tag. Bodies that
29//! contain another array of integers contain a whole chunk, tag and all, which is what makes the
30//! decoder a fold and what makes the cascade free: nothing in `Rle` knows what its run lengths are
31//! encoded as. The header is fixed width little endian rather than a varint, because 5 bytes per
32//! chunk against a chunk that holds a row group is not worth the branch on the decode path.
33//!
34//! ## What the chooser does, and what it will have to do instead
35//!
36//! It encodes every candidate and keeps the smallest. That is the honest baseline for M1, which is
37//! a measurement of what the format can do rather than of how fast a writer can decide, and it is
38//! not what a write path can afford. Section 6.3 describes the real thing: evaluate the candidates
39//! on a systematic sample, not the first N rows, because column data is frequently clustered and
40//! the first 1024 rows of a sorted column look constant. Building that first would mean the numbers
41//! this milestone produces are the sampler's numbers rather than the format's, and there would be
42//! no way to tell how much the sampler is leaving behind.
43
44use rudb_common::{Error, Result};
45
46use crate::reader::Reader;
47
48use crate::bitpack::{self, VALUES};
49
50/// How deep a cascade is allowed to go.
51///
52/// Three levels is what section 6.3 says captures most of what a general compressor would find:
53/// dictionary, then bit packed codes, then nothing left worth doing. The limit exists because the
54/// chooser is exhaustive and a cascade that could nest forever would be exponential, and because a
55/// fourth level has never once been the smallest candidate in anything measured so far.
56const MAX_DEPTH: u8 = 3;
57
58/// What a chunk is encoded as. The discriminant is the tag byte in the serialized form and is part
59/// of the format, so the numbers are written down rather than left to the compiler.
60#[derive(Debug, Clone, Copy, PartialEq, Eq)]
61pub enum Kind {
62    /// One value repeated. The whole chunk is the tag, the count and the value.
63    Constant = 0,
64    /// Frame of reference then bit packed, per 1024 values. Covers plain bit packing at base zero
65    /// and a raw copy at width 64.
66    Packed = 1,
67    /// Differences between neighbours, zigzagged so a decreasing column is as cheap as an
68    /// increasing one, then encoded as a chunk in its own right.
69    Delta = 2,
70    /// Run values and run lengths, each encoded as a chunk in its own right.
71    Rle = 3,
72    /// A dictionary of the distinct values and an array of codes into it, both encoded as chunks in
73    /// their own right.
74    Dict = 4,
75    /// One dominant value with an exception list of positions and values.
76    Sparse = 5,
77}
78
79impl Kind {
80    fn tag(self) -> u8 {
81        self as u8
82    }
83
84    fn from_tag(tag: u8) -> Result<Self> {
85        match tag {
86            0 => Ok(Self::Constant),
87            1 => Ok(Self::Packed),
88            2 => Ok(Self::Delta),
89            3 => Ok(Self::Rle),
90            4 => Ok(Self::Dict),
91            5 => Ok(Self::Sparse),
92            other => Err(Error::internal(format!("unknown encoding tag {other}"))),
93        }
94    }
95
96    /// The name that goes in a report.
97    #[must_use]
98    pub fn name(self) -> &'static str {
99        match self {
100            Self::Constant => "CONSTANT",
101            Self::Packed => "FOR+BITPACK",
102            Self::Delta => "DELTA",
103            Self::Rle => "RLE",
104            Self::Dict => "DICT",
105            Self::Sparse => "SPARSE",
106        }
107    }
108}
109
110/// Encodes a chunk of integers, choosing the cascade that comes out smallest.
111///
112/// # Errors
113///
114/// If the chunk is longer than `u32::MAX`, or if an encoding produces something its own decoder
115/// would not accept, which is an internal inconsistency rather than a caller error.
116pub fn encode(values: &[i64]) -> Result<Vec<u8>> {
117    encode_at(values, 0)
118}
119
120/// Decodes a chunk written by [`encode`].
121///
122/// # Errors
123///
124/// If the bytes are truncated, carry an unknown tag, or describe a chunk whose parts do not agree
125/// with each other.
126pub fn decode(bytes: &[u8]) -> Result<Vec<i64>> {
127    let mut reader = Reader::new(bytes);
128    let values = decode_chunk(&mut reader)?;
129    if reader.remaining() != 0 {
130        return Err(Error::internal(format!(
131            "{} bytes left over after decoding a chunk",
132            reader.remaining()
133        )));
134    }
135    Ok(values)
136}
137
138/// Decodes a chunk that sits at the front of a longer buffer, and says how many bytes it took.
139///
140/// A string column holds integer chunks inside its own body, and the reader on that side cannot
141/// know where the nested chunk ends until it has been read. A chunk is self delimiting, so this is
142/// the same work [`decode`] does without the check that nothing follows.
143///
144/// # Errors
145///
146/// As [`decode`], except that trailing bytes are what the caller asked about rather than an error.
147pub fn decode_prefix(bytes: &[u8]) -> Result<(Vec<i64>, usize)> {
148    let mut reader = Reader::new(bytes);
149    let values = decode_chunk(&mut reader)?;
150    Ok((values, reader.used()))
151}
152
153/// [`describe`] over a chunk at the front of a longer buffer, and how many bytes it took.
154///
155/// # Errors
156///
157/// As [`decode_prefix`].
158pub fn describe_prefix(bytes: &[u8]) -> Result<(String, usize)> {
159    let mut reader = Reader::new(bytes);
160    let text = describe_chunk(&mut reader)?;
161    Ok((text, reader.used()))
162}
163
164/// The size in bytes of every candidate, for a report that wants to say what the cascade was
165/// chosen over rather than only what it chose. A candidate that does not apply is absent.
166///
167/// # Errors
168///
169/// As [`encode`].
170pub fn candidate_sizes(values: &[i64]) -> Result<Vec<(Kind, usize)>> {
171    let mut sizes = Vec::new();
172    for kind in candidates(values, 0) {
173        if let Some(bytes) = encode_as(kind, values, 0)? {
174            sizes.push((kind, bytes.len()));
175        }
176    }
177    Ok(sizes)
178}
179
180/// The cascade a chunk was encoded as, as a line of text like `DICT(PACKED, PACKED)`.
181///
182/// # Errors
183///
184/// As [`decode`].
185pub fn describe(bytes: &[u8]) -> Result<String> {
186    let mut reader = Reader::new(bytes);
187    describe_chunk(&mut reader)
188}
189
190fn encode_at(values: &[i64], depth: u8) -> Result<Vec<u8>> {
191    let mut best: Option<Vec<u8>> = None;
192    for kind in candidates(values, depth) {
193        let Some(bytes) = encode_as(kind, values, depth)? else {
194            continue;
195        };
196        if best.as_ref().is_none_or(|current| bytes.len() < current.len()) {
197            best = Some(bytes);
198        }
199    }
200    // `Packed` applies to every input including the empty one, so the chooser always has at least
201    // one candidate and this cannot be reached without a bug in `candidates`.
202    best.ok_or_else(|| Error::internal("no encoding applied to the chunk"))
203}
204
205/// Which candidates are worth encoding for this input.
206///
207/// The filters here are not the cost model. They are the cases where the encoding cannot be
208/// expressed at all, or is provably larger than `Packed` on the same data, so that the exhaustive
209/// chooser does not spend a dictionary build on a column of 100,000 distinct values to discover
210/// what its distinct count already said.
211fn candidates(values: &[i64], depth: u8) -> Vec<Kind> {
212    let mut kinds = vec![Kind::Packed];
213    if depth >= MAX_DEPTH || values.is_empty() {
214        return kinds;
215    }
216    if values.iter().all(|value| *value == values[0]) {
217        // Nothing else can beat 13 bytes, so this is the whole answer rather than a candidate.
218        return vec![Kind::Constant];
219    }
220    if values.len() >= 2 && deltas(values).is_some() {
221        kinds.push(Kind::Delta);
222    }
223    if run_count(values) * 4 <= values.len() * 3 {
224        kinds.push(Kind::Rle);
225    }
226    let distinct = distinct_values(values);
227    if distinct.len() * 2 <= values.len() {
228        kinds.push(Kind::Dict);
229    }
230    // Written as a match rather than as a chained `if let` because the minimum supported Rust
231    // version is 1.85 and let chains landed in 1.88.
232    match dominant_value(values) {
233        Some((_, count)) if count * 10 >= values.len() * 8 => kinds.push(Kind::Sparse),
234        _ => {}
235    }
236    kinds
237}
238
239/// `None` when the encoding does not apply to this input, which the caller treats as a candidate
240/// that did not run rather than as a failure.
241fn encode_as(kind: Kind, values: &[i64], depth: u8) -> Result<Option<Vec<u8>>> {
242    let mut out = Vec::new();
243    put_u8(&mut out, kind.tag());
244    put_u32(&mut out, u32::try_from(values.len()).map_err(|_| too_long(values.len()))?);
245    match kind {
246        Kind::Constant => {
247            let Some(first) = values.first() else {
248                return Ok(None);
249            };
250            if values.iter().any(|value| value != first) {
251                return Ok(None);
252            }
253            put_i64(&mut out, *first);
254        }
255        Kind::Packed => encode_packed(values, &mut out)?,
256        Kind::Delta => {
257            let Some(deltas) = deltas(values) else {
258                return Ok(None);
259            };
260            put_i64(&mut out, values[0]);
261            out.extend_from_slice(&encode_at(&deltas, depth + 1)?);
262        }
263        Kind::Rle => {
264            let (run_values, run_lengths) = runs(values);
265            if run_values.is_empty() {
266                return Ok(None);
267            }
268            out.extend_from_slice(&encode_at(&run_values, depth + 1)?);
269            out.extend_from_slice(&encode_at(&run_lengths, depth + 1)?);
270        }
271        Kind::Dict => {
272            let dictionary = distinct_values(values);
273            if dictionary.is_empty() {
274                return Ok(None);
275            }
276            let codes = codes_over(values, &dictionary);
277            out.extend_from_slice(&encode_at(&dictionary, depth + 1)?);
278            out.extend_from_slice(&encode_at(&codes, depth + 1)?);
279        }
280        Kind::Sparse => {
281            let Some((value, _)) = dominant_value(values) else {
282                return Ok(None);
283            };
284            let mut positions = Vec::new();
285            let mut exceptions = Vec::new();
286            for (index, other) in values.iter().enumerate() {
287                if *other != value {
288                    positions.push(index as i64);
289                    exceptions.push(*other);
290                }
291            }
292            put_i64(&mut out, value);
293            put_u32(
294                &mut out,
295                u32::try_from(positions.len()).map_err(|_| too_long(positions.len()))?,
296            );
297            out.extend_from_slice(&encode_at(&positions, depth + 1)?);
298            out.extend_from_slice(&encode_at(&exceptions, depth + 1)?);
299        }
300    }
301    Ok(Some(out))
302}
303
304/// Frame of reference and bit packing, one base and one width per 1024 values.
305///
306/// A base per unit rather than per chunk is most of what makes this work on real columns. A
307/// timestamp column over a day drifts across a range that needs 47 bits, and the same column inside
308/// any one unit spans a few seconds and needs 12. One base per row group would pay the 47 on every
309/// value.
310///
311/// A unit shorter than 1024 values, which is the last one of any chunk whose length is not a
312/// multiple of the unit and is the only one of every short array in a cascade, goes through
313/// [`bitpack::pack_tail`] instead. The transposed layout has no partial form and would charge a
314/// five entry dictionary for 1024 entries.
315fn encode_packed(values: &[i64], out: &mut Vec<u8>) -> Result<()> {
316    for unit in values.chunks(VALUES) {
317        let base = unit.iter().copied().min().unwrap_or(0);
318        let offsets: Vec<u64> = unit.iter().map(|value| offset_from(*value, base)).collect();
319        let width = bitpack::required_width(&offsets);
320        put_i64(out, base);
321        put_u8(out, u8::try_from(width).map_err(|_| Error::internal("impossible width"))?);
322        if unit.len() == VALUES {
323            let mut packed = vec![0u64; bitpack::packed_len::<u64>(width)];
324            bitpack::pack(&offsets, width, &mut packed)?;
325            for word in packed {
326                put_u64(out, word);
327            }
328        } else {
329            bitpack::pack_tail(&offsets, width, out)?;
330        }
331    }
332    Ok(())
333}
334
335fn decode_chunk(reader: &mut Reader<'_>) -> Result<Vec<i64>> {
336    let kind = Kind::from_tag(reader.u8()?)?;
337    let count = reader.u32()? as usize;
338    match kind {
339        Kind::Constant => Ok(vec![reader.i64()?; count]),
340        Kind::Packed => {
341            let mut values = Vec::with_capacity(count);
342            while values.len() < count {
343                let base = reader.i64()?;
344                let width = reader.u8()? as usize;
345                let wanted = (count - values.len()).min(VALUES);
346                if wanted == VALUES {
347                    let mut packed = vec![0u64; bitpack::packed_len::<u64>(width)];
348                    for word in &mut packed {
349                        *word = reader.u64()?;
350                    }
351                    let mut unit = vec![0u64; VALUES];
352                    bitpack::unpack(&packed, width, &mut unit)?;
353                    values.extend(unit.iter().map(|offset| value_from(*offset, base)));
354                } else {
355                    let bytes = reader.bytes(bitpack::tail_len(wanted, width))?;
356                    let unit = bitpack::unpack_tail(bytes, width, wanted)?;
357                    values.extend(unit.iter().map(|offset| value_from(*offset, base)));
358                }
359            }
360            Ok(values)
361        }
362        Kind::Delta => {
363            let first = reader.i64()?;
364            let deltas = decode_chunk(reader)?;
365            let mut values = Vec::with_capacity(count);
366            values.push(first);
367            let mut current = first;
368            for delta in deltas {
369                current = current.wrapping_add(unzigzag(delta as u64));
370                values.push(current);
371            }
372            check_count(values.len(), count)?;
373            Ok(values)
374        }
375        Kind::Rle => {
376            let run_values = decode_chunk(reader)?;
377            let run_lengths = decode_chunk(reader)?;
378            if run_values.len() != run_lengths.len() {
379                return Err(Error::internal("an RLE chunk has more runs than run lengths"));
380            }
381            let mut values = Vec::with_capacity(count);
382            for (value, length) in run_values.into_iter().zip(run_lengths) {
383                let length = usize::try_from(length)
384                    .map_err(|_| Error::internal("a negative RLE run length"))?;
385                values.extend(std::iter::repeat_n(value, length));
386            }
387            check_count(values.len(), count)?;
388            Ok(values)
389        }
390        Kind::Dict => {
391            let dictionary = decode_chunk(reader)?;
392            let codes = decode_chunk(reader)?;
393            let mut values = Vec::with_capacity(count);
394            for code in codes {
395                let index =
396                    usize::try_from(code).ok().and_then(|index| dictionary.get(index)).ok_or_else(
397                        || Error::internal(format!("code {code} is not in the dictionary")),
398                    )?;
399                values.push(*index);
400            }
401            check_count(values.len(), count)?;
402            Ok(values)
403        }
404        Kind::Sparse => {
405            let value = reader.i64()?;
406            let exception_count = reader.u32()? as usize;
407            let positions = decode_chunk(reader)?;
408            let exceptions = decode_chunk(reader)?;
409            if positions.len() != exception_count || exceptions.len() != exception_count {
410                return Err(Error::internal("a sparse chunk disagrees about its exception count"));
411            }
412            let mut values = vec![value; count];
413            for (position, exception) in positions.into_iter().zip(exceptions) {
414                let position = usize::try_from(position)
415                    .ok()
416                    .filter(|position| *position < count)
417                    .ok_or_else(|| {
418                        Error::internal(format!("exception at {position} is outside the chunk"))
419                    })?;
420                values[position] = exception;
421            }
422            Ok(values)
423        }
424    }
425}
426
427fn describe_chunk(reader: &mut Reader<'_>) -> Result<String> {
428    let kind = Kind::from_tag(reader.u8()?)?;
429    let count = reader.u32()? as usize;
430    Ok(match kind {
431        Kind::Constant => {
432            reader.i64()?;
433            "CONSTANT".to_string()
434        }
435        Kind::Packed => {
436            let mut widths = Vec::new();
437            let mut seen = 0;
438            while seen < count {
439                reader.i64()?;
440                let width = reader.u8()? as usize;
441                let wanted = (count - seen).min(VALUES);
442                if wanted == VALUES {
443                    for _ in 0..bitpack::packed_len::<u64>(width) {
444                        reader.u64()?;
445                    }
446                } else {
447                    reader.bytes(bitpack::tail_len(wanted, width))?;
448                }
449                widths.push(width);
450                seen += wanted;
451            }
452            let low = widths.iter().copied().min().unwrap_or(0);
453            let high = widths.iter().copied().max().unwrap_or(0);
454            // Square brackets rather than round ones, so that a reader and a test can both take a
455            // parenthesis to mean one more level of cascade and nothing else.
456            if low == high {
457                format!("FOR+BITPACK[{low}]")
458            } else {
459                format!("FOR+BITPACK[{low}..{high}]")
460            }
461        }
462        Kind::Delta => {
463            reader.i64()?;
464            format!("DELTA({})", describe_chunk(reader)?)
465        }
466        Kind::Rle => {
467            let values = describe_chunk(reader)?;
468            let lengths = describe_chunk(reader)?;
469            format!("RLE({values}, {lengths})")
470        }
471        Kind::Dict => {
472            let dictionary = describe_chunk(reader)?;
473            let codes = describe_chunk(reader)?;
474            format!("DICT({dictionary}, {codes})")
475        }
476        Kind::Sparse => {
477            reader.i64()?;
478            reader.u32()?;
479            let positions = describe_chunk(reader)?;
480            let exceptions = describe_chunk(reader)?;
481            format!("SPARSE({positions}, {exceptions})")
482        }
483    })
484}
485
486/// The distance from the frame of reference base, which is always representable in a `u64` because
487/// both ends came from an `i64` and the width of the difference is at most 65 bits minus the sign.
488fn offset_from(value: i64, base: i64) -> u64 {
489    (i128::from(value) - i128::from(base)) as u64
490}
491
492fn value_from(offset: u64, base: i64) -> i64 {
493    (i128::from(base) + i128::from(offset)) as i64
494}
495
496/// Zigzag, so that a column that counts down packs as narrowly as one that counts up. Without it a
497/// delta of -1 is 64 bits of ones.
498fn zigzag(value: i64) -> u64 {
499    ((value << 1) ^ (value >> 63)) as u64
500}
501
502fn unzigzag(value: u64) -> i64 {
503    ((value >> 1) as i64) ^ -((value & 1) as i64)
504}
505
506/// The zigzagged differences, or `None` if any difference is too wide to be one.
507///
508/// A column holding both `i64::MIN` and `i64::MAX` has a difference that does not fit in an `i64`,
509/// and rather than widening every delta array to 128 bits for a case that does not occur in data,
510/// the encoding declines to apply. `Packed` covers it.
511fn deltas(values: &[i64]) -> Option<Vec<i64>> {
512    let mut deltas = Vec::with_capacity(values.len().saturating_sub(1));
513    for pair in values.windows(2) {
514        let difference = i128::from(pair[1]) - i128::from(pair[0]);
515        let difference = i64::try_from(difference).ok()?;
516        deltas.push(zigzag(difference) as i64);
517    }
518    Some(deltas)
519}
520
521fn run_count(values: &[i64]) -> usize {
522    let mut runs = 0;
523    let mut previous = None;
524    for value in values {
525        if previous != Some(value) {
526            runs += 1;
527            previous = Some(value);
528        }
529    }
530    runs
531}
532
533fn runs(values: &[i64]) -> (Vec<i64>, Vec<i64>) {
534    let mut run_values: Vec<i64> = Vec::new();
535    let mut run_lengths: Vec<i64> = Vec::new();
536    for value in values {
537        if run_values.last() == Some(value) {
538            *run_lengths.last_mut().expect("a run length exists beside every run value") += 1;
539        } else {
540            run_values.push(*value);
541            run_lengths.push(1);
542        }
543    }
544    (run_values, run_lengths)
545}
546
547/// The distinct values in sorted order.
548///
549/// Sorted rather than in order of first appearance, because an ordered dictionary is what lets a
550/// range predicate become a code range instead of a code set, per section 6.7, and because the
551/// codes of a clustered column then run in order and delta encode.
552fn distinct_values(values: &[i64]) -> Vec<i64> {
553    let mut distinct = values.to_vec();
554    distinct.sort_unstable();
555    distinct.dedup();
556    distinct
557}
558
559fn codes_over(values: &[i64], dictionary: &[i64]) -> Vec<i64> {
560    values
561        .iter()
562        .map(|value| {
563            dictionary
564                .binary_search(value)
565                .expect("the dictionary is the distinct values of this chunk") as i64
566        })
567        .collect()
568}
569
570/// The most frequent value and how often it occurs, found without a hash map because the caller
571/// only asks when the column is already suspected of being nearly constant.
572fn dominant_value(values: &[i64]) -> Option<(i64, usize)> {
573    let mut sorted = values.to_vec();
574    sorted.sort_unstable();
575    let mut best: Option<(i64, usize)> = None;
576    let mut index = 0;
577    while index < sorted.len() {
578        let value = sorted[index];
579        let mut end = index;
580        while end < sorted.len() && sorted[end] == value {
581            end += 1;
582        }
583        let count = end - index;
584        if best.is_none_or(|(_, seen)| count > seen) {
585            best = Some((value, count));
586        }
587        index = end;
588    }
589    best
590}
591
592fn check_count(actual: usize, expected: usize) -> Result<()> {
593    if actual == expected {
594        Ok(())
595    } else {
596        Err(Error::internal(format!(
597            "a chunk says it holds {expected} values and decoded to {actual}"
598        )))
599    }
600}
601
602fn too_long(len: usize) -> Error {
603    Error::internal(format!("a chunk of {len} values is longer than the format allows"))
604}
605
606fn put_u8(out: &mut Vec<u8>, value: u8) {
607    out.push(value);
608}
609
610fn put_u32(out: &mut Vec<u8>, value: u32) {
611    out.extend_from_slice(&value.to_le_bytes());
612}
613
614fn put_u64(out: &mut Vec<u8>, value: u64) {
615    out.extend_from_slice(&value.to_le_bytes());
616}
617
618fn put_i64(out: &mut Vec<u8>, value: i64) {
619    out.extend_from_slice(&value.to_le_bytes());
620}
621
622#[cfg(test)]
623mod tests {
624    use super::*;
625
626    fn round_trip(values: &[i64]) -> Vec<u8> {
627        let bytes = encode(values).unwrap();
628        assert_eq!(decode(&bytes).unwrap(), values, "{}", describe(&bytes).unwrap());
629        bytes
630    }
631
632    fn kind_of(bytes: &[u8]) -> Kind {
633        Kind::from_tag(bytes[0]).unwrap()
634    }
635
636    /// The same xorshift the bit packing tests use, for the same reason.
637    struct Random(u64);
638
639    impl Random {
640        fn new() -> Self {
641            Self(0x9e37_79b9_7f4a_7c15)
642        }
643
644        fn next(&mut self) -> u64 {
645            self.0 ^= self.0 << 13;
646            self.0 ^= self.0 >> 7;
647            self.0 ^= self.0 << 17;
648            self.0
649        }
650    }
651
652    #[test]
653    fn an_empty_chunk_round_trips() {
654        let bytes = round_trip(&[]);
655        assert_eq!(bytes.len(), 5);
656    }
657
658    #[test]
659    fn a_constant_column_costs_thirteen_bytes_however_long_it_is() {
660        let bytes = round_trip(&vec![42; 1_000_000]);
661        assert_eq!(kind_of(&bytes), Kind::Constant);
662        assert_eq!(bytes.len(), 13);
663    }
664
665    #[test]
666    fn a_narrow_range_is_packed_at_the_width_of_the_range_and_not_of_the_type() {
667        // 100_000 values between 1000 and 1063 is 6 bits each, plus 9 bytes of header per 1024.
668        let mut random = Random::new();
669        let values: Vec<i64> = (0..100_000).map(|_| 1000 + (random.next() % 64) as i64).collect();
670        let bytes = round_trip(&values);
671        assert_eq!(kind_of(&bytes), Kind::Packed);
672        let packed = 100_000 * 6 / 8;
673        assert!(bytes.len() < packed + 2000, "{} bytes for {packed} of payload", bytes.len());
674        assert!(bytes.len() > packed, "{} bytes cannot hold {packed}", bytes.len());
675    }
676
677    #[test]
678    fn a_counter_becomes_deltas_and_then_a_constant() {
679        // The classic case and the reason DELTA exists. A million consecutive integers is a
680        // difference of 1 a million times, which is a constant chunk under the delta.
681        let values: Vec<i64> = (0..1_000_000).collect();
682        let bytes = round_trip(&values);
683        assert_eq!(kind_of(&bytes), Kind::Delta);
684        assert_eq!(describe(&bytes).unwrap(), "DELTA(CONSTANT)");
685        assert!(bytes.len() < 40, "{} bytes for a counter", bytes.len());
686    }
687
688    #[test]
689    fn a_column_that_counts_down_is_as_cheap_as_one_that_counts_up() {
690        // What zigzag is for. Without it every delta is -1, which is 64 bits of ones.
691        let up: Vec<i64> = (0..100_000).collect();
692        let down: Vec<i64> = (0..100_000).rev().collect();
693        assert_eq!(round_trip(&up).len(), round_trip(&down).len());
694    }
695
696    #[test]
697    fn long_runs_become_rle() {
698        let mut values = Vec::new();
699        for run in 0..1000 {
700            values.extend(std::iter::repeat_n(run % 7, 200));
701        }
702        let bytes = round_trip(&values);
703        assert_eq!(kind_of(&bytes), Kind::Rle);
704        assert!(bytes.len() < 2000, "{} bytes for 1000 runs", bytes.len());
705    }
706
707    #[test]
708    fn a_low_cardinality_column_becomes_a_dictionary() {
709        // Values that are far apart so that packing them directly is 30 bits each, and only 40 of
710        // them so that the codes are 6 bits each. The dictionary has to win by a factor of five.
711        let mut random = Random::new();
712        let dictionary: Vec<i64> = (0..40).map(|index| 1_000_000_000 + index * 7919).collect();
713        let values: Vec<i64> =
714            (0..100_000).map(|_| dictionary[(random.next() % 40) as usize]).collect();
715        let bytes = round_trip(&values);
716        assert_eq!(kind_of(&bytes), Kind::Dict);
717        assert!(bytes.len() < 100_000, "{} bytes", bytes.len());
718    }
719
720    #[test]
721    fn a_nearly_constant_column_becomes_sparse() {
722        let mut values = vec![0i64; 100_000];
723        for index in 0..300 {
724            values[index * 331] = 1 << 40;
725        }
726        let bytes = round_trip(&values);
727        assert_eq!(kind_of(&bytes), Kind::Sparse);
728        assert!(bytes.len() < 3000, "{} bytes for 300 exceptions", bytes.len());
729    }
730
731    #[test]
732    fn the_cascade_goes_more_than_one_level_deep() {
733        // The whole point of section 6.3. A dictionary over a clustered column produces codes that
734        // run in long stretches, and the run lengths of those are themselves compressible.
735        let mut values = Vec::new();
736        for index in 0..2000i64 {
737            values.extend(std::iter::repeat_n(1_000_000 + (index % 5) * 104_729, 100));
738        }
739        let bytes = round_trip(&values);
740        let shape = describe(&bytes).unwrap();
741        assert!(shape.contains('('), "{shape} is not a cascade");
742        assert!(bytes.len() < 4000, "{} bytes: {shape}", bytes.len());
743    }
744
745    #[test]
746    fn random_data_is_packed_at_full_width_and_costs_what_it_costs() {
747        // The case where nothing works, which has to come out at eight bytes a value plus change
748        // rather than at eight bytes a value plus a dictionary of every value in the column.
749        let mut random = Random::new();
750        let values: Vec<i64> = (0..10_000).map(|_| random.next() as i64).collect();
751        let bytes = round_trip(&values);
752        assert_eq!(kind_of(&bytes), Kind::Packed);
753        assert!(bytes.len() < 10_000 * 8 + 1000, "{} bytes", bytes.len());
754    }
755
756    #[test]
757    fn the_extremes_of_the_type_survive() {
758        // Every offset and every delta in here overflows something if the arithmetic is done in 64
759        // bits, which is why it is done in 128.
760        let values = vec![i64::MIN, i64::MAX, 0, -1, i64::MIN, i64::MAX];
761        round_trip(&values);
762        round_trip(&[i64::MIN; 3]);
763        round_trip(&[i64::MIN, i64::MIN + 1]);
764    }
765
766    #[test]
767    fn a_chunk_that_is_not_a_multiple_of_the_unit_round_trips() {
768        for len in [1, 2, 1023, 1024, 1025, 2047, 2049] {
769            let values: Vec<i64> = (0..len).map(|index| (index * 31 % 97) as i64).collect();
770            round_trip(&values);
771        }
772    }
773
774    #[test]
775    fn a_partial_unit_costs_its_own_values_and_not_a_whole_unit() {
776        // Three values that need 40 bits each. In the transposed layout a unit is 1024 values
777        // whether it holds them or not, so this would be 5 KB, and every nested array in a cascade
778        // is this short. It is 15 bytes of payload and 14 of header.
779        let values = vec![1i64 << 39, (1 << 39) + 7, 1 << 38];
780        let bytes = encode_as(Kind::Packed, &values, 0).unwrap().unwrap();
781        assert_eq!(bytes.len(), 5 + 9 + 15);
782        assert_eq!(decode(&bytes).unwrap(), values);
783    }
784
785    #[test]
786    fn the_frame_of_reference_is_per_unit_and_not_per_chunk() {
787        // A column that drifts, which is what a timestamp column and a clustered key both do. Each
788        // unit here spans 1023 and packs at 10 bits, and a base per chunk would pay the 22 bits the
789        // whole chunk spans on every value in it.
790        let values: Vec<i64> =
791            (0..4096i64).map(|index| (index / 1024) * 1_000_000 + (index % 1024)).collect();
792        let bytes = encode_as(Kind::Packed, &values, 0).unwrap().unwrap();
793        assert_eq!(describe(&bytes).unwrap(), "FOR+BITPACK[10]");
794        assert_eq!(decode(&bytes).unwrap(), values);
795    }
796
797    #[test]
798    fn every_candidate_that_applies_decodes_to_the_input() {
799        // The chooser only ever hands back the smallest, so without this the other five are only
800        // tested when they happen to win. Any of them being wrong is a wrong answer that appears
801        // when a column's distribution shifts.
802        let mut values = vec![5i64; 3000];
803        for (index, value) in values.iter_mut().enumerate() {
804            if index % 500 == 0 {
805                *value = index as i64;
806            }
807        }
808        let applicable = candidates(&values, 0);
809        assert!(applicable.len() >= 4, "{applicable:?}");
810        for kind in applicable {
811            let bytes = encode_as(kind, &values, 0).unwrap().unwrap();
812            assert_eq!(decode(&bytes).unwrap(), values, "{}", kind.name());
813        }
814    }
815
816    #[test]
817    fn the_chooser_picks_the_smallest_candidate_rather_than_the_first_that_applies() {
818        let mut values = vec![5i64; 3000];
819        values[1500] = 9;
820        let chosen = encode(&values).unwrap();
821        for (_, size) in candidate_sizes(&values).unwrap() {
822            assert!(chosen.len() <= size);
823        }
824    }
825
826    #[test]
827    fn a_truncated_chunk_is_an_error_and_not_a_panic() {
828        let bytes = encode(&[1, 2, 3, 4, 5]).unwrap();
829        for len in 0..bytes.len() {
830            let error = decode(&bytes[..len]).unwrap_err();
831            assert!(error.message().contains("chunk"), "{error}");
832        }
833    }
834
835    #[test]
836    fn trailing_bytes_are_an_error() {
837        let mut bytes = encode(&[1, 2, 3]).unwrap();
838        bytes.push(0);
839        let error = decode(&bytes).unwrap_err();
840        assert!(error.message().contains("left over"), "{error}");
841    }
842
843    #[test]
844    fn an_unknown_tag_is_an_error() {
845        let error = decode(&[99, 0, 0, 0, 0]).unwrap_err();
846        assert!(error.message().contains("unknown encoding tag"), "{error}");
847    }
848
849    #[test]
850    fn a_dictionary_code_outside_the_dictionary_is_an_error() {
851        // A corrupted or malicious chunk must not index out of bounds, and this is the one place in
852        // the decoder where a number that came off the disk is used as an index. Built by hand
853        // rather than by corrupting a real chunk, because a byte offset into an encoding that the
854        // chooser is free to change is a test that breaks for the wrong reason.
855        let mut bytes = vec![Kind::Dict.tag()];
856        put_u32(&mut bytes, 1);
857        bytes.extend_from_slice(&encode(&[10]).unwrap());
858        bytes.extend_from_slice(&encode(&[5]).unwrap());
859        let error = decode(&bytes).unwrap_err();
860        assert!(error.message().contains("not in the dictionary"), "{error}");
861    }
862
863    #[test]
864    fn a_negative_run_length_is_an_error() {
865        // The other number off the disk that the decoder would otherwise trust, and the one that
866        // would turn into an allocation of nine quintillion values.
867        let mut bytes = vec![Kind::Rle.tag()];
868        put_u32(&mut bytes, 4);
869        bytes.extend_from_slice(&encode(&[7]).unwrap());
870        bytes.extend_from_slice(&encode(&[-4]).unwrap());
871        let error = decode(&bytes).unwrap_err();
872        assert!(error.message().contains("negative"), "{error}");
873    }
874
875    #[test]
876    fn the_cascade_depth_is_bounded() {
877        // Without the limit a chooser that finds a dictionary of a dictionary of a dictionary would
878        // recurse until the values ran out, and the encode time of a wide column would be a
879        // surprise rather than a number.
880        let values: Vec<i64> = (0..50_000).map(|index| (index / 100) % 250).collect();
881        let bytes = round_trip(&values);
882        let shape = describe(&bytes).unwrap();
883        let depth = shape.matches('(').count();
884        assert!(depth <= MAX_DEPTH as usize, "{shape} is {depth} deep");
885    }
886
887    #[test]
888    fn candidate_sizes_reports_what_the_chooser_looked_at() {
889        let values: Vec<i64> = (0..5000).map(|index| index % 17).collect();
890        let sizes = candidate_sizes(&values).unwrap();
891        assert!(sizes.iter().any(|(kind, _)| *kind == Kind::Dict));
892        assert!(sizes.iter().any(|(kind, _)| *kind == Kind::Packed));
893        assert!(sizes.iter().all(|(_, size)| *size > 0));
894    }
895
896    #[test]
897    fn a_chunk_can_be_read_from_the_front_of_a_longer_buffer() {
898        // What a string column does. It writes an integer chunk of lengths into the middle of its
899        // own body and has to find the end of it again on the way back.
900        let first = encode(&[1, 2, 3]).unwrap();
901        let second: Vec<i64> = (0..3000).map(|index| index % 11).collect();
902        let second_bytes = encode(&second).unwrap();
903        let mut joined = first.clone();
904        joined.extend_from_slice(&second_bytes);
905        joined.extend_from_slice(b"and then something else");
906
907        let (values, used) = decode_prefix(&joined).unwrap();
908        assert_eq!(values, vec![1, 2, 3]);
909        assert_eq!(used, first.len());
910        let (more, used_again) = decode_prefix(&joined[used..]).unwrap();
911        assert_eq!(more, second);
912        assert_eq!(used_again, second_bytes.len());
913
914        let (text, described) = describe_prefix(&joined).unwrap();
915        assert_eq!(described, first.len());
916        assert_eq!(text, describe(&first).unwrap());
917    }
918
919    #[test]
920    fn a_truncated_chunk_is_still_an_error_when_read_as_a_prefix() {
921        let bytes = encode(&(0..2000).collect::<Vec<i64>>()).unwrap();
922        for len in 0..bytes.len() {
923            assert!(decode_prefix(&bytes[..len]).is_err(), "{len} bytes decoded");
924        }
925    }
926}