1use std::path::Path;
16use std::time::{Duration, Instant};
17
18use rudb_common::{LogicalType, Result, Value};
19use rudb_graph::{Form, KeyMap, Keys, wire};
20use rudb_vector::Chunk;
21
22use crate::section::{self, Attachment};
23use crate::{Catalog, Reader, invalid, type_tag};
24
25#[derive(Debug)]
32pub struct KeyColumn<'a> {
33 reader: &'a Reader,
34 column: usize,
35}
36
37impl<'a> KeyColumn<'a> {
38 pub fn new(reader: &'a Reader, column: usize) -> Result<Self> {
47 let fields = reader.table().fields();
48 let Some(field) = fields.get(column) else {
49 return Err(invalid(&format!(
50 "column {column} is past the {} of table {}",
51 fields.len(),
52 reader.table().name()
53 )));
54 };
55 if !mappable(&field.ty) {
56 return Err(invalid(&format!(
57 "a key map over {} needs an integer key form, and {} has none",
58 field.name, field.ty
59 )));
60 }
61 Ok(Self { reader, column })
62 }
63}
64
65impl Keys for KeyColumn<'_> {
66 fn scan(&self, each: &mut dyn FnMut(Option<i128>) -> Result<()>) -> Result<()> {
67 for part in 0..self.reader.parts() {
68 let chunk = self.reader.read(part, &[self.column])?;
69 let values = chunk.column(0)?;
70 for row in 0..chunk.len() {
71 each(key_at(&chunk, values, row)?)?;
72 }
73 }
74 Ok(())
75 }
76}
77
78fn mappable(ty: &LogicalType) -> bool {
80 matches!(
81 ty,
82 LogicalType::TinyInt
83 | LogicalType::SmallInt
84 | LogicalType::Integer
85 | LogicalType::BigInt
86 | LogicalType::HugeInt
87 | LogicalType::UTinyInt
88 | LogicalType::USmallInt
89 | LogicalType::UInteger
90 | LogicalType::UBigInt
91 | LogicalType::Date
92 | LogicalType::Decimal { .. }
93 )
94}
95
96fn key_at(chunk: &Chunk, values: &rudb_vector::Vector, row: usize) -> Result<Option<i128>> {
104 if let Some(key) = values.signed_at(row) {
105 return Ok(Some(key));
106 }
107 match chunk.value_at(row, 0) {
108 Value::Null => Ok(None),
109 Value::TinyInt(key) => Ok(Some(i128::from(key))),
110 Value::SmallInt(key) => Ok(Some(i128::from(key))),
111 Value::Integer(key) | Value::Date(key) => Ok(Some(i128::from(key))),
112 Value::BigInt(key) | Value::Time(key) | Value::Timestamp(key) => Ok(Some(i128::from(key))),
113 Value::HugeInt(key) | Value::Decimal { unscaled: key, .. } => Ok(Some(key)),
114 Value::UTinyInt(key) => Ok(Some(i128::from(key))),
115 Value::USmallInt(key) => Ok(Some(i128::from(key))),
116 Value::UInteger(key) => Ok(Some(i128::from(key))),
117 Value::UBigInt(key) => Ok(Some(i128::from(key))),
118 other => Err(invalid(&format!("a key column holds {other}, which is not a key"))),
119 }
120}
121
122#[derive(Debug, Clone, Copy)]
129pub struct Built {
130 pub column: usize,
132 pub form: Form,
134 pub rows: u64,
136 pub distinct: bool,
139 pub bytes: usize,
141 pub column_bytes: u64,
143 pub built: bool,
147 pub build: Duration,
149}
150
151pub fn build_key_map(reader: &Reader, column: usize) -> Result<KeyMap> {
157 KeyMap::build_from(&KeyColumn::new(reader, column)?)
158}
159
160pub const BUDGET_SHARE: u64 = 10;
168
169pub const BUDGET_FLOOR: u64 = 64 * 1024;
182
183pub fn build_key_maps(path: &Path, table: &str, columns: &[usize]) -> Result<Vec<Built>> {
193 build_key_maps_within(path, table, columns, BUDGET_SHARE)
194}
195
196pub fn build_key_maps_within(
214 path: &Path,
215 table: &str,
216 columns: &[usize],
217 share: u64,
218) -> Result<Vec<Built>> {
219 let reader = Catalog::open(path)?.table(table)?;
220 let column_bytes = reader.layout().columns_total();
221 let allowance = (column_bytes.saturating_mul(share) / 100).max(BUDGET_FLOOR);
222 let mut spent = held_bytes(&reader, columns)?;
223 let mut report = Vec::with_capacity(columns.len());
224 let mut payloads = Vec::with_capacity(columns.len());
225 for &column in columns {
226 let start = Instant::now();
227 let map = build_key_map(&reader, column)?;
228 let payload = wire::encode(&map, type_tag(&reader.table().fields()[column].ty)?)?;
229 report.push(Built {
230 column,
231 form: map.form(),
232 rows: map.observed().rows,
233 distinct: map.observed().distinct,
234 bytes: payload.bytes.len(),
235 column_bytes,
236 built: false,
237 build: start.elapsed(),
238 });
239 payloads.push((column, payload));
240 }
241 let mut order = (0..payloads.len()).collect::<Vec<_>>();
244 order.sort_by_key(|&at| payloads[at].1.bytes.len());
245 let mut keep = vec![false; payloads.len()];
246 for at in order {
247 if !report[at].distinct {
252 continue;
253 }
254 let cost = payloads[at].1.bytes.len() as u64;
255 if spent.saturating_add(cost) <= allowance {
256 spent += cost;
257 keep[at] = true;
258 report[at].built = true;
259 }
260 }
261 drop(reader);
265 let attachments = payloads
266 .iter()
267 .zip(&keep)
268 .filter(|&(_, &keep)| keep)
269 .map(|((column, payload), _)| {
270 Ok(Attachment {
271 kind: *section::KEY_MAP,
272 id: u64::try_from(*column).map_err(|_| invalid("column index overflow"))?,
273 flags: payload.flags,
274 header_bytes: payload.header_bytes,
275 bytes: &payload.bytes,
276 })
277 })
278 .collect::<Result<Vec<_>>>()?;
279 crate::attach(path, table, &attachments)?;
280 Ok(report)
281}
282
283fn held_bytes(reader: &Reader, replacing: &[usize]) -> Result<u64> {
293 let mut total = 0;
294 for held in reader.table().sections() {
295 if !held.among(section::GRAPH_KINDS) {
296 continue;
297 }
298 let replaced = held.kind == *section::KEY_MAP
299 && replacing.iter().any(|&column| u64::try_from(column) == Ok(held.id));
300 if replaced || !held.usable(reader.table().generation()) {
301 continue;
302 }
303 let Ok(extents) = reader.extents(held) else { continue };
304 total += extents.iter().map(|extent| u64::from(extent.length)).sum::<u64>();
305 }
306 Ok(total)
307}
308
309#[must_use]
318pub fn key_map(reader: &Reader, column: usize) -> Option<KeyMap> {
319 let table = reader.table();
320 let id = u64::try_from(column).ok()?;
321 let held = table
322 .sections()
323 .iter()
324 .find(|section| section.kind == *section::KEY_MAP && section.id == id)?;
325 if !held.usable(table.generation()) {
326 return None;
327 }
328 let (map, tag) = wire::decode(&reader.payload(held).ok()?).ok()?;
329 if tag != type_tag(&table.fields().get(column)?.ty).ok()? {
334 return None;
335 }
336 Some(map)
337}
338
339#[cfg(test)]
340mod tests {
341 use std::fs;
342 use std::path::PathBuf;
343 use std::time::{SystemTime, UNIX_EPOCH};
344
345 use rudb_common::Field;
346 use rudb_graph::Rid;
347 use rudb_vector::Vector;
348
349 use super::*;
350 use crate::Writer;
351
352 fn path(label: &str) -> PathBuf {
353 let stamp = SystemTime::now().duration_since(UNIX_EPOCH).expect("time advances").as_nanos();
354 std::env::temp_dir().join(format!("rudb-graph-{label}-{}-{stamp}.rdb", std::process::id()))
355 }
356
357 fn table_of(label: &str, keys: &[Option<i64>]) -> PathBuf {
359 let path = path(label);
360 let mut writer =
361 Writer::create(&path, "parent", vec![Field::new("key", LogicalType::BigInt)])
362 .expect("new file");
363 for part in keys.chunks(1000) {
364 let values =
365 part.iter().map(|key| key.map_or(Value::Null, Value::BigInt)).collect::<Vec<_>>();
366 let chunk =
367 Chunk::new(vec![Vector::from_values(LogicalType::BigInt, &values).expect("keys")])
368 .expect("one column");
369 writer.append(&chunk).expect("a part");
370 }
371 writer.finish().expect("commit");
372 path
373 }
374
375 fn resolves(keys: &[Option<i64>], map: &KeyMap) {
377 for (rid, key) in keys.iter().enumerate() {
378 let Some(key) = *key else { continue };
379 let found =
380 map.lookup(i128::from(key)).expect("lookup").expect("a key in the column resolves");
381 assert_eq!(found, rid as Rid, "key {key} resolved to {found} rather than {rid}");
382 }
383 }
384
385 #[test]
386 fn a_key_map_built_over_a_file_resolves_every_key_to_its_own_row() {
387 let keys = (1..=3000_i64).map(Some).collect::<Vec<_>>();
392 let path = table_of("identity", &keys);
393 let built = build_key_maps(&path, "parent", &[0]).expect("build");
394 assert_eq!(built.len(), 1);
395 assert_eq!(built[0].form, Form::Identity);
396 assert_eq!(built[0].rows, 3000);
397 assert!(built[0].distinct);
398 assert_eq!(built[0].bytes, wire::HEADER_BYTES, "identity is a header and nothing else");
399
400 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
401 let map = key_map(&reader, 0).expect("the map is in the file");
402 assert_eq!(map.form(), Form::Identity);
403 resolves(&keys, &map);
404 assert_eq!(map.lookup(0).expect("a key below the column"), None);
405 assert_eq!(map.lookup(3001).expect("a key past the column"), None);
406
407 fs::remove_file(&path).expect("clean up");
408 }
409
410 #[test]
411 fn a_column_with_gaps_takes_the_bitmap_form_and_still_resolves() {
412 let keys = (0..2000_i64).map(|value| Some(value * 4 + 7)).collect::<Vec<_>>();
413 let path = table_of("dense", &keys);
414 let built = build_key_maps(&path, "parent", &[0]).expect("build");
415 assert_eq!(built[0].form, Form::Dense);
416
417 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
418 let map = key_map(&reader, 0).expect("the map is in the file");
419 resolves(&keys, &map);
420 assert_eq!(map.lookup(8).expect("a value in the range but not the column"), None);
421
422 fs::remove_file(&path).expect("clean up");
423 }
424
425 #[test]
426 fn a_column_out_of_order_takes_the_sorted_form_and_still_resolves() {
427 let keys = (0..1500_i64).map(|value| Some((value * 7919) % 100_003)).collect::<Vec<_>>();
428 let path = table_of("sorted", &keys);
429 let built = build_key_maps(&path, "parent", &[0]).expect("build");
430 assert_eq!(built[0].form, Form::Sorted);
431 assert!(built[0].distinct, "the sort settles distinctness for an unordered column");
432
433 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
434 let map = key_map(&reader, 0).expect("the map is in the file");
435 resolves(&keys, &map);
436
437 fs::remove_file(&path).expect("clean up");
438 }
439
440 #[test]
441 fn a_null_in_the_key_column_does_not_shift_the_rows_after_it() {
442 let mut keys = (1..=1200_i64).map(Some).collect::<Vec<_>>();
447 keys[3] = None;
448 keys[900] = None;
449 let path = table_of("nulls", &keys);
450 let built = build_key_maps(&path, "parent", &[0]).expect("build");
451 assert_eq!(built[0].rows, 1198, "a null is not a key");
452
453 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
454 let map = key_map(&reader, 0).expect("the map is in the file");
455 resolves(&keys, &map);
456
457 fs::remove_file(&path).expect("clean up");
458 }
459
460 #[test]
461 fn a_column_with_a_repeat_in_it_is_mapped_and_reported_as_no_parent() {
462 let mut keys = (1..=500_i64).map(Some).collect::<Vec<_>>();
467 keys[200] = Some(7);
468 let path = table_of("repeat", &keys);
469 let built = build_key_maps(&path, "parent", &[0]).expect("build");
470 assert!(!built[0].distinct, "a repeat is observed rather than declared away");
471 assert!(!built[0].built, "and a map no rid can be resolved through is not kept");
472 assert!(built[0].bytes > 0, "what it would have cost is still reported");
473
474 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
475 assert!(key_map(&reader, 0).is_none(), "nothing was written to read back");
476 assert!(reader.table().sections().is_empty());
477
478 fs::remove_file(&path).expect("clean up");
479 }
480
481 #[test]
482 fn a_table_with_no_key_map_answers_with_none_rather_than_an_error() {
483 let path = table_of("absent", &[Some(1), Some(2)]);
486 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
487 assert!(key_map(&reader, 0).is_none());
488 assert!(key_map(&reader, 99).is_none(), "a column that does not exist is not a panic");
489 fs::remove_file(&path).expect("clean up");
490 }
491
492 #[test]
493 fn a_stale_key_map_is_ignored_and_the_table_still_reads() {
494 let path = table_of("stale", &(1..=100_i64).map(Some).collect::<Vec<_>>());
495 build_key_maps(&path, "parent", &[0]).expect("build");
496
497 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
500 assert!(key_map(&reader, 0).is_some());
501 let generation = reader.table().generation();
502 drop(reader);
503
504 let mut entry = Catalog::open(&path)
506 .expect("reopen")
507 .table("parent")
508 .expect("the table")
509 .table()
510 .sections()[0];
511 assert!(entry.usable(generation));
512 entry.generation = generation + 1;
513 assert!(!entry.usable(generation), "a rewrite invalidates rather than corrupts");
514
515 fs::remove_file(&path).expect("clean up");
516 }
517
518 #[test]
519 fn a_torn_key_map_costs_the_shortcut_and_not_the_query() {
520 let keys = (1..=200_i64).map(Some).collect::<Vec<_>>();
521 let path = table_of("torn", &keys);
522 build_key_maps(&path, "parent", &[0]).expect("build");
523
524 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
525 let extent = reader
526 .extents(&reader.table().sections()[0])
527 .expect("extent table")
528 .first()
529 .copied()
530 .expect("one extent");
531 drop(reader);
532 let file = fs::OpenOptions::new().write(true).open(&path).expect("reopen to corrupt");
533 crate::write_at(&file, extent.offset, &[0xff; 8]).expect("flip the header");
534 drop(file);
535
536 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
537 assert!(key_map(&reader, 0).is_none(), "a payload that does not checksum is not a map");
538 assert_eq!(reader.table().rows(), 200, "and the table is untouched");
539
540 fs::remove_file(&path).expect("clean up");
541 }
542
543 #[test]
544 fn a_column_with_no_integer_key_form_is_refused_by_name() {
545 let path = path("varchar");
546 let mut writer =
547 Writer::create(&path, "parent", vec![Field::new("name", LogicalType::Varchar)])
548 .expect("new file");
549 let chunk = Chunk::new(vec![
550 Vector::from_values(LogicalType::Varchar, &[Value::Varchar("a".into())])
551 .expect("one name"),
552 ])
553 .expect("one column");
554 writer.append(&chunk).expect("a part");
555 writer.finish().expect("commit");
556
557 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
558 let error = KeyColumn::new(&reader, 0).expect_err("a string key needs its codes");
559 assert!(error.to_string().contains("integer key form"), "{error}");
560
561 fs::remove_file(&path).expect("clean up");
562 }
563
564 #[test]
565 fn several_columns_are_mapped_in_one_commit() {
566 let path = path("two_columns");
567 let mut writer = Writer::create(
568 &path,
569 "parent",
570 vec![
571 Field::required("id", LogicalType::BigInt),
572 Field::required("code", LogicalType::Integer),
573 ],
574 )
575 .expect("new file");
576 let ids = (1..=400_i64).map(Value::BigInt).collect::<Vec<_>>();
577 let codes = (1..=400_i32).map(|code| Value::Integer(code * 3)).collect::<Vec<_>>();
578 let chunk = Chunk::new(vec![
579 Vector::from_values(LogicalType::BigInt, &ids).expect("ids"),
580 Vector::from_values(LogicalType::Integer, &codes).expect("codes"),
581 ])
582 .expect("two columns");
583 writer.append(&chunk).expect("a part");
584 writer.finish().expect("commit");
585
586 let built = build_key_maps(&path, "parent", &[0, 1]).expect("build both");
587 assert_eq!(built.len(), 2);
588 assert_eq!(built[0].form, Form::Identity);
589 assert_eq!(built[1].form, Form::Dense);
590
591 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
592 assert_eq!(reader.table().sections().len(), 2, "one commit and two entries");
593 assert_eq!(key_map(&reader, 0).expect("the id map").form(), Form::Identity);
594 assert_eq!(key_map(&reader, 1).expect("the code map").form(), Form::Dense);
595 assert_eq!(
596 key_map(&reader, 1).expect("the code map").lookup(9).expect("lookup"),
597 Some(2),
598 "the third code is the third row"
599 );
600
601 fs::remove_file(&path).expect("clean up");
602 }
603
604 #[test]
605 fn the_statistics_sections_do_not_count_against_the_graph_budget() {
606 let keys = (1..=3000_i64).map(Some).collect::<Vec<_>>();
612 let path = table_of("apart", &keys);
613 crate::stats::build_stats(&path, "parent", &[0]).expect("summaries first");
614
615 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
616 let statistics = reader
617 .table()
618 .sections()
619 .iter()
620 .filter(|held| held.among(section::STATISTICS_KINDS))
621 .count();
622 assert_eq!(statistics, 2, "a summary and a sketch are in the file");
623 assert_eq!(held_bytes(&reader, &[0]).expect("held"), 0, "and neither is the graph's");
624
625 drop(reader);
626 fs::remove_file(&path).expect("clean up");
627 }
628
629 #[test]
630 fn a_map_that_does_not_fit_the_budget_is_measured_and_not_written() {
631 let keys = (0..100_000_i64).map(|value| Some(value * 8)).collect::<Vec<_>>();
638 let path = table_of("budget", &keys);
639 let built = build_key_maps(&path, "parent", &[0]).expect("build");
640 assert_eq!(built[0].form, Form::Dense);
641 assert!(!built[0].built, "a map ten times its column does not fit a tenth of it");
642 assert!(built[0].bytes as u64 > built[0].column_bytes, "{built:?}");
643
644 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
645 assert!(reader.table().sections().is_empty(), "and nothing was written");
646 assert!(key_map(&reader, 0).is_none());
647 drop(reader);
648
649 let built = build_key_maps_within(&path, "parent", &[0], 100_000).expect("build");
652 assert!(built[0].built);
653 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
654 let map = key_map(&reader, 0).expect("the map is in the file");
655 resolves(&keys, &map);
656
657 fs::remove_file(&path).expect("clean up");
658 }
659
660 #[test]
661 fn the_budget_admits_the_cheapest_maps_it_can_fit() {
662 let path = path("budget_order");
668 let mut writer = Writer::create(
669 &path,
670 "parent",
671 vec![
672 Field::required("id", LogicalType::BigInt),
673 Field::required("code", LogicalType::BigInt),
674 ],
675 )
676 .expect("new file");
677 let ids = (1..=100_000_i64).map(Value::BigInt).collect::<Vec<_>>();
678 let codes = (1..=100_000_i64)
679 .map(|code| Value::BigInt((code * 2_147_483_647) % 999_999_937))
680 .collect::<Vec<_>>();
681 for part in 0..100 {
682 let at = part * 1000;
683 let chunk = Chunk::new(vec![
684 Vector::from_values(LogicalType::BigInt, &ids[at..at + 1000]).expect("ids"),
685 Vector::from_values(LogicalType::BigInt, &codes[at..at + 1000]).expect("codes"),
686 ])
687 .expect("two columns");
688 writer.append(&chunk).expect("a part");
689 }
690 writer.finish().expect("commit");
691
692 let built = build_key_maps(&path, "parent", &[1, 0]).expect("build");
693 assert_eq!(built[0].column, 1, "the report is in the order it was asked in");
694 assert_eq!(built[0].form, Form::Sorted);
695 assert!(!built[0].built, "the sorted map did not fit: {built:?}");
696 assert!(built[1].built, "the identity map did, and was reached second: {built:?}");
697
698 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
699 assert!(key_map(&reader, 0).is_some());
700 assert!(key_map(&reader, 1).is_none());
701
702 fs::remove_file(&path).expect("clean up");
703 }
704}