Skip to main content

rudb_encoding/
string.rs

1//! The string column, which is offsets, bytes, and the choice between compressing the bytes and
2//! not storing most of them at all.
3//!
4//! ClickBench `hits` is a string dataset before it is anything else. `URL`, `Referer`, `Title` and
5//! the referer derived columns are most of the 20.46 GB DuckDB writes for it, so most of what
6//! `spec/02-the-goal.md` promises on the resource axis has to come out of this file.
7//!
8//! ## The five shapes
9//!
10//! `CONSTANT` when every value is the same. `PLAIN`, which is lengths and raw bytes and is the
11//! baseline the others have to beat. `FSST`, which is a symbol table and the same lengths over
12//! compressed bytes. `DICT`, which is the distinct values and an array of codes. `FRONT`, which is
13//! the length of the prefix each value shares with the one before it and the rest of the value.
14//!
15//! `DICT_FSST` from the section 6.2 table is not a sixth shape. A dictionary's entries are a string
16//! column, and encoding them goes back through the same chooser, so a dictionary whose entries are
17//! FSST compressed is what the chooser produces on its own whenever that is smaller. The same
18//! recursion gives run length encoding of strings for free, because the codes are an integer chunk
19//! and `crate::integer` already knows what to do with a column of long runs.
20//!
21//! ## Why front coding is here
22//!
23//! The whole file measurement in M1 says the chooser produces 11.65 GB for `hits` against Parquet's
24//! 13.76 GB, and that `URL`, `Referer` and `OriginalURL` are 6.11 GB of it, and that on those three
25//! the chooser loses to Parquet's Snappy. The shape it picked on all three was `DICT(FSST[255])`,
26//! so the cascade was working and FSST was still losing.
27//!
28//! The reason is structural. FSST compresses each value on its own against a 255 symbol table, and
29//! a block compressor has the previous few kilobytes of the page to point back into. Two URLs that
30//! share a host and half a path are most of a back reference to each other and are nothing at all
31//! to a symbol table, which can only spend eight bytes of a symbol on the part they share and has
32//! to spend it again on every value. On a sorted dictionary of URLs the value before is the closest
33//! thing in the column to the value in hand, and the bytes they share are the redundancy Snappy was
34//! finding. Front coding is what reaches those bytes, and it composes with everything else here:
35//! the suffixes it leaves behind are a string column and go back through the chooser, so
36//! `DICT(FRONT(FSST))` is a shape the chooser can arrive at without anyone naming it.
37//!
38//! The chain has no restarts, so reading entry `n` means walking from entry zero. That is the right
39//! trade while a dictionary is decoded whole, which is what `decode` does. When something wants one
40//! entry out of a dictionary without materialising the rest, the answer is a restart every so many
41//! entries, and it costs one full value per block.
42//!
43//! ## Lengths, not offsets
44//!
45//! The usual layout is `n + 1` offsets and Arrow does it that way because a slice of an array has
46//! to be free. On disk the offsets are a monotonically increasing sequence whose differences are
47//! the lengths, and the differences are what compress: URL lengths in a real column are a few dozen
48//! distinct values in a narrow band, which the integer cascade turns into a handful of bits each,
49//! while the offsets themselves need enough bits to address the whole chunk. The integer cascade
50//! would find that by choosing DELTA, and storing lengths directly gets to the same place without
51//! spending a level of the cascade on it. Offsets are a prefix sum away and that is a decode time
52//! cost of one add per value.
53//!
54//! ## What is not here
55//!
56//! Nulls. A chunk here is N byte strings and an empty string is a value like any other. Validity is
57//! a bitmap that belongs to the column rather than to the encoding, per `spec/05-storage.md`, and
58//! `ROARING` in the section 6.2 table is what encodes it.
59//!
60//! Shared symbol tables and shared dictionaries across columns, which are section 6.4 and are the
61//! measurement this milestone exists for. Everything here is one column on its own, which is the
62//! baseline they get compared against.
63
64use rudb_common::{Error, Result};
65
66use crate::chooser::{Chooser, EXHAUSTIVE};
67use crate::fsst::SymbolTable;
68use crate::integer;
69use crate::lz;
70use crate::reader::Reader;
71
72/// How deep the recursion goes. A dictionary of a dictionary is not a thing, so this only has to
73/// stop the dictionary's own entries from being dictionary encoded again.
74const MAX_DEPTH: u8 = 2;
75
76/// How little sharing between neighbours is still worth offering front coding for, as one over
77/// this. A twentieth of the column is around where the prefix lengths start paying for themselves,
78/// and below it the candidate is an encode of the whole column that loses.
79const SHARE_DIVISOR: usize = 20;
80
81/// How few bytes is too few to bother looking for repeats in.
82///
83/// The matcher costs a hash table and a pass over the bytes whether it wins or not, and the chooser
84/// is exhaustive, so an ungated candidate is a tax on every string column in the database. Four
85/// kilobytes is about where a 32 KiB window has enough behind it to find anything.
86const LZ_FLOOR: usize = 4096;
87
88/// How many bytes of a column the symbol table is trained on.
89///
90/// The paper trains on about 16 KB. This is four times that, because training happens once per
91/// chunk here rather than once per block, and because the cost of a symbol that is only in the
92/// sample by accident is paid on every value in the chunk.
93pub(crate) const SAMPLE_BYTES: usize = 64 * 1024;
94
95/// What a string chunk is encoded as. The discriminant is the tag byte and is part of the format.
96#[derive(Debug, Clone, Copy, PartialEq, Eq)]
97pub enum Kind {
98    /// One value repeated.
99    Constant = 0,
100    /// Lengths and raw bytes.
101    Plain = 1,
102    /// Lengths, a symbol table, and FSST compressed bytes.
103    Fsst = 2,
104    /// The distinct values as a string chunk of their own, and codes into it as an integer chunk.
105    Dict = 3,
106    /// Shared prefix lengths as an integer chunk, and what is left of each value as a string chunk.
107    Front = 4,
108    /// Value lengths, copy lengths and copy offsets as integer chunks, and the bytes no copy
109    /// covered as a string chunk. See the `lz` module for what the matcher does and why it is here.
110    Lz = 5,
111}
112
113impl Kind {
114    fn tag(self) -> u8 {
115        self as u8
116    }
117
118    fn from_tag(tag: u8) -> Result<Self> {
119        match tag {
120            0 => Ok(Self::Constant),
121            1 => Ok(Self::Plain),
122            2 => Ok(Self::Fsst),
123            3 => Ok(Self::Dict),
124            4 => Ok(Self::Front),
125            5 => Ok(Self::Lz),
126            other => Err(Error::internal(format!("unknown string encoding tag {other}"))),
127        }
128    }
129
130    /// The name that goes in a report.
131    #[must_use]
132    pub fn name(self) -> &'static str {
133        match self {
134            Self::Constant => "CONSTANT",
135            Self::Plain => "PLAIN",
136            Self::Fsst => "FSST",
137            Self::Dict => "DICT",
138            Self::Front => "FRONT",
139            Self::Lz => "LZ",
140        }
141    }
142}
143
144/// Encodes a chunk of strings, choosing whatever comes out smallest.
145///
146/// Every candidate that applies is encoded in full and the smallest is kept, which is what this has
147/// always done and is what every size this crate has reported came out of. [`encode_with`] is the
148/// same thing with the search made swappable.
149///
150/// # Errors
151///
152/// If the chunk is longer than `u32::MAX` values, or if an encoding produces something its own
153/// decoder would not accept.
154pub fn encode(values: &[&[u8]]) -> Result<Vec<u8>> {
155    encode_with(values, &EXHAUSTIVE)
156}
157
158/// [`encode`] with somebody else deciding which candidates are worth encoding in full.
159///
160/// A chooser narrows the list and nothing else. It cannot offer a candidate that does not apply, so
161/// whatever it picks still has to encode the whole chunk and still has to decode, and the worst a
162/// bad one can do is come out bigger than [`encode`] would have.
163///
164/// # Errors
165///
166/// As [`encode`].
167pub fn encode_with(values: &[&[u8]], chooser: &dyn Chooser) -> Result<Vec<u8>> {
168    encode_at(values, 0, chooser)
169}
170
171/// A decoded chunk as one buffer with the values laid end to end, and where each one ends in it.
172///
173/// This is what the decoder builds and [`decode`] is a copy out of it. The cascade is why: a nest
174/// like `FRONT(LZ(FSST))` decodes three levels to produce one, and a level that hands its caller a
175/// `Vec<Vec<u8>>` has allocated once per value and copied every byte it holds. Three levels of that
176/// on a chunk of a thousand URLs is three thousand allocations to produce a thousand strings that
177/// the caller almost always wants back to back anyway.
178///
179/// It also makes the levels cheaper on their own terms. `PLAIN` is one `memcpy` of the whole
180/// payload because the values are already end to end in the file. `FRONT` copies a shared prefix
181/// out of the buffer it is writing into, so the previous value never has to be somewhere else.
182/// `LZ` replays straight into the buffer, which is what its copy offsets meant in the first place.
183#[derive(Debug, Clone, Default, PartialEq, Eq)]
184pub struct Flat {
185    bytes: Vec<u8>,
186    /// Where each value ends, so a value starts where the one before it ended and the last entry
187    /// is the length of `bytes`. Ends rather than offsets because a value is appended and its end
188    /// is what is known at that moment.
189    ends: Vec<usize>,
190}
191
192impl Flat {
193    fn with_capacity(count: usize, bytes: usize) -> Self {
194        Self { bytes: Vec::with_capacity(bytes), ends: Vec::with_capacity(count) }
195    }
196
197    fn push(&mut self, value: &[u8]) {
198        self.bytes.extend_from_slice(value);
199        self.ends.push(self.bytes.len());
200    }
201
202    /// Where the value at `index` starts, which is where the one before it ended.
203    fn start(&self, index: usize) -> usize {
204        if index == 0 { 0 } else { self.ends[index - 1] }
205    }
206
207    /// How many values the chunk holds.
208    #[must_use]
209    pub fn len(&self) -> usize {
210        self.ends.len()
211    }
212
213    /// Whether the chunk holds no values at all, which is not the same as holding empty ones.
214    #[must_use]
215    pub fn is_empty(&self) -> bool {
216        self.ends.is_empty()
217    }
218
219    /// The values laid end to end. A caller that already knows the boundaries, which is what a
220    /// global dictionary's offsets are, needs nothing else.
221    #[must_use]
222    pub fn bytes(&self) -> &[u8] {
223        &self.bytes
224    }
225
226    /// The value at `index`, or `None` past the end.
227    #[must_use]
228    pub fn get(&self, index: usize) -> Option<&[u8]> {
229        let end = *self.ends.get(index)?;
230        self.bytes.get(self.start(index)..end)
231    }
232
233    /// Every value in order.
234    pub fn iter(&self) -> impl Iterator<Item = &[u8]> {
235        let mut at = 0;
236        self.ends.iter().map(move |end| {
237            let value = self.bytes.get(at..*end).unwrap_or_default();
238            at = *end;
239            value
240        })
241    }
242
243    /// The buffer on its own, for a caller that wanted the bytes rather than the values.
244    #[must_use]
245    pub fn into_bytes(self) -> Vec<u8> {
246        self.bytes
247    }
248
249    /// The buffer and the ends that divide it, for a caller building its own layout over them.
250    ///
251    /// [`into_bytes`](Self::into_bytes) is enough for a caller that already knows where the values
252    /// end, which is what a global dictionary's stored offsets are. A caller that does not know has
253    /// only [`iter`](Self::iter), and walking that to build a run of boundaries copies out numbers
254    /// the chunk already holds. This hands both halves over and keeps the one allocation each.
255    #[must_use]
256    pub fn into_parts(self) -> (Vec<u8>, Vec<usize>) {
257        (self.bytes, self.ends)
258    }
259
260    fn into_values(self) -> Vec<Vec<u8>> {
261        let mut values = Vec::with_capacity(self.len());
262        let mut at = 0;
263        for end in &self.ends {
264            values.push(self.bytes[at..*end].to_vec());
265            at = *end;
266        }
267        values
268    }
269}
270
271/// Decodes a chunk written by [`encode`] without taking it apart into a value each.
272///
273/// # Errors
274///
275/// As [`decode`].
276pub fn decode_flat(bytes: &[u8]) -> Result<Flat> {
277    let mut reader = Reader::new(bytes);
278    let flat = decode_chunk(&mut reader)?;
279    if reader.remaining() != 0 {
280        return Err(Error::internal(format!(
281            "{} bytes left over after decoding a string chunk",
282            reader.remaining()
283        )));
284    }
285    Ok(flat)
286}
287
288/// Decodes a chunk that sits at the front of a longer buffer, and says how many bytes it took.
289///
290/// A column group holds one of these per column, and the decoder on that side cannot know where
291/// one ends until it has been read.
292///
293/// # Errors
294///
295/// As [`decode`], except that trailing bytes are what the caller asked about rather than an error.
296pub fn decode_prefix(bytes: &[u8]) -> Result<(Vec<Vec<u8>>, usize)> {
297    let mut reader = Reader::new(bytes);
298    let values = decode_chunk(&mut reader)?;
299    Ok((values.into_values(), reader.used()))
300}
301
302/// [`describe`] over a chunk at the front of a longer buffer, and how many bytes it took.
303///
304/// # Errors
305///
306/// As [`decode_prefix`].
307pub fn describe_prefix(bytes: &[u8]) -> Result<(String, usize)> {
308    let mut reader = Reader::new(bytes);
309    let text = describe_chunk(&mut reader)?;
310    Ok((text, reader.used()))
311}
312
313/// Decodes a chunk written by [`encode`].
314///
315/// # Errors
316///
317/// If the bytes are truncated, carry an unknown tag, or describe a chunk whose parts disagree.
318pub fn decode(bytes: &[u8]) -> Result<Vec<Vec<u8>>> {
319    Ok(decode_flat(bytes)?.into_values())
320}
321
322/// The size of every candidate that applies, for a report that wants to say what was chosen over
323/// what.
324///
325/// # Errors
326///
327/// As [`encode`].
328pub fn candidate_sizes(values: &[&[u8]]) -> Result<Vec<(Kind, usize)>> {
329    let mut sizes = Vec::new();
330    for kind in candidates(values, 0) {
331        if let Some(bytes) = encode_as(kind, values, 0, &EXHAUSTIVE)? {
332            sizes.push((kind, bytes.len()));
333        }
334    }
335    Ok(sizes)
336}
337
338/// Which candidates [`encode`] would try on this chunk, in the order it tries them.
339///
340/// The chooser is exhaustive, so this is also the list of encodes it pays for to return one of
341/// them. A caller measuring where the encode time goes needs the list separately from the sizes,
342/// because a candidate that is offered and turns out not to apply still costs whatever it spent
343/// finding that out.
344#[must_use]
345pub fn offered(values: &[&[u8]]) -> Vec<Kind> {
346    candidates(values, 0)
347}
348
349/// One candidate on its own, which is what the chooser calls once per entry in [`offered`].
350///
351/// `None` when the encoding does not apply, which is what the chooser treats as a candidate that
352/// did not run rather than as a failure. This is here so that the time the chooser spends can be
353/// attributed to the candidate that spent it, which is the measurement F2 wants before anybody
354/// replaces the exhaustive search with a sampled one. It is not how a writer encodes a chunk:
355/// [`encode`] is, and picking a kind by hand gives up the only thing the chooser is for.
356///
357/// # Errors
358///
359/// As [`encode`].
360pub fn encode_only(kind: Kind, values: &[&[u8]]) -> Result<Option<Vec<u8>>> {
361    encode_as(kind, values, 0, &EXHAUSTIVE)
362}
363
364/// How big one candidate comes out, which is all a sampling chooser needs from it.
365///
366/// The bytes are thrown away, so this says nothing [`encode_only`] does not. It is `pub(crate)` and
367/// separate so that the sampler in [`crate::chooser`] is not handing back buffers it will not read.
368pub(crate) fn size_as(kind: Kind, values: &[&[u8]], depth: u8) -> Result<Option<usize>> {
369    Ok(encode_as(kind, values, depth, &EXHAUSTIVE)?.map(|bytes| bytes.len()))
370}
371
372/// The shape a chunk was encoded as, as a line of text like `DICT(FSST, RLE(...))`.
373///
374/// # Errors
375///
376/// As [`decode`].
377pub fn describe(bytes: &[u8]) -> Result<String> {
378    let mut reader = Reader::new(bytes);
379    describe_chunk(&mut reader)
380}
381
382fn encode_at(values: &[&[u8]], depth: u8, chooser: &dyn Chooser) -> Result<Vec<u8>> {
383    let offered = candidates(values, depth);
384    let mut best: Option<Vec<u8>> = None;
385    for kind in chooser.narrow_strings(values, &offered, depth) {
386        let Some(bytes) = encode_as(kind, values, depth, chooser)? else {
387            continue;
388        };
389        if best.as_ref().is_none_or(|current| bytes.len() < current.len()) {
390            best = Some(bytes);
391        }
392    }
393    best.ok_or_else(|| Error::internal("no string encoding applied to the chunk"))
394}
395
396fn candidates(values: &[&[u8]], depth: u8) -> Vec<Kind> {
397    let mut kinds = vec![Kind::Plain];
398    if values.is_empty() {
399        return kinds;
400    }
401    if values.iter().all(|value| *value == values[0]) {
402        return vec![Kind::Constant];
403    }
404    kinds.push(Kind::Fsst);
405    if depth < MAX_DEPTH && has_duplicates(values) {
406        kinds.push(Kind::Dict);
407    }
408    if depth < MAX_DEPTH && sharing_of(values) >= total_len(values) / SHARE_DIVISOR {
409        kinds.push(Kind::Front);
410    }
411    if depth < MAX_DEPTH && total_len(values) >= LZ_FLOOR {
412        kinds.push(Kind::Lz);
413    }
414    kinds
415}
416
417/// How many bytes each value shares with the value before it, added up.
418///
419/// This is a full pass over the column, and it is here rather than on a sample because it is byte
420/// comparisons that stop at the first difference, which on a column with nothing to share stops
421/// immediately. Against training a symbol table and compressing the whole column, which is what
422/// offering the candidate would cost, it is not worth sampling.
423fn sharing_of(values: &[&[u8]]) -> usize {
424    let mut shared = 0;
425    for pair in values.windows(2) {
426        shared += shared_prefix(pair[0], pair[1]);
427    }
428    shared
429}
430
431/// Every value split into the bytes it shares with the value before it and the bytes it does not.
432///
433/// The suffixes point into the values, so this costs the prefix lengths and nothing else. It is
434/// shared with [`crate::multi`], which front codes a column before compressing it against a symbol
435/// table that belongs to the whole group.
436pub(crate) fn front_code<'a>(values: &[&'a [u8]]) -> (Vec<i64>, Vec<&'a [u8]>) {
437    let mut prefixes = Vec::with_capacity(values.len());
438    let mut suffixes: Vec<&'a [u8]> = Vec::with_capacity(values.len());
439    let mut previous: &[u8] = b"";
440    for value in values {
441        let value: &'a [u8] = value;
442        let shared = shared_prefix(previous, value);
443        prefixes.push(shared as i64);
444        suffixes.push(&value[shared..]);
445        previous = value;
446    }
447    (prefixes, suffixes)
448}
449
450/// The other half. The suffixes are consumed because the values are built out of them.
451///
452/// # Errors
453///
454/// If a prefix is negative or is longer than the value it is a prefix of, which is what a corrupt
455/// or hand written chunk looks like from here.
456pub(crate) fn front_decode(prefixes: &[i64], suffixes: Vec<Vec<u8>>) -> Result<Vec<Vec<u8>>> {
457    let mut values: Vec<Vec<u8>> = Vec::with_capacity(suffixes.len());
458    for (index, suffix) in suffixes.into_iter().enumerate() {
459        let shared = usize::try_from(prefixes[index])
460            .map_err(|_| Error::internal("a negative shared prefix length"))?;
461        let previous: &[u8] = if index == 0 { b"" } else { &values[index - 1] };
462        if shared > previous.len() {
463            return Err(Error::internal(format!(
464                "a value shares {shared} bytes with a value {} bytes long",
465                previous.len()
466            )));
467        }
468        let mut value = Vec::with_capacity(shared + suffix.len());
469        value.extend_from_slice(&previous[..shared]);
470        value.extend_from_slice(&suffix);
471        values.push(value);
472    }
473    Ok(values)
474}
475
476fn shared_prefix(previous: &[u8], value: &[u8]) -> usize {
477    let limit = previous.len().min(value.len());
478    let mut shared = 0;
479    while shared < limit && previous[shared] == value[shared] {
480        shared += 1;
481    }
482    shared
483}
484
485fn total_len(values: &[&[u8]]) -> usize {
486    values.iter().map(|value| value.len()).sum()
487}
488
489fn encode_as(
490    kind: Kind,
491    values: &[&[u8]],
492    depth: u8,
493    chooser: &dyn Chooser,
494) -> Result<Option<Vec<u8>>> {
495    let mut out = vec![kind.tag()];
496    put_u32(&mut out, u32::try_from(values.len()).map_err(|_| too_long(values.len()))?);
497    match kind {
498        Kind::Constant => {
499            let Some(first) = values.first() else {
500                return Ok(None);
501            };
502            if values.iter().any(|value| value != first) {
503                return Ok(None);
504            }
505            put_u32(&mut out, u32::try_from(first.len()).map_err(|_| too_long(first.len()))?);
506            out.extend_from_slice(first);
507        }
508        Kind::Plain => {
509            out.extend_from_slice(&encode_lengths(values, chooser)?);
510            for value in values {
511                out.extend_from_slice(value);
512            }
513        }
514        Kind::Fsst => {
515            let sample = sample_of(values);
516            let table = SymbolTable::train(&sample);
517            if table.is_empty() {
518                return Ok(None);
519            }
520            let mut compressed = Vec::new();
521            let mut lengths = Vec::with_capacity(values.len());
522            for value in values {
523                let before = compressed.len();
524                table.compress(value, &mut compressed);
525                lengths.push((compressed.len() - before) as i64);
526            }
527            table.serialize(&mut out);
528            out.extend_from_slice(&integer::encode_with(&lengths, chooser)?);
529            out.extend_from_slice(&compressed);
530        }
531        Kind::Dict => {
532            let (entries, codes) = dictionary_of(values);
533            if entries.is_empty() {
534                return Ok(None);
535            }
536            out.extend_from_slice(&encode_at(&entries, depth + 1, chooser)?);
537            out.extend_from_slice(&integer::encode_with(&codes, chooser)?);
538        }
539        Kind::Front => {
540            let (prefixes, suffixes) = front_code(values);
541            out.extend_from_slice(&integer::encode_with(&prefixes, chooser)?);
542            out.extend_from_slice(&encode_at(&suffixes, depth + 1, chooser)?);
543        }
544        Kind::Lz => {
545            let mut joined = Vec::with_capacity(total_len(values));
546            let mut sizes = Vec::with_capacity(values.len());
547            for value in values {
548                joined.extend_from_slice(value);
549                sizes.push(value.len() as i64);
550            }
551            let tokens = lz::tokens_of(&joined);
552            out.extend_from_slice(&integer::encode_with(&sizes, chooser)?);
553            out.extend_from_slice(&integer::encode_with(&tokens.lengths, chooser)?);
554            out.extend_from_slice(&integer::encode_with(&tokens.offsets, chooser)?);
555            out.extend_from_slice(&encode_at(&tokens.literals, depth + 1, chooser)?);
556        }
557    }
558    Ok(Some(out))
559}
560
561fn decode_chunk(reader: &mut Reader<'_>) -> Result<Flat> {
562    let kind = Kind::from_tag(reader.u8()?)?;
563    let count = reader.u32()? as usize;
564    match kind {
565        Kind::Constant => {
566            let len = reader.u32()? as usize;
567            let value = reader.bytes(len)?;
568            let mut flat = Flat::with_capacity(count, len.saturating_mul(count));
569            for _ in 0..count {
570                flat.push(value);
571            }
572            Ok(flat)
573        }
574        Kind::Plain => {
575            let lengths = decode_lengths(reader, count)?;
576            // One copy of the whole payload rather than one a value, which the file already laid
577            // out end to end and which is the layout wanted back.
578            let total = sum_of(&lengths)?;
579            let payload = reader.bytes(total)?;
580            let mut flat = Flat::with_capacity(count, total);
581            flat.bytes.extend_from_slice(payload);
582            let mut at = 0;
583            for length in lengths {
584                at += length;
585                flat.ends.push(at);
586            }
587            Ok(flat)
588        }
589        Kind::Fsst => {
590            let runs = read_compressed(reader, count)?;
591            let mut flat = Flat::with_capacity(count, runs.payload.len());
592            let mut at = 0;
593            for index in 0..count {
594                runs.run_into(index, &mut at, &mut flat.bytes)?;
595                flat.ends.push(flat.bytes.len());
596            }
597            Ok(flat)
598        }
599        Kind::Dict => {
600            let dictionary = decode_chunk(reader)?;
601            let codes = decode_integers(reader)?;
602            if codes.len() != count {
603                return Err(Error::internal(format!(
604                    "a dictionary chunk says it holds {count} values and has {} codes",
605                    codes.len()
606                )));
607            }
608            let mut flat = Flat::with_capacity(count, dictionary.bytes.len());
609            for code in codes {
610                let entry =
611                    usize::try_from(code).ok().and_then(|index| dictionary.get(index)).ok_or_else(
612                        || Error::internal(format!("code {code} is not in the dictionary")),
613                    )?;
614                flat.push(entry);
615            }
616            Ok(flat)
617        }
618        Kind::Front => {
619            let prefixes = decode_integers(reader)?;
620            let suffixes = decode_chunk(reader)?;
621            if prefixes.len() != count || suffixes.len() != count {
622                return Err(Error::internal(format!(
623                    "a front coded chunk says it holds {count} values and has {} prefixes and {} suffixes",
624                    prefixes.len(),
625                    suffixes.len()
626                )));
627            }
628            // The shared prefix is copied out of the buffer being written into, so a value never
629            // has to exist anywhere but where it belongs.
630            let mut flat = Flat::with_capacity(count, suffixes.bytes.len());
631            for (index, prefix) in prefixes.iter().enumerate() {
632                let shared = usize::try_from(*prefix)
633                    .map_err(|_| Error::internal("a negative shared prefix length"))?;
634                let (from, previous) = if index == 0 {
635                    (0, 0)
636                } else {
637                    (flat.start(index - 1), flat.ends[index - 1] - flat.start(index - 1))
638                };
639                if shared > previous {
640                    return Err(Error::internal(format!(
641                        "a value shares {shared} bytes with a value {previous} bytes long"
642                    )));
643                }
644                flat.bytes.extend_from_within(from..from + shared);
645                flat.bytes.extend_from_slice(suffixes.get(index).expect("in range"));
646                flat.ends.push(flat.bytes.len());
647            }
648            Ok(flat)
649        }
650        Kind::Lz => {
651            let sizes = decode_integers(reader)?;
652            let lengths = decode_integers(reader)?;
653            let offsets = decode_integers(reader)?;
654            if sizes.len() != count {
655                return Err(Error::internal(format!(
656                    "a matched chunk says it holds {count} values and has {} lengths",
657                    sizes.len()
658                )));
659            }
660            let mut total = 0usize;
661            let mut widths = Vec::with_capacity(count);
662            for size in sizes {
663                let width = usize::try_from(size)
664                    .map_err(|_| Error::internal("a negative string length"))?;
665                total = total
666                    .checked_add(width)
667                    .ok_or_else(|| Error::internal("a string chunk longer than memory"))?;
668                widths.push(width);
669            }
670            // The copies point back into the bytes already replayed, which is the buffer the values
671            // are going into, so the replay is the decode and there is nothing to cut up after it.
672            let mut flat = Flat::with_capacity(count, total);
673            replay_literals(reader, &lengths, &offsets, &mut flat.bytes)?;
674            if flat.bytes.len() != total {
675                return Err(Error::internal(format!(
676                    "a matched chunk rebuilt {} bytes where its lengths add up to {total}",
677                    flat.bytes.len()
678                )));
679            }
680            let mut at = 0;
681            for width in widths {
682                at += width;
683                flat.ends.push(at);
684            }
685            Ok(flat)
686        }
687    }
688}
689
690/// A compressed chunk's symbol table and its runs, left where the file put them.
691///
692/// Reading a compressed chunk into this rather than straight into a buffer is what lets a run be
693/// decompressed where the run belongs. The payload is one slice, the run boundaries come from the
694/// length array, and so asking for a run is a decompress of a subslice and nothing else.
695struct Compressed<'a> {
696    /// The table the runs were compressed against.
697    table: SymbolTable,
698    /// How many compressed bytes each run holds, in order.
699    lengths: Vec<usize>,
700    /// Every run's compressed bytes, end to end.
701    payload: &'a [u8],
702}
703
704impl Compressed<'_> {
705    /// Decompresses run `index` onto the end of `out`, with `at` saying where the run starts.
706    ///
707    /// The caller carries the offset because the runs are asked for in order, and adding a length
708    /// per run is cheaper than the prefix sum the alternative wants.
709    ///
710    /// # Errors
711    ///
712    /// If there is no such run, if it runs off the end of the payload, or if it does not decompress.
713    fn run_into(&self, index: usize, at: &mut usize, out: &mut Vec<u8>) -> Result<()> {
714        let length = *self
715            .lengths
716            .get(index)
717            .ok_or_else(|| Error::internal(format!("run {index} is not in the chunk")))?;
718        let end = at
719            .checked_add(length)
720            .ok_or_else(|| Error::internal("a compressed chunk longer than memory"))?;
721        let run = self
722            .payload
723            .get(*at..end)
724            .ok_or_else(|| Error::internal("a compressed run is past the end of its chunk"))?;
725        *at = end;
726        self.table.decompress(run, out)
727    }
728}
729
730/// Reads a compressed chunk's table, run lengths and payload without decompressing any of it.
731///
732/// The tag and the count have already been read.
733///
734/// # Errors
735///
736/// If the table does not deserialize, if the length array is not `count` long, or if the lengths
737/// add up to more than the chunk has left.
738fn read_compressed<'a>(reader: &mut Reader<'a>, count: usize) -> Result<Compressed<'a>> {
739    let (table, used) = SymbolTable::deserialize(reader.rest())?;
740    reader.skip(used)?;
741    let lengths = decode_lengths(reader, count)?;
742    // The compressed total is what the payload holds and it is also the only sane guess at the
743    // decompressed one, so it is checked before it is believed.
744    let compressed_len = sum_of(&lengths)?;
745    if compressed_len > reader.remaining() {
746        return Err(Error::internal(format!(
747            "a compressed chunk says it holds {compressed_len} bytes and has {}",
748            reader.remaining()
749        )));
750    }
751    let payload = reader.bytes(compressed_len)?;
752    Ok(Compressed { table, lengths, payload })
753}
754
755/// Replays a matched chunk's tokens, reading the literal runs out of the nested chunk holding them.
756///
757/// The nested chunk is decoded into a buffer and copied out of, the way anything nested is, unless
758/// it is compressed. On the ClickBench `URL` column it always is, and there a block of a thousand
759/// values holds about eight thousand seven hundred literal runs, so that buffer is the whole
760/// block's bytes and copying the runs out of it writes every one of them a second time.
761/// Decompressing a run straight to where it belongs skips the buffer, the length array that would
762/// cut it up, and that second pass over the bytes.
763///
764/// # Errors
765///
766/// Whatever reading the literals or replaying the tokens reports.
767fn replay_literals(
768    reader: &mut Reader<'_>,
769    lengths: &[i64],
770    offsets: &[i64],
771    out: &mut Vec<u8>,
772) -> Result<()> {
773    if reader.rest().first() == Some(&Kind::Fsst.tag()) {
774        reader.u8()?;
775        let runs = reader.u32()? as usize;
776        let compressed = read_compressed(reader, runs)?;
777        let mut at = 0;
778        return lz::replay(
779            runs,
780            |index, into| compressed.run_into(index, &mut at, into),
781            lengths,
782            offsets,
783            out,
784        );
785    }
786    let literals = decode_chunk(reader)?;
787    lz::rebuild_into(&literals, lengths, offsets, out)
788}
789
790fn describe_chunk(reader: &mut Reader<'_>) -> Result<String> {
791    let kind = Kind::from_tag(reader.u8()?)?;
792    let count = reader.u32()? as usize;
793    Ok(match kind {
794        Kind::Constant => {
795            let len = reader.u32()? as usize;
796            reader.bytes(len)?;
797            "CONSTANT".to_string()
798        }
799        Kind::Plain => {
800            let (shape, lengths) = describe_lengths(reader, count)?;
801            reader.skip(lengths.iter().sum())?;
802            format!("PLAIN({shape})")
803        }
804        Kind::Fsst => {
805            let (table, used) = SymbolTable::deserialize(reader.rest())?;
806            reader.skip(used)?;
807            let (shape, lengths) = describe_lengths(reader, count)?;
808            reader.skip(lengths.iter().sum())?;
809            format!("FSST[{}]({shape})", table.len())
810        }
811        Kind::Dict => {
812            let entries = describe_chunk(reader)?;
813            let codes = describe_integers(reader)?;
814            format!("DICT({entries}, {codes})")
815        }
816        Kind::Front => {
817            let prefixes = describe_integers(reader)?;
818            let suffixes = describe_chunk(reader)?;
819            format!("FRONT({prefixes}, {suffixes})")
820        }
821        Kind::Lz => {
822            let sizes = describe_integers(reader)?;
823            let lengths = describe_integers(reader)?;
824            let offsets = describe_integers(reader)?;
825            let literals = describe_chunk(reader)?;
826            format!("LZ({sizes}, {lengths}, {offsets}, {literals})")
827        }
828    })
829}
830
831/// The shape of the length array and the lengths themselves, because a describe has to walk past
832/// the payload to leave the reader where the next chunk starts and the payload size is the sum of
833/// the lengths.
834fn describe_lengths(reader: &mut Reader<'_>, count: usize) -> Result<(String, Vec<usize>)> {
835    let (shape, _) = integer::describe_prefix(reader.rest())?;
836    let lengths = decode_lengths(reader, count)?;
837    Ok((shape, lengths))
838}
839
840fn encode_lengths(values: &[&[u8]], chooser: &dyn Chooser) -> Result<Vec<u8>> {
841    let lengths: Vec<i64> = values.iter().map(|value| value.len() as i64).collect();
842    integer::encode_with(&lengths, chooser)
843}
844
845fn decode_lengths(reader: &mut Reader<'_>, count: usize) -> Result<Vec<usize>> {
846    let lengths = decode_integers(reader)?;
847    if lengths.len() != count {
848        return Err(Error::internal(format!(
849            "a string chunk says it holds {count} values and has {} lengths",
850            lengths.len()
851        )));
852    }
853    lengths
854        .into_iter()
855        .map(|length| {
856            usize::try_from(length).map_err(|_| Error::internal("a negative string length"))
857        })
858        .collect()
859}
860
861/// How long the values add up to, refusing a length array that adds up to more than memory.
862///
863/// A truncated chunk used to be caught by the read of the value that ran off the end. Reading the
864/// payload in one go means the total has to be trusted before the read rather than after it, and a
865/// corrupt length array is the only thing that could overflow it.
866fn sum_of(lengths: &[usize]) -> Result<usize> {
867    lengths
868        .iter()
869        .try_fold(0usize, |total, length| total.checked_add(*length))
870        .ok_or_else(|| Error::internal("a string chunk longer than memory"))
871}
872
873/// Reads one nested integer chunk. The integer decoder wants a slice of exactly its own chunk and
874/// the reader does not know how long that is, so it decodes from the rest of the buffer and is told
875/// afterwards how much it used.
876fn decode_integers(reader: &mut Reader<'_>) -> Result<Vec<i64>> {
877    let (values, used) = integer::decode_prefix(reader.rest())?;
878    reader.skip(used)?;
879    Ok(values)
880}
881
882fn describe_integers(reader: &mut Reader<'_>) -> Result<String> {
883    let (text, used) = integer::describe_prefix(reader.rest())?;
884    reader.skip(used)?;
885    Ok(text)
886}
887
888/// A sample of the column spread across the whole of it, taken at random skips rather than at a
889/// fixed stride.
890///
891/// Section 6.3 makes the point about choosing an encoding from a sample and it applies at least as
892/// much to training a symbol table. Column data is frequently sorted or clustered, so the first
893/// 64 KB of a URL column is the hosts that sort first and a table trained on it escapes most of the
894/// rest of the column.
895///
896/// The skips are random rather than fixed because a fixed stride aliases. Column data is also
897/// frequently periodic, and a stride that shares a factor with the period samples one phase of it
898/// and never sees the others. That is not a hypothetical: the first version of this took every
899/// `n`th value, and on a test column whose values cycle with a period that the stride happened to
900/// divide, the table it trained was 3.4 times worse than one trained on the whole column, because
901/// it learned eight byte symbols that only line up with the phase it saw and had no shorter symbols
902/// left to fall back on.
903///
904/// The generator is a fixed seed xorshift, so the sample is a function of the column and encoding
905/// the same values twice produces the same bytes.
906pub(crate) fn sample_of<'a>(values: &[&'a [u8]]) -> Vec<&'a [u8]> {
907    sample_bytes_of(values, SAMPLE_BYTES)
908}
909
910/// [`sample_of`] with the byte budget spelled out, for a caller training one table over several
911/// columns that has to split the budget between them.
912pub(crate) fn sample_bytes_of<'a>(values: &[&'a [u8]], budget: usize) -> Vec<&'a [u8]> {
913    let budget = budget.max(1);
914    let total: usize = values.iter().map(|value| value.len()).sum();
915    if total <= budget {
916        return values.to_vec();
917    }
918    let stride = total.div_ceil(budget).max(1);
919    let span = (stride * 2 - 1).max(1) as u64;
920    let mut state = 0x2545_f491_4f6c_dd1du64;
921    let mut sample = Vec::with_capacity(values.len() / stride + 1);
922    let mut at = 0usize;
923    while at < values.len() {
924        sample.push(values[at]);
925        state ^= state << 13;
926        state ^= state >> 7;
927        state ^= state << 17;
928        at += 1 + (state % span) as usize;
929    }
930    sample
931}
932
933/// The distinct values in sorted order and the code of every value, in one pass over one sort.
934///
935/// The dictionary is sorted for the same reason the integer one is: an ordered dictionary turns a
936/// range predicate into a code range rather than a code set, and front coding over the entries needs
937/// them sorted anyway.
938///
939/// It sorts a permutation of indices rather than the values, which is the whole point. Sorting the
940/// values means copying every one of them onto the heap first, and the codes then have to be found
941/// by searching the dictionary back for each value, which is a binary search of string comparisons
942/// per row. Walking the permutation gives the codes away for free, because the position a value
943/// sorted to is the position its code was assigned at.
944fn dictionary_of<'a>(values: &[&'a [u8]]) -> (Vec<&'a [u8]>, Vec<i64>) {
945    let mut order: Vec<u32> = (0..values.len() as u32).collect();
946    order.sort_unstable_by(|left, right| values[*left as usize].cmp(values[*right as usize]));
947    let mut entries: Vec<&'a [u8]> = Vec::new();
948    let mut codes = vec![0i64; values.len()];
949    for &index in &order {
950        let value = values[index as usize];
951        if entries.last() != Some(&value) {
952            entries.push(value);
953        }
954        codes[index as usize] = (entries.len() - 1) as i64;
955    }
956    (entries, codes)
957}
958
959/// Whether any value appears twice, which is the only thing the candidate list wants to know.
960///
961/// This used to build the whole sorted dictionary and compare its length against the input, which
962/// is a copy of the chunk and a sort of it paid on every chunk at every level whether the dictionary
963/// was ever encoded or not. It is a linear probe over hashes instead: expected O(n), no allocation
964/// per value, and it stops at the first duplicate it finds, which on a column with any repetition at
965/// all is immediately.
966///
967/// A hash collision is resolved by comparing the bytes, so the answer is exact rather than probable.
968fn has_duplicates(values: &[&[u8]]) -> bool {
969    let Some(slots) = values.len().checked_mul(2).map(usize::next_power_of_two) else {
970        return false;
971    };
972    let mask = slots - 1;
973    let mut table = vec![u32::MAX; slots];
974    for (index, value) in values.iter().enumerate() {
975        let mut at = hash_of(value) as usize & mask;
976        loop {
977            let held = table[at];
978            if held == u32::MAX {
979                table[at] = index as u32;
980                break;
981            }
982            if values[held as usize] == *value {
983                return true;
984            }
985            at = (at + 1) & mask;
986        }
987    }
988    false
989}
990
991/// FNV-1a over the bytes, eight at a time.
992///
993/// Good enough for a table that verifies every hit, and it is not part of the format, so nothing
994/// depends on which hash this is. Eight bytes at a time because a URL column is long values and a
995/// byte at a time over a hundred bytes of every one of 122,880 rows is the loop this is here to
996/// avoid.
997fn hash_of(value: &[u8]) -> u64 {
998    let mut hash = 0xcbf2_9ce4_8422_2325_u64;
999    let mut chunks = value.chunks_exact(8);
1000    for chunk in &mut chunks {
1001        let word = u64::from_le_bytes(chunk.try_into().expect("chunks_exact(8) gives eight bytes"));
1002        hash = (hash ^ word).wrapping_mul(0x1_0000_01b3);
1003    }
1004    for byte in chunks.remainder() {
1005        hash = (hash ^ u64::from(*byte)).wrapping_mul(0x1_0000_01b3);
1006    }
1007    (hash ^ (value.len() as u64)).wrapping_mul(0x1_0000_01b3)
1008}
1009
1010fn too_long(len: usize) -> Error {
1011    Error::internal(format!("a string chunk of {len} is longer than the format allows"))
1012}
1013
1014fn put_u32(out: &mut Vec<u8>, value: u32) {
1015    out.extend_from_slice(&value.to_le_bytes());
1016}
1017
1018#[cfg(test)]
1019mod tests {
1020    use super::*;
1021
1022    fn urls(count: usize) -> Vec<Vec<u8>> {
1023        let hosts = ["www.example.com", "shop.example.com", "news.other.example.org"];
1024        let paths = ["/index.html", "/catalog/item", "/search", "/user/profile/settings"];
1025        (0..count)
1026            .map(|index| {
1027                let host = hosts[index % hosts.len()];
1028                let path = paths[(index / 3) % paths.len()];
1029                format!("http://{host}{path}?session={}&ref=google", index * 7).into_bytes()
1030            })
1031            .collect()
1032    }
1033
1034    /// The same values with a scrambled identifier stuck on the front of each, for the tests that
1035    /// need neighbouring values to have nothing in common. Shuffling the order is not enough,
1036    /// because two URLs picked at random still agree on a scheme and often on a host.
1037    fn keyed(values: Vec<Vec<u8>>) -> Vec<Vec<u8>> {
1038        values
1039            .into_iter()
1040            .enumerate()
1041            .map(|(index, value)| {
1042                let key = (index as u64).wrapping_mul(0x9e37_79b9_7f4a_7c15) % 1_000_000_007;
1043                let mut out = format!("{key:010}/").into_bytes();
1044                out.extend_from_slice(&value);
1045                out
1046            })
1047            .collect()
1048    }
1049
1050    fn borrow(values: &[Vec<u8>]) -> Vec<&[u8]> {
1051        values.iter().map(Vec::as_slice).collect()
1052    }
1053
1054    fn round_trip(values: &[Vec<u8>]) -> Vec<u8> {
1055        let borrowed = borrow(values);
1056        let bytes = encode(&borrowed).unwrap();
1057        let back = decode(&bytes).unwrap();
1058        assert_eq!(back, values, "{}", describe(&bytes).unwrap());
1059        check_flat(&bytes, values);
1060        bytes
1061    }
1062
1063    /// The flat form holds the same values and lays them out the way a caller with its own offsets
1064    /// expects. Called from [`round_trip`], so every shape any test in here reaches is checked.
1065    fn check_flat(bytes: &[u8], values: &[Vec<u8>]) {
1066        let flat = decode_flat(bytes).unwrap();
1067        let shape = describe(bytes).unwrap();
1068        assert_eq!(flat.len(), values.len(), "{shape}");
1069        assert_eq!(flat.iter().collect::<Vec<_>>(), borrow(values), "{shape}");
1070        assert_eq!(flat.bytes(), values.concat(), "{shape}");
1071        assert_eq!(flat.get(values.len()), None, "{shape}");
1072    }
1073
1074    fn kind_of(bytes: &[u8]) -> Kind {
1075        Kind::from_tag(bytes[0]).unwrap()
1076    }
1077
1078    #[test]
1079    fn every_shape_decodes_flat_to_what_it_decodes_split() {
1080        // round_trip only sees the shape the chooser picked, which on any one column is one of the
1081        // six. This walks all of them, so PLAIN reading its payload in one go and FRONT copying a
1082        // prefix out of the buffer it is filling are both covered on data they apply to.
1083        let columns =
1084            [urls(600), keyed(urls(600)), vec![b"same".to_vec(); 400], vec![Vec::new(); 7]];
1085        for values in &columns {
1086            let borrowed = borrow(values);
1087            for kind in offered(&borrowed) {
1088                let Some(bytes) = encode_only(kind, &borrowed).unwrap() else {
1089                    continue;
1090                };
1091                assert_eq!(decode(&bytes).unwrap(), *values, "{}", kind.name());
1092                let flat = decode_flat(&bytes).unwrap();
1093                assert_eq!(flat.iter().collect::<Vec<_>>(), borrowed, "{}", kind.name());
1094                assert_eq!(flat.bytes(), values.concat(), "{}", kind.name());
1095            }
1096        }
1097    }
1098
1099    #[test]
1100    fn a_front_coded_chunk_that_shares_more_than_it_has_is_an_error() {
1101        // The prefix chain is the one place the flat decoder reads back out of the buffer it is
1102        // filling, so a prefix longer than the value before it is what would hand back somebody
1103        // else's bytes rather than fail. Built by hand because no encoder produces one.
1104        let suffixes: [&[u8]; 2] = [b"abc", b"x"];
1105        let mut bytes = vec![Kind::Front.tag()];
1106        put_u32(&mut bytes, 2);
1107        bytes.extend_from_slice(&integer::encode(&[0, 9]).unwrap());
1108        bytes.extend_from_slice(&encode_only(Kind::Plain, &suffixes).unwrap().unwrap());
1109        let error = decode_flat(&bytes).expect_err("a nine byte prefix of a three byte value");
1110        assert_eq!(error.message(), "a value shares 9 bytes with a value 3 bytes long");
1111        assert_eq!(decode(&bytes).unwrap_err().message(), error.message());
1112    }
1113
1114    #[test]
1115    fn the_dictionary_is_sorted_and_the_codes_point_back_at_the_values() {
1116        // The two things the dictionary path has to get right, and the reason it is one function
1117        // now rather than a sort followed by a binary search per row.
1118        let values = vec![
1119            b"pear".to_vec(),
1120            b"apple".to_vec(),
1121            b"pear".to_vec(),
1122            b"cherry".to_vec(),
1123            b"apple".to_vec(),
1124        ];
1125        let borrowed = borrow(&values);
1126        let (entries, codes) = dictionary_of(&borrowed);
1127        assert_eq!(entries, vec![b"apple".as_slice(), b"cherry".as_slice(), b"pear".as_slice()]);
1128        assert_eq!(codes, vec![2, 0, 2, 1, 0]);
1129        for (code, value) in codes.iter().zip(&borrowed) {
1130            assert_eq!(entries[*code as usize], *value);
1131        }
1132    }
1133
1134    #[test]
1135    fn a_column_with_nothing_repeated_has_no_duplicates_and_one_with_anything_does() {
1136        let distinct: Vec<Vec<u8>> =
1137            (0..5000).map(|index| format!("value-{index}").into_bytes()).collect();
1138        assert!(!has_duplicates(&borrow(&distinct)));
1139
1140        // One repeat at the far end, so a check that gave up early would miss it.
1141        let mut repeated = distinct.clone();
1142        repeated.push(b"value-0".to_vec());
1143        assert!(has_duplicates(&borrow(&repeated)));
1144
1145        assert!(!has_duplicates(&borrow(&Vec::new())));
1146        assert!(!has_duplicates(&borrow(&[b"one".to_vec()])));
1147        assert!(has_duplicates(&borrow(&vec![b"same".to_vec(); 2])));
1148    }
1149
1150    #[test]
1151    fn long_values_that_differ_only_at_the_end_are_not_confused_for_each_other() {
1152        // The hash is eight bytes at a time and the table verifies every hit, so this is the case
1153        // that says the verify is really there rather than the hash being trusted.
1154        let stem = "http://www.example.com/a/very/long/path/that/goes/on?session=";
1155        let values: Vec<Vec<u8>> =
1156            (0..2000).map(|index| format!("{stem}{index}").into_bytes()).collect();
1157        assert!(!has_duplicates(&borrow(&values)));
1158        let (entries, codes) = dictionary_of(&borrow(&values));
1159        assert_eq!(entries.len(), values.len());
1160        assert_eq!(codes.len(), values.len());
1161    }
1162
1163    #[test]
1164    fn what_the_chooser_returns_is_the_smallest_of_what_it_was_offered() {
1165        // `offered` and `encode_only` are what `cargo xtask encode` splits the chooser's seconds
1166        // with, so they have to describe the chooser that actually runs rather than a second copy
1167        // of its rules that drifts. This is the assertion that keeps the two the same thing: walk
1168        // the list, encode each one alone, and the smallest has to be byte for byte what `encode`
1169        // came back with.
1170        for values in [urls(400), keyed(urls(400)), vec![b"same".to_vec(); 50], Vec::new()] {
1171            let borrowed = borrow(&values);
1172            let chosen = encode(&borrowed).unwrap();
1173            let mut smallest: Option<Vec<u8>> = None;
1174            for kind in offered(&borrowed) {
1175                let Some(bytes) = encode_only(kind, &borrowed).unwrap() else {
1176                    continue;
1177                };
1178                if smallest.as_ref().is_none_or(|best| bytes.len() < best.len()) {
1179                    smallest = Some(bytes);
1180                }
1181            }
1182            assert_eq!(smallest.as_deref(), Some(chosen.as_slice()), "{}", values.len());
1183        }
1184    }
1185
1186    fn raw_size(values: &[Vec<u8>]) -> usize {
1187        values.iter().map(Vec::len).sum::<usize>() + values.len() * 4
1188    }
1189
1190    #[test]
1191    fn a_matched_chunk_replays_literals_whether_or_not_they_are_compressed() {
1192        // The literals of a matched chunk are a chunk of their own, and when that chunk is
1193        // compressed the replay decompresses each run straight into the output instead of into a
1194        // buffer it then copies out of. Both columns here are checked value for value by
1195        // round_trip, so what is left is to show that one of them takes the fused path and the
1196        // other takes the one that decodes the literals first, and that the two agree.
1197        let compressed = describe(&round_trip(&keyed(urls(20_000)))).unwrap();
1198        assert!(compressed.starts_with("LZ(") && compressed.contains(", FSST["), "{compressed}");
1199
1200        let buffered = describe(&round_trip(&keyed(urls(300)))).unwrap();
1201        assert!(buffered.starts_with("LZ(") && buffered.contains(", PLAIN("), "{buffered}");
1202    }
1203
1204    #[test]
1205    fn an_empty_chunk_round_trips() {
1206        let bytes = round_trip(&[]);
1207        assert_eq!(kind_of(&bytes), Kind::Plain);
1208    }
1209
1210    #[test]
1211    fn a_constant_column_costs_what_one_value_costs() {
1212        let values = vec![b"https://www.example.com/".to_vec(); 100_000];
1213        let bytes = round_trip(&values);
1214        assert_eq!(kind_of(&bytes), Kind::Constant);
1215        assert_eq!(bytes.len(), 9 + 24);
1216    }
1217
1218    #[test]
1219    fn a_url_column_of_unique_values_is_matched_rather_than_only_compressed() {
1220        // Every value distinct, so a dictionary is the values plus an index and cannot win, and
1221        // every value starts with an identifier of its own, so neighbours share nothing and front
1222        // coding cannot win either. This used to be the case that fell back to FSST, on the
1223        // reasoning that a symbol table was the only thing that could reach repeated vocabulary
1224        // with no structure around it. That reasoning was wrong and #575 is the measurement: the
1225        // vocabulary repeats at a distance, and a match finder reaches distance where a 255 symbol
1226        // table of at most eight bytes each does not.
1227        let values = keyed(urls(20_000));
1228        let bytes = round_trip(&values);
1229        assert_eq!(kind_of(&bytes), Kind::Lz);
1230
1231        // Against the encoding that used to win, on the same values, so the claim is a comparison
1232        // and not just a label.
1233        let borrowed: Vec<&[u8]> = values.iter().map(Vec::as_slice).collect();
1234        let fsst = encode_as(Kind::Fsst, &borrowed, 0, &EXHAUSTIVE).unwrap().unwrap();
1235        assert!(bytes.len() < fsst.len(), "{} against FSST {}", bytes.len(), fsst.len());
1236
1237        // Eleven bytes of every value are the identifier and a separator and nothing compresses
1238        // them, so the ratio here is lower than the one FSST gets on the URLs on their own.
1239        let ratio = raw_size(&values) as f64 / bytes.len() as f64;
1240        assert!(ratio > 4.0, "{ratio:.2}x");
1241    }
1242
1243    #[test]
1244    fn a_sample_of_a_periodic_column_learns_every_phase_of_it() {
1245        // This column is periodic and its period is what a fixed stride would have divided. The
1246        // sample has to see all of it, because a table trained on one phase learns eight byte
1247        // symbols that only line up with that phase and has nothing shorter to fall back on. The
1248        // measured cost of getting this wrong was 3.4 times the compressed size.
1249        let values = urls(20_000);
1250        let borrowed = borrow(&values);
1251        let sample = sample_of(&borrowed);
1252        let mut phases: Vec<&[u8]> = sample
1253            .iter()
1254            .map(|value| {
1255                let query =
1256                    value.iter().position(|byte| *byte == b'?').expect("every value has a query");
1257                &value[..query]
1258            })
1259            .collect();
1260        phases.sort_unstable();
1261        phases.dedup();
1262        // Three hosts and four paths, and the sample has to contain all twelve of the combinations.
1263        assert_eq!(phases.len(), 12);
1264        let whole = SymbolTable::train(&borrowed);
1265        let sampled = SymbolTable::train(&sample);
1266        let mut on_whole = Vec::new();
1267        let mut on_sample = Vec::new();
1268        for value in &borrowed {
1269            whole.compress(value, &mut on_whole);
1270            sampled.compress(value, &mut on_sample);
1271        }
1272        // Training on a twentieth of the column is allowed to cost something. It is not allowed to
1273        // cost a factor.
1274        assert!(
1275            on_sample.len() < on_whole.len() * 5 / 4,
1276            "{} against {}",
1277            on_sample.len(),
1278            on_whole.len()
1279        );
1280    }
1281
1282    #[test]
1283    fn a_repeating_column_becomes_a_dictionary_of_compressed_entries() {
1284        // The DICT_FSST row of the section 6.2 table, which is not an encoding of its own here: it
1285        // is a dictionary whose entries went back through the chooser. What the entries then get
1286        // is whatever wins on them, and since #575 that is the match finder rather than front
1287        // coding with the leftovers FSST compressed. The point of the test is unchanged: nobody
1288        // named the shape and the chooser arrived at it.
1289        let distinct = urls(500);
1290        let values: Vec<Vec<u8>> =
1291            (0..50_000).map(|index| distinct[index * 7919 % distinct.len()].clone()).collect();
1292        let bytes = round_trip(&values);
1293        assert_eq!(kind_of(&bytes), Kind::Dict);
1294        let shape = describe(&bytes).unwrap();
1295        assert!(shape.starts_with("DICT(LZ("), "{shape}");
1296        let ratio = raw_size(&values) as f64 / bytes.len() as f64;
1297        assert!(ratio > 20.0, "{ratio:.2}x, {shape}");
1298    }
1299
1300    #[test]
1301    fn a_column_of_long_runs_costs_almost_nothing() {
1302        // A dictionary makes the codes an integer chunk, and the integer chunk knows what to do
1303        // with runs, so run length encoding of strings falls out of the recursion.
1304        let distinct = urls(50);
1305        let mut values = Vec::new();
1306        for entry in &distinct {
1307            values.extend(std::iter::repeat_n(entry.clone(), 1000));
1308        }
1309        let bytes = round_trip(&values);
1310        let shape = describe(&bytes).unwrap();
1311        assert!(shape.contains("RLE"), "{shape}");
1312        assert!(bytes.len() < 2000, "{} bytes: {shape}", bytes.len());
1313    }
1314
1315    #[test]
1316    fn incompressible_strings_stay_close_to_their_own_size() {
1317        // The case where nothing works. It has to land on PLAIN or on an FSST that is not much
1318        // worse, rather than on a dictionary of every value in the column.
1319        let mut state = 0x2545_f491_4f6c_dd1du64;
1320        let values: Vec<Vec<u8>> = (0..2000)
1321            .map(|_| {
1322                (0..32)
1323                    .map(|_| {
1324                        state ^= state << 13;
1325                        state ^= state >> 7;
1326                        state ^= state << 17;
1327                        state as u8
1328                    })
1329                    .collect()
1330            })
1331            .collect();
1332        let bytes = round_trip(&values);
1333        assert!(bytes.len() < 2000 * 32 + 3000, "{} bytes", bytes.len());
1334    }
1335
1336    #[test]
1337    fn lengths_are_stored_rather_than_offsets() {
1338        // Every value is 24 bytes, so the lengths are a constant chunk and cost 13 bytes for the
1339        // whole column. Offsets would be 100,000 increasing integers.
1340        let values: Vec<Vec<u8>> =
1341            (0..100_000).map(|index| format!("{index:024}").into_bytes()).collect();
1342        let borrowed = borrow(&values);
1343        let bytes = encode_only(Kind::Plain, &borrowed).unwrap().unwrap();
1344        assert_eq!(bytes.len(), 5 + 13 + 100_000 * 24);
1345    }
1346
1347    #[test]
1348    fn empty_strings_are_values_and_not_nulls() {
1349        let values = vec![Vec::new(), b"a".to_vec(), Vec::new(), b"bb".to_vec()];
1350        round_trip(&values);
1351    }
1352
1353    #[test]
1354    fn a_chunk_with_one_value_round_trips() {
1355        round_trip(&[b"only".to_vec()]);
1356    }
1357
1358    #[test]
1359    fn every_candidate_that_applies_decodes_to_the_input() {
1360        let values = urls(3000);
1361        let borrowed = borrow(&values);
1362        let applicable = candidates(&borrowed, 0);
1363        assert!(applicable.len() >= 2, "{applicable:?}");
1364        for kind in applicable {
1365            let bytes = encode_only(kind, &borrowed).unwrap().unwrap();
1366            assert_eq!(decode(&bytes).unwrap(), values, "{}", kind.name());
1367        }
1368    }
1369
1370    #[test]
1371    fn the_chooser_picks_the_smallest_candidate() {
1372        let values = urls(2000);
1373        let borrowed = borrow(&values);
1374        let chosen = encode(&borrowed).unwrap();
1375        for (_, size) in candidate_sizes(&borrowed).unwrap() {
1376            assert!(chosen.len() <= size);
1377        }
1378    }
1379
1380    #[test]
1381    fn a_truncated_chunk_is_an_error_and_not_a_panic() {
1382        let values = urls(40);
1383        let bytes = encode(&borrow(&values)).unwrap();
1384        for len in 0..bytes.len() {
1385            assert!(decode(&bytes[..len]).is_err(), "{len} bytes decoded");
1386        }
1387    }
1388
1389    #[test]
1390    fn trailing_bytes_are_an_error() {
1391        let mut bytes = encode(&borrow(&urls(10))).unwrap();
1392        bytes.push(0);
1393        let error = decode(&bytes).unwrap_err();
1394        assert!(error.message().contains("left over"), "{error}");
1395    }
1396
1397    #[test]
1398    fn an_unknown_tag_is_an_error() {
1399        let error = decode(&[99, 0, 0, 0, 0]).unwrap_err();
1400        assert!(error.message().contains("unknown string encoding tag"), "{error}");
1401    }
1402
1403    #[test]
1404    fn a_dictionary_code_outside_the_dictionary_is_an_error() {
1405        let mut bytes = vec![Kind::Dict.tag()];
1406        put_u32(&mut bytes, 1);
1407        bytes.extend_from_slice(&encode(&[b"one".as_slice()]).unwrap());
1408        bytes.extend_from_slice(&integer::encode(&[9]).unwrap());
1409        let error = decode(&bytes).unwrap_err();
1410        assert!(error.message().contains("not in the dictionary"), "{error}");
1411    }
1412
1413    #[test]
1414    fn a_sorted_column_of_urls_is_front_coded() {
1415        // The M1 finding, in a test. Sorted URLs share a host and most of a path with the URL next
1416        // to them, FSST cannot reach those bytes because it compresses each value on its own, and
1417        // front coding is the shape that reaches them.
1418        let mut values = urls(20_000);
1419        values.sort();
1420        let bytes = round_trip(&values);
1421        assert_eq!(kind_of(&bytes), Kind::Front);
1422        let shape = describe(&bytes).unwrap();
1423        let mut plain = Vec::new();
1424        let borrowed = borrow(&values);
1425        for (kind, size) in candidate_sizes(&borrowed).unwrap() {
1426            if kind == Kind::Fsst {
1427                plain.push(size);
1428            }
1429        }
1430        let fsst = plain[0];
1431        assert!(bytes.len() * 2 < fsst, "{} against FSST {fsst}: {shape}", bytes.len());
1432    }
1433
1434    #[test]
1435    fn a_column_with_nothing_to_share_is_not_offered_front_coding() {
1436        // The candidate costs an encode of the whole column, so a column whose neighbours have
1437        // nothing in common must not be paying for it.
1438        let mut state = 0x9e37_79b9_7f4a_7c15u64;
1439        let values: Vec<Vec<u8>> = (0..2000)
1440            .map(|_| {
1441                (0..24)
1442                    .map(|_| {
1443                        state ^= state << 13;
1444                        state ^= state >> 7;
1445                        state ^= state << 17;
1446                        (state % 251) as u8
1447                    })
1448                    .collect()
1449            })
1450            .collect();
1451        let borrowed = borrow(&values);
1452        assert!(!candidates(&borrowed, 0).contains(&Kind::Front));
1453    }
1454
1455    #[test]
1456    fn a_prefix_longer_than_the_value_before_it_is_an_error() {
1457        let mut bytes = vec![Kind::Front.tag()];
1458        put_u32(&mut bytes, 2);
1459        bytes.extend_from_slice(&integer::encode(&[0, 9]).unwrap());
1460        bytes.extend_from_slice(&encode(&[b"one".as_slice(), b"two".as_slice()]).unwrap());
1461        let error = decode(&bytes).unwrap_err();
1462        assert!(error.message().contains("shares 9 bytes"), "{error}");
1463    }
1464
1465    #[test]
1466    fn a_negative_prefix_is_an_error() {
1467        let mut bytes = vec![Kind::Front.tag()];
1468        put_u32(&mut bytes, 1);
1469        bytes.extend_from_slice(&integer::encode(&[-1]).unwrap());
1470        bytes.extend_from_slice(&encode(&[b"one".as_slice()]).unwrap());
1471        let error = decode(&bytes).unwrap_err();
1472        assert!(error.message().contains("negative shared prefix"), "{error}");
1473    }
1474
1475    #[test]
1476    fn a_negative_length_is_an_error() {
1477        let mut bytes = vec![Kind::Plain.tag()];
1478        put_u32(&mut bytes, 1);
1479        bytes.extend_from_slice(&integer::encode(&[-1]).unwrap());
1480        let error = decode(&bytes).unwrap_err();
1481        assert!(error.message().contains("negative string length"), "{error}");
1482    }
1483
1484    #[test]
1485    fn the_sample_is_spread_across_the_chunk_and_not_taken_from_the_front() {
1486        // A sorted column whose first 64 KB says nothing about the rest of it. If the sample were
1487        // the front, the table would learn `aaaa` and escape every `zzzz`.
1488        let mut values: Vec<Vec<u8>> = Vec::new();
1489        for index in 0..20_000 {
1490            let head = if index < 10_000 { "aaaaaaaaaaaaaaaa" } else { "zzzzzzzzzzzzzzzz" };
1491            values.push(format!("{head}/{index:08}").into_bytes());
1492        }
1493        let borrowed = borrow(&values);
1494        let sample = sample_of(&borrowed);
1495        let first_half = sample.iter().filter(|value| value.starts_with(b"aaaa")).count();
1496        let second_half = sample.len() - first_half;
1497        assert!(first_half > 0 && second_half > 0, "{first_half} and {second_half}");
1498        let bytes = round_trip(&values);
1499        let ratio = raw_size(&values) as f64 / bytes.len() as f64;
1500        assert!(ratio > 4.0, "{ratio:.2}x");
1501    }
1502}