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