Skip to main content

rudb_encoding/
multi.rs

1//! Columns encoded together instead of one at a time.
2//!
3//! `spec/06-compression.md` section 6.4 calls this multi-column compression and says it is the
4//! mechanism that has no equivalent in DuckDB. The idea is that columns are not independent, so
5//! encoding them independently throws away real redundancy. On ClickBench `hits` the obvious case
6//! is `URL` and `Referer`, which are both URLs drawn from the same universe, and the less obvious
7//! one is that `URL`, `Referer`, `Title` and the referer derived columns are all the same alphabet
8//! and could share one symbol table.
9//!
10//! ## The three strategies
11//!
12//! `INDEPENDENT` is each column encoded on its own by [`crate::string`]. It is the baseline the
13//! other two have to beat, and it is what a group falls back to when they do not.
14//!
15//! `SHARED_TABLE` is one FSST symbol table trained on a sample of all the columns, with every
16//! column front coded and the leftovers compressed against it. Storing one 255 symbol table rather
17//! than six saves almost nothing by itself. The effect that matters is that a table trained on the
18//! union has more evidence per symbol, so it compresses each column better than a table trained on
19//! that column alone would, and the columns that gain most are the small ones that never had enough
20//! bytes to train on.
21//!
22//! The front coding is there because `INDEPENDENT` is the baseline and `INDEPENDENT` can front
23//! code, so a shared table that could not would be losing to a baseline it was never measured
24//! against. On the generated URLs in the tests here that gap was a factor of two. Front coding a
25//! column whose neighbours share nothing costs a run of zeros, which the integer chunk stores in a
26//! few bytes, so nothing has to decide whether to do it.
27//!
28//! `SHARED_DICT` is one dictionary holding the union of the values, with every column becoming an
29//! array of codes into it. A value that appears in three columns is stored once rather than three
30//! times. The dictionary is itself a string column, so it goes back through the string chooser and
31//! comes out FSST compressed, which means a shared dictionary is also a shared symbol table.
32//!
33//! ## What decides
34//!
35//! Here, measuring all three and keeping the smallest, for the same reason [`crate::string`] does
36//! it that way: M1 is measuring what the format can do rather than how fast a writer can decide.
37//! A write path cannot afford this. Section 6.4 says detection is by sampling pairs over a global
38//! sample at table level, recorded in the catalog as hints, with each row group checking only the
39//! hinted pairs. [`dictionary_groups`] is that pruning step, and it runs on sketches rather than on
40//! data, so the 5,460 pairs of a 105 column table cost a Jaccard estimate each rather than a pass
41//! over the column.
42//!
43//! ## What is not here
44//!
45//! Correlation encodings, which are the third form in section 6.4: column B stored as a function of
46//! column A, either as a per dictionary entry lookup for a functional dependency or as B minus f(A)
47//! for a numeric one. [`crate::sketch::dependence`] is the detection half of that and the encoding
48//! half is its own piece of work.
49//!
50//! Global dictionaries, which are section 6.5 and are a dictionary across the whole table rather
51//! than across a group of columns in one chunk. The two compose, and the reason they are separate
52//! is that a global dictionary needs an incremental builder that can spill, which is open question
53//! five and is the thing most likely to make the idea impractical.
54
55use rudb_common::{Error, Result};
56
57use crate::fsst::SymbolTable;
58use crate::integer;
59use crate::reader::Reader;
60use crate::sketch::Sketch;
61use crate::string;
62
63/// The least a column contributes to a shared sample, however small it is next to the rest of the
64/// group. Enough to see an alphabet, not enough to matter against the 64 KB the group gets.
65const FLOOR_SAMPLE_BYTES: usize = 2 * 1024;
66
67/// How a column group is encoded. The discriminant is the tag byte and is part of the format.
68#[derive(Debug, Clone, Copy, PartialEq, Eq)]
69pub enum Strategy {
70    /// Each column encoded on its own.
71    Independent = 0,
72    /// One symbol table over all of them.
73    SharedTable = 1,
74    /// One dictionary over the union of their values.
75    SharedDict = 2,
76}
77
78impl Strategy {
79    fn tag(self) -> u8 {
80        self as u8
81    }
82
83    fn from_tag(tag: u8) -> Result<Self> {
84        match tag {
85            0 => Ok(Self::Independent),
86            1 => Ok(Self::SharedTable),
87            2 => Ok(Self::SharedDict),
88            other => Err(Error::internal(format!("unknown column group tag {other}"))),
89        }
90    }
91
92    /// The name that goes in a report.
93    #[must_use]
94    pub fn name(self) -> &'static str {
95        match self {
96            Self::Independent => "INDEPENDENT",
97            Self::SharedTable => "SHARED_TABLE",
98            Self::SharedDict => "SHARED_DICT",
99        }
100    }
101}
102
103/// Encodes a group of string columns together, choosing whatever comes out smallest.
104///
105/// The columns do not have to be the same length. Nothing here pairs values up by row, so a group
106/// is a set of columns that share an alphabet or a value universe rather than a set of columns from
107/// the same table.
108///
109/// # Errors
110///
111/// If the group holds more than `u32::MAX` columns or a column holds more than `u32::MAX` values,
112/// or if an encoding produces something its own decoder would not accept.
113pub fn encode_group(columns: &[&[&[u8]]]) -> Result<Vec<u8>> {
114    let mut best: Option<Vec<u8>> = None;
115    for strategy in [Strategy::Independent, Strategy::SharedTable, Strategy::SharedDict] {
116        let Some(bytes) = encode_as(strategy, columns)? else {
117            continue;
118        };
119        if best.as_ref().is_none_or(|current| bytes.len() < current.len()) {
120            best = Some(bytes);
121        }
122    }
123    best.ok_or_else(|| Error::internal("no strategy applied to the column group"))
124}
125
126/// Decodes a group written by [`encode_group`].
127///
128/// # Errors
129///
130/// If the bytes are truncated, carry an unknown tag, or describe a group whose parts disagree.
131pub fn decode_group(bytes: &[u8]) -> Result<Vec<Vec<Vec<u8>>>> {
132    let mut reader = Reader::new(bytes);
133    let columns = decode_at(&mut reader)?;
134    if reader.remaining() != 0 {
135        return Err(Error::internal(format!(
136            "{} bytes left over after decoding a column group",
137            reader.remaining()
138        )));
139    }
140    Ok(columns)
141}
142
143/// The size of every strategy that applies, which is the measurement section 6.4 is asking for.
144///
145/// # Errors
146///
147/// As [`encode_group`].
148pub fn strategy_sizes(columns: &[&[&[u8]]]) -> Result<Vec<(Strategy, usize)>> {
149    let mut sizes = Vec::new();
150    for strategy in [Strategy::Independent, Strategy::SharedTable, Strategy::SharedDict] {
151        if let Some(bytes) = encode_as(strategy, columns)? {
152            sizes.push((strategy, bytes.len()));
153        }
154    }
155    Ok(sizes)
156}
157
158/// The shape a group was encoded as, as a line of text.
159///
160/// # Errors
161///
162/// As [`decode_group`].
163pub fn describe(bytes: &[u8]) -> Result<String> {
164    let mut reader = Reader::new(bytes);
165    describe_at(&mut reader)
166}
167
168/// Which columns should be considered for a shared dictionary, from one sketch per column.
169///
170/// Two columns are put in the same group when their estimated Jaccard similarity is at least
171/// `threshold`, and grouping is transitive: if A overlaps B and B overlaps C then all three end up
172/// together, even if A and C do not overlap each other. That is deliberate. A dictionary is a set
173/// union, so a chain of overlapping columns still stores fewer values once than separately, and
174/// insisting that every pair in a group overlaps would turn this into a clique problem for no
175/// benefit.
176///
177/// Every column appears in exactly one group, and a column that overlaps nothing is a group of one.
178/// The groups come back in the order of their lowest column index, and the columns inside a group
179/// in index order, so the result does not depend on the order the pairs were tested in.
180///
181/// This is a pruning step and not a decision. Two columns can overlap heavily and still be better
182/// off apart, which is why what comes out of here goes to [`strategy_sizes`] rather than straight
183/// into a writer.
184///
185/// # Errors
186///
187/// If the sketches were not all built at the same k, since then no estimate over a pair means
188/// anything.
189pub fn dictionary_groups(sketches: &[Sketch], threshold: f64) -> Result<Vec<Vec<usize>>> {
190    let mut parent: Vec<usize> = (0..sketches.len()).collect();
191    for left in 0..sketches.len() {
192        for right in (left + 1)..sketches.len() {
193            if sketches[left].jaccard(&sketches[right])? >= threshold {
194                let (a, b) = (find(&mut parent, left), find(&mut parent, right));
195                if a != b {
196                    // The lower index wins, so the group's root is its first column and the output
197                    // order is a function of the columns rather than of the loop.
198                    let (low, high) = if a < b { (a, b) } else { (b, a) };
199                    parent[high] = low;
200                }
201            }
202        }
203    }
204    let mut groups: Vec<Vec<usize>> = Vec::new();
205    let mut roots: Vec<usize> = Vec::new();
206    for column in 0..sketches.len() {
207        let root = find(&mut parent, column);
208        match roots.iter().position(|seen| *seen == root) {
209            Some(at) => groups[at].push(column),
210            None => {
211                roots.push(root);
212                groups.push(vec![column]);
213            }
214        }
215    }
216    Ok(groups)
217}
218
219fn find(parent: &mut [usize], mut node: usize) -> usize {
220    while parent[node] != node {
221        parent[node] = parent[parent[node]];
222        node = parent[node];
223    }
224    node
225}
226
227fn encode_as(strategy: Strategy, columns: &[&[&[u8]]]) -> Result<Option<Vec<u8>>> {
228    let mut out = vec![strategy.tag()];
229    put_u32(&mut out, u32::try_from(columns.len()).map_err(|_| too_many(columns.len()))?);
230    match strategy {
231        Strategy::Independent => {
232            for column in columns {
233                out.extend_from_slice(&string::encode(column)?);
234            }
235        }
236        Strategy::SharedTable => {
237            if columns.len() < 2 {
238                return Ok(None);
239            }
240            // Front coded first, and the table trained on what is left. A shared table has to beat
241            // encoding the columns apart, and encoding a column apart can front code it, so a
242            // shared table that cannot is comparing itself against a better baseline than the one
243            // it was written for. On the generated URLs in the tests here that gap was a factor of
244            // two. Front coding a column with nothing to share costs a run of zeros, which the
245            // integer chunk stores in a few bytes, so this is not a decision anything has to make.
246            let coded: Vec<(Vec<i64>, Vec<&[u8]>)> =
247                columns.iter().map(|column| string::front_code(column)).collect();
248            let suffixes: Vec<&[&[u8]]> =
249                coded.iter().map(|(_, suffixes)| suffixes.as_slice()).collect();
250            let table = SymbolTable::train(&shared_sample(&suffixes));
251            if table.is_empty() {
252                return Ok(None);
253            }
254            table.serialize(&mut out);
255            for (prefixes, suffixes) in &coded {
256                put_u32(
257                    &mut out,
258                    u32::try_from(prefixes.len()).map_err(|_| too_many(prefixes.len()))?,
259                );
260                out.extend_from_slice(&integer::encode(prefixes)?);
261                let mut compressed = Vec::new();
262                let mut lengths = Vec::with_capacity(suffixes.len());
263                for value in suffixes {
264                    let before = compressed.len();
265                    table.compress(value, &mut compressed);
266                    lengths.push((compressed.len() - before) as i64);
267                }
268                out.extend_from_slice(&integer::encode(&lengths)?);
269                out.extend_from_slice(&compressed);
270            }
271        }
272        Strategy::SharedDict => {
273            if columns.len() < 2 {
274                return Ok(None);
275            }
276            let dictionary = union_values(columns);
277            let total: usize = columns.iter().map(|column| column.len()).sum();
278            // A dictionary holding as many values as the columns do stores everything once and adds
279            // an index on top, so it cannot win and is not worth the encode.
280            if dictionary.is_empty() || dictionary.len() >= total {
281                return Ok(None);
282            }
283            let entries: Vec<&[u8]> = dictionary.iter().map(Vec::as_slice).collect();
284            out.extend_from_slice(&string::encode(&entries)?);
285            for column in columns {
286                let codes = codes_over(column, &dictionary);
287                out.extend_from_slice(&integer::encode(&codes)?);
288            }
289        }
290    }
291    Ok(Some(out))
292}
293
294fn decode_at(reader: &mut Reader<'_>) -> Result<Vec<Vec<Vec<u8>>>> {
295    let strategy = Strategy::from_tag(reader.u8()?)?;
296    let count = reader.u32()? as usize;
297    let mut columns = Vec::with_capacity(count.min(1024));
298    match strategy {
299        Strategy::Independent => {
300            for _ in 0..count {
301                let (values, used) = string::decode_prefix(reader.rest())?;
302                reader.skip(used)?;
303                columns.push(values);
304            }
305        }
306        Strategy::SharedTable => {
307            let (table, used) = SymbolTable::deserialize(reader.rest())?;
308            reader.skip(used)?;
309            for _ in 0..count {
310                let rows = reader.u32()? as usize;
311                let (prefixes, used) = integer::decode_prefix(reader.rest())?;
312                reader.skip(used)?;
313                let (lengths, used) = integer::decode_prefix(reader.rest())?;
314                reader.skip(used)?;
315                if lengths.len() != rows || prefixes.len() != rows {
316                    return Err(Error::internal(format!(
317                        "a column says it holds {rows} values and has {} prefixes and {} lengths",
318                        prefixes.len(),
319                        lengths.len()
320                    )));
321                }
322                let mut suffixes = Vec::with_capacity(rows);
323                for length in lengths {
324                    let length = usize::try_from(length)
325                        .map_err(|_| Error::internal("a negative compressed length"))?;
326                    let compressed = reader.bytes(length)?;
327                    let mut value = Vec::new();
328                    table.decompress(compressed, &mut value)?;
329                    suffixes.push(value);
330                }
331                columns.push(string::front_decode(&prefixes, suffixes)?);
332            }
333        }
334        Strategy::SharedDict => {
335            let (dictionary, used) = string::decode_prefix(reader.rest())?;
336            reader.skip(used)?;
337            for _ in 0..count {
338                let (codes, used) = integer::decode_prefix(reader.rest())?;
339                reader.skip(used)?;
340                let mut values = Vec::with_capacity(codes.len());
341                for code in codes {
342                    let entry = usize::try_from(code)
343                        .ok()
344                        .and_then(|index| dictionary.get(index))
345                        .ok_or_else(|| {
346                            Error::internal(format!("code {code} is not in the shared dictionary"))
347                        })?;
348                    values.push(entry.clone());
349                }
350                columns.push(values);
351            }
352        }
353    }
354    Ok(columns)
355}
356
357fn describe_at(reader: &mut Reader<'_>) -> Result<String> {
358    let strategy = Strategy::from_tag(reader.u8()?)?;
359    let count = reader.u32()? as usize;
360    let mut parts = Vec::with_capacity(count.min(1024));
361    let head = match strategy {
362        Strategy::Independent => {
363            for _ in 0..count {
364                let (text, used) = string::describe_prefix(reader.rest())?;
365                reader.skip(used)?;
366                parts.push(text);
367            }
368            "INDEPENDENT".to_string()
369        }
370        Strategy::SharedTable => {
371            let (table, used) = SymbolTable::deserialize(reader.rest())?;
372            reader.skip(used)?;
373            for _ in 0..count {
374                let rows = reader.u32()? as usize;
375                let (prefixes, used) = integer::describe_prefix(reader.rest())?;
376                reader.skip(used)?;
377                // The same chunk twice, once for its shape and once for the lengths themselves,
378                // which is how the describe knows how far past the payload to step.
379                let (text, _) = integer::describe_prefix(reader.rest())?;
380                let (lengths, used) = integer::decode_prefix(reader.rest())?;
381                reader.skip(used)?;
382                if lengths.len() != rows {
383                    return Err(Error::internal("a column group disagrees with itself"));
384                }
385                let bytes: i64 = lengths.iter().sum();
386                reader.skip(usize::try_from(bytes).map_err(|_| {
387                    Error::internal("a column group has a negative compressed size")
388                })?)?;
389                parts.push(format!("FRONT({prefixes}, {text})"));
390            }
391            format!("SHARED_TABLE[{}]", table.len())
392        }
393        Strategy::SharedDict => {
394            let (text, used) = string::describe_prefix(reader.rest())?;
395            reader.skip(used)?;
396            for _ in 0..count {
397                let (codes, used) = integer::describe_prefix(reader.rest())?;
398                reader.skip(used)?;
399                parts.push(codes);
400            }
401            format!("SHARED_DICT({text})")
402        }
403    };
404    Ok(format!("{head}({})", parts.join(", ")))
405}
406
407/// A sample of all the columns together, with the byte budget shared out in proportion to how big
408/// the columns are and a floor so that a small column still gets looked at.
409///
410/// The total budget is the same one a single column gets, because a shared table that trained on
411/// six times the sample would win partly on the sample size and the report would credit sharing
412/// with something sharing did not do.
413///
414/// Splitting that budget evenly is the obvious thing and it is wrong. A group of one column of
415/// 20,000 values and one of 40 is dominated by the first, and giving the first half the budget
416/// costs it more than the second gains: measured on exactly that pair, an even split made the
417/// shared table 220,755 bytes against 218,099 for encoding the two columns independently, so
418/// sharing lost. In proportion, the dominant column keeps nearly all of the budget and its table is
419/// nearly the table it would have had alone, while the floor is what buys the small column a table
420/// trained on far more evidence than its own 40 values could provide.
421fn shared_sample<'a>(columns: &[&[&'a [u8]]]) -> Vec<&'a [u8]> {
422    let sizes: Vec<usize> =
423        columns.iter().map(|column| column.iter().map(|value| value.len()).sum()).collect();
424    let total: usize = sizes.iter().sum();
425    let floor = FLOOR_SAMPLE_BYTES;
426    let mut sample = Vec::new();
427    for (column, bytes) in columns.iter().zip(&sizes) {
428        let share = if total == 0 {
429            floor
430        } else {
431            (string::SAMPLE_BYTES as u128 * *bytes as u128 / total as u128) as usize
432        };
433        sample.extend(string::sample_bytes_of(column, share.max(floor)));
434    }
435    sample
436}
437
438/// The distinct values across all the columns, sorted, which is the shared dictionary.
439fn union_values(columns: &[&[&[u8]]]) -> Vec<Vec<u8>> {
440    let mut values: Vec<Vec<u8>> =
441        columns.iter().flat_map(|column| column.iter().map(|value| value.to_vec())).collect();
442    values.sort_unstable();
443    values.dedup();
444    values
445}
446
447fn codes_over(values: &[&[u8]], dictionary: &[Vec<u8>]) -> Vec<i64> {
448    values
449        .iter()
450        .map(|value| {
451            dictionary
452                .binary_search_by(|entry| entry.as_slice().cmp(value))
453                .expect("the dictionary is the union of the columns in this group")
454                as i64
455        })
456        .collect()
457}
458
459fn too_many(count: usize) -> Error {
460    Error::internal(format!("a column group of {count} is larger than the format allows"))
461}
462
463fn put_u32(out: &mut Vec<u8>, value: u32) {
464    out.extend_from_slice(&value.to_le_bytes());
465}
466
467#[cfg(test)]
468mod tests {
469    use super::*;
470
471    /// URLs from one host, so two columns built from different hosts share an alphabet and no
472    /// values at all.
473    fn urls(host: &str, count: usize, from: usize) -> Vec<Vec<u8>> {
474        let paths = ["/index.html", "/catalog/item", "/search", "/user/profile/settings"];
475        (from..from + count)
476            .map(|index| {
477                let path = paths[(index / 3) % paths.len()];
478                format!("http://{host}{path}?session={}&ref=google", index * 7).into_bytes()
479            })
480            .collect()
481    }
482
483    fn borrow(values: &[Vec<u8>]) -> Vec<&[u8]> {
484        values.iter().map(Vec::as_slice).collect()
485    }
486
487    fn group<'a>(columns: &'a [Vec<&'a [u8]>]) -> Vec<&'a [&'a [u8]]> {
488        columns.iter().map(Vec::as_slice).collect()
489    }
490
491    fn round_trip(columns: &[&[&[u8]]]) -> Vec<u8> {
492        let bytes = encode_group(columns).unwrap();
493        let back = decode_group(&bytes).unwrap();
494        assert_eq!(back.len(), columns.len());
495        for (decoded, original) in back.iter().zip(columns) {
496            assert_eq!(decoded.len(), original.len(), "{}", describe(&bytes).unwrap());
497            for (left, right) in decoded.iter().zip(*original) {
498                assert_eq!(left.as_slice(), *right, "{}", describe(&bytes).unwrap());
499            }
500        }
501        bytes
502    }
503
504    fn strategy_of(bytes: &[u8]) -> Strategy {
505        Strategy::from_tag(bytes[0]).unwrap()
506    }
507
508    fn size_of(sizes: &[(Strategy, usize)], strategy: Strategy) -> usize {
509        sizes
510            .iter()
511            .find(|(kind, _)| *kind == strategy)
512            .map(|(_, size)| *size)
513            .unwrap_or_else(|| panic!("{} did not apply", strategy.name()))
514    }
515
516    #[test]
517    fn two_columns_from_the_same_universe_share_a_dictionary() {
518        // The `URL` and `Referer` case. Both columns are URLs and most of the values in one are in
519        // the other, so the union is stored once instead of the intersection being stored twice.
520        let left = urls("www.example.com", 20_000, 0);
521        let right = urls("www.example.com", 20_000, 5_000);
522        let columns = [borrow(&left), borrow(&right)];
523        let group = group(&columns);
524        let bytes = round_trip(&group);
525        assert_eq!(strategy_of(&bytes), Strategy::SharedDict);
526        let sizes = strategy_sizes(&group).unwrap();
527        let independent = size_of(&sizes, Strategy::Independent);
528        // Three quarters of the values are in both columns, and the saving is a third rather than
529        // a half because the codes still cost something once the values stop being repeated.
530        assert!(
531            bytes.len() * 4 < independent * 3,
532            "{} against {independent} independent",
533            bytes.len()
534        );
535    }
536
537    #[test]
538    fn two_columns_of_the_same_alphabet_share_a_symbol_table() {
539        // No value appears in both columns, so a dictionary has nothing to share. The alphabet is
540        // the same, which is what a symbol table can still share.
541        //
542        // A thousand values a column, and that number is doing work. Sharing wins here because a
543        // 255 symbol table is a fixed cost and two columns of this size pay it twice. It stops
544        // winning once each column is big enough to train a table of its own that fits it better
545        // than a joint one does, and on this data the crossover is between two and four thousand
546        // values a column: 21,319 shared against 22,223 apart at two thousand, and 41,002 against
547        // 39,750 at four. The gain from sharing a table is a small column's gain, which is what
548        // the next test is about.
549        let left = urls("www.example.com", 1_000, 0);
550        let right = urls("news.other.example.org", 1_000, 500_000);
551        let columns = [borrow(&left), borrow(&right)];
552        let group = group(&columns);
553        let bytes = round_trip(&group);
554        let sizes = strategy_sizes(&group).unwrap();
555        let shared = size_of(&sizes, Strategy::SharedTable);
556        let independent = size_of(&sizes, Strategy::Independent);
557        assert!(shared < independent, "{shared} against {independent} independent");
558        assert_eq!(strategy_of(&bytes), Strategy::SharedTable);
559    }
560
561    #[test]
562    fn a_small_column_gains_most_from_a_shared_table() {
563        // The reason sharing a table is worth more than the bytes of the table. A column of 40
564        // values has nothing to train on, and a table trained on the group is a table that has.
565        let big = urls("www.example.com", 20_000, 0);
566        let small = urls("www.example.com", 40, 900_000);
567        let columns = [borrow(&big), borrow(&small)];
568        let sizes = strategy_sizes(&group(&columns)).unwrap();
569        let shared = size_of(&sizes, Strategy::SharedTable);
570        let independent = size_of(&sizes, Strategy::Independent);
571        assert!(shared < independent, "group {shared} against {independent} apart");
572    }
573
574    #[test]
575    fn unrelated_columns_are_left_alone() {
576        let urls = urls("www.example.com", 4_000, 0);
577        let numbers: Vec<Vec<u8>> = (0..4_000)
578            .map(|index| format!("{:016x}", index * 2_654_435_761u64).into_bytes())
579            .collect();
580        let columns = [borrow(&urls), borrow(&numbers)];
581        let group = group(&columns);
582        let bytes = round_trip(&group);
583        assert_eq!(strategy_of(&bytes), Strategy::Independent);
584    }
585
586    #[test]
587    fn a_group_of_one_is_the_column_on_its_own() {
588        let column = urls("www.example.com", 2_000, 0);
589        let columns = [borrow(&column)];
590        let bytes = round_trip(&group(&columns));
591        assert_eq!(strategy_of(&bytes), Strategy::Independent);
592        assert_eq!(bytes.len(), 5 + string::encode(&borrow(&column)).unwrap().len());
593    }
594
595    #[test]
596    fn an_empty_group_round_trips() {
597        let bytes = round_trip(&[]);
598        assert_eq!(decode_group(&bytes).unwrap().len(), 0);
599    }
600
601    #[test]
602    fn columns_do_not_have_to_be_the_same_length() {
603        let left = urls("www.example.com", 3_000, 0);
604        let right = urls("www.example.com", 700, 1_000);
605        let columns = [borrow(&left), borrow(&right)];
606        round_trip(&group(&columns));
607    }
608
609    #[test]
610    fn an_empty_column_in_a_group_round_trips() {
611        let left = urls("www.example.com", 1_000, 0);
612        let empty: Vec<Vec<u8>> = Vec::new();
613        let columns = [borrow(&left), borrow(&empty)];
614        round_trip(&group(&columns));
615    }
616
617    #[test]
618    fn every_strategy_that_applies_decodes_to_the_input() {
619        let left = urls("www.example.com", 3_000, 0);
620        let right = urls("www.example.com", 3_000, 1_000);
621        let columns = [borrow(&left), borrow(&right)];
622        let group = group(&columns);
623        for strategy in [Strategy::Independent, Strategy::SharedTable, Strategy::SharedDict] {
624            let bytes = encode_as(strategy, &group).unwrap().unwrap();
625            let back = decode_group(&bytes).unwrap();
626            assert_eq!(back[0].len(), left.len(), "{}", strategy.name());
627            assert_eq!(back[1][7], right[7], "{}", strategy.name());
628        }
629    }
630
631    #[test]
632    fn the_chooser_picks_the_smallest_strategy() {
633        let left = urls("www.example.com", 2_000, 0);
634        let right = urls("www.example.com", 2_000, 1_000);
635        let columns = [borrow(&left), borrow(&right)];
636        let group = group(&columns);
637        let chosen = encode_group(&group).unwrap();
638        for (_, size) in strategy_sizes(&group).unwrap() {
639            assert!(chosen.len() <= size);
640        }
641    }
642
643    #[test]
644    fn describe_says_what_every_column_came_out_as() {
645        let left = urls("www.example.com", 2_000, 0);
646        let right = urls("www.example.com", 2_000, 1_000);
647        let columns = [borrow(&left), borrow(&right)];
648        let group = group(&columns);
649        for strategy in [Strategy::Independent, Strategy::SharedTable, Strategy::SharedDict] {
650            let bytes = encode_as(strategy, &group).unwrap().unwrap();
651            let shape = describe(&bytes).unwrap();
652            assert!(shape.starts_with(strategy.name()), "{shape}");
653            assert!(shape.contains(", "), "{shape}");
654        }
655    }
656
657    #[test]
658    fn a_truncated_group_is_an_error_and_not_a_panic() {
659        let left = urls("www.example.com", 40, 0);
660        let right = urls("www.example.com", 40, 20);
661        let columns = [borrow(&left), borrow(&right)];
662        let group = group(&columns);
663        for strategy in [Strategy::Independent, Strategy::SharedTable, Strategy::SharedDict] {
664            let bytes = encode_as(strategy, &group).unwrap().unwrap();
665            for len in 0..bytes.len() {
666                assert!(
667                    decode_group(&bytes[..len]).is_err(),
668                    "{} decoded at {len} bytes",
669                    strategy.name()
670                );
671            }
672        }
673    }
674
675    #[test]
676    fn trailing_bytes_are_an_error() {
677        let column = urls("www.example.com", 10, 0);
678        let columns = [borrow(&column)];
679        let mut bytes = encode_group(&group(&columns)).unwrap();
680        bytes.push(0);
681        let error = decode_group(&bytes).unwrap_err();
682        assert!(error.message().contains("left over"), "{error}");
683    }
684
685    #[test]
686    fn an_unknown_tag_is_an_error() {
687        let error = decode_group(&[9, 0, 0, 0, 0]).unwrap_err();
688        assert!(error.message().contains("unknown column group tag"), "{error}");
689    }
690
691    #[test]
692    fn a_code_outside_the_shared_dictionary_is_an_error() {
693        let mut bytes = vec![Strategy::SharedDict.tag()];
694        put_u32(&mut bytes, 1);
695        bytes.extend_from_slice(&string::encode(&[b"one".as_slice()]).unwrap());
696        bytes.extend_from_slice(&integer::encode(&[4]).unwrap());
697        let error = decode_group(&bytes).unwrap_err();
698        assert!(error.message().contains("not in the shared dictionary"), "{error}");
699    }
700
701    #[test]
702    fn overlapping_columns_are_grouped_and_the_rest_are_not() {
703        let first = urls("www.example.com", 20_000, 0);
704        let second = urls("news.other.example.org", 20_000, 0);
705        let third = urls("www.example.com", 20_000, 4_000);
706        let fourth = urls("news.other.example.org", 20_000, 4_000);
707        let sketches: Vec<Sketch> = [&first, &second, &third, &fourth]
708            .iter()
709            .map(|column| Sketch::of(&borrow(column)))
710            .collect();
711        let groups = dictionary_groups(&sketches, 0.5).unwrap();
712        assert_eq!(groups, vec![vec![0, 2], vec![1, 3]]);
713    }
714
715    #[test]
716    fn a_column_that_overlaps_nothing_is_a_group_of_one() {
717        let sketches: Vec<Sketch> = (0..4)
718            .map(|index| Sketch::of(&borrow(&urls("www.example.com", 5_000, index * 100_000))))
719            .collect();
720        let groups = dictionary_groups(&sketches, 0.5).unwrap();
721        assert_eq!(groups, vec![vec![0], vec![1], vec![2], vec![3]]);
722    }
723
724    #[test]
725    fn grouping_is_transitive_and_does_not_need_every_pair_to_overlap() {
726        // A overlaps B by half, B overlaps C by half, A and C not at all. They still go together,
727        // because a dictionary is a union and a chain of overlaps still stores fewer values once
728        // than three sets store separately.
729        let a = urls("www.example.com", 20_000, 0);
730        let b = urls("www.example.com", 20_000, 10_000);
731        let c = urls("www.example.com", 20_000, 20_000);
732        let sketches: Vec<Sketch> =
733            [&a, &b, &c].iter().map(|column| Sketch::of(&borrow(column))).collect();
734        assert!(sketches[0].jaccard(&sketches[2]).unwrap() < 0.01);
735        let groups = dictionary_groups(&sketches, 0.3).unwrap();
736        assert_eq!(groups, vec![vec![0, 1, 2]]);
737    }
738
739    #[test]
740    fn grouping_needs_sketches_of_the_same_size() {
741        let sketches = [Sketch::new(16).unwrap(), Sketch::new(32).unwrap()];
742        assert!(dictionary_groups(&sketches, 0.5).is_err());
743    }
744}