1use rudb_common::{Error, Result};
56
57use crate::fsst::SymbolTable;
58use crate::integer;
59use crate::reader::Reader;
60use crate::sketch::Sketch;
61use crate::string;
62
63const FLOOR_SAMPLE_BYTES: usize = 2 * 1024;
66
67#[derive(Debug, Clone, Copy, PartialEq, Eq)]
69pub enum Strategy {
70 Independent = 0,
72 SharedTable = 1,
74 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 #[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
103pub 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
126pub 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
143pub 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
158pub fn describe(bytes: &[u8]) -> Result<String> {
164 let mut reader = Reader::new(bytes);
165 describe_at(&mut reader)
166}
167
168pub 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 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 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 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 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
407fn 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
438fn 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 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 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 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 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 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 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}