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> {
289 let mut total = 0;
290 for held in reader.table().sections() {
291 let replaced = held.kind == *section::KEY_MAP
292 && replacing.iter().any(|&column| u64::try_from(column) == Ok(held.id));
293 if replaced || !held.usable(reader.table().generation()) {
294 continue;
295 }
296 let Ok(extents) = reader.extents(held) else { continue };
297 total += extents.iter().map(|extent| u64::from(extent.length)).sum::<u64>();
298 }
299 Ok(total)
300}
301
302#[must_use]
311pub fn key_map(reader: &Reader, column: usize) -> Option<KeyMap> {
312 let table = reader.table();
313 let id = u64::try_from(column).ok()?;
314 let held = table
315 .sections()
316 .iter()
317 .find(|section| section.kind == *section::KEY_MAP && section.id == id)?;
318 if !held.usable(table.generation()) {
319 return None;
320 }
321 let (map, tag) = wire::decode(&reader.payload(held).ok()?).ok()?;
322 if tag != type_tag(&table.fields().get(column)?.ty).ok()? {
327 return None;
328 }
329 Some(map)
330}
331
332#[cfg(test)]
333mod tests {
334 use std::fs;
335 use std::path::PathBuf;
336 use std::time::{SystemTime, UNIX_EPOCH};
337
338 use rudb_common::Field;
339 use rudb_graph::Rid;
340 use rudb_vector::Vector;
341
342 use super::*;
343 use crate::Writer;
344
345 fn path(label: &str) -> PathBuf {
346 let stamp = SystemTime::now().duration_since(UNIX_EPOCH).expect("time advances").as_nanos();
347 std::env::temp_dir().join(format!("rudb-graph-{label}-{}-{stamp}.rdb", std::process::id()))
348 }
349
350 fn table_of(label: &str, keys: &[Option<i64>]) -> PathBuf {
352 let path = path(label);
353 let mut writer =
354 Writer::create(&path, "parent", vec![Field::new("key", LogicalType::BigInt)])
355 .expect("new file");
356 for part in keys.chunks(1000) {
357 let values =
358 part.iter().map(|key| key.map_or(Value::Null, Value::BigInt)).collect::<Vec<_>>();
359 let chunk =
360 Chunk::new(vec![Vector::from_values(LogicalType::BigInt, &values).expect("keys")])
361 .expect("one column");
362 writer.append(&chunk).expect("a part");
363 }
364 writer.finish().expect("commit");
365 path
366 }
367
368 fn resolves(keys: &[Option<i64>], map: &KeyMap) {
370 for (rid, key) in keys.iter().enumerate() {
371 let Some(key) = *key else { continue };
372 let found =
373 map.lookup(i128::from(key)).expect("lookup").expect("a key in the column resolves");
374 assert_eq!(found, rid as Rid, "key {key} resolved to {found} rather than {rid}");
375 }
376 }
377
378 #[test]
379 fn a_key_map_built_over_a_file_resolves_every_key_to_its_own_row() {
380 let keys = (1..=3000_i64).map(Some).collect::<Vec<_>>();
385 let path = table_of("identity", &keys);
386 let built = build_key_maps(&path, "parent", &[0]).expect("build");
387 assert_eq!(built.len(), 1);
388 assert_eq!(built[0].form, Form::Identity);
389 assert_eq!(built[0].rows, 3000);
390 assert!(built[0].distinct);
391 assert_eq!(built[0].bytes, wire::HEADER_BYTES, "identity is a header and nothing else");
392
393 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
394 let map = key_map(&reader, 0).expect("the map is in the file");
395 assert_eq!(map.form(), Form::Identity);
396 resolves(&keys, &map);
397 assert_eq!(map.lookup(0).expect("a key below the column"), None);
398 assert_eq!(map.lookup(3001).expect("a key past the column"), None);
399
400 fs::remove_file(&path).expect("clean up");
401 }
402
403 #[test]
404 fn a_column_with_gaps_takes_the_bitmap_form_and_still_resolves() {
405 let keys = (0..2000_i64).map(|value| Some(value * 4 + 7)).collect::<Vec<_>>();
406 let path = table_of("dense", &keys);
407 let built = build_key_maps(&path, "parent", &[0]).expect("build");
408 assert_eq!(built[0].form, Form::Dense);
409
410 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
411 let map = key_map(&reader, 0).expect("the map is in the file");
412 resolves(&keys, &map);
413 assert_eq!(map.lookup(8).expect("a value in the range but not the column"), None);
414
415 fs::remove_file(&path).expect("clean up");
416 }
417
418 #[test]
419 fn a_column_out_of_order_takes_the_sorted_form_and_still_resolves() {
420 let keys = (0..1500_i64).map(|value| Some((value * 7919) % 100_003)).collect::<Vec<_>>();
421 let path = table_of("sorted", &keys);
422 let built = build_key_maps(&path, "parent", &[0]).expect("build");
423 assert_eq!(built[0].form, Form::Sorted);
424 assert!(built[0].distinct, "the sort settles distinctness for an unordered column");
425
426 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
427 let map = key_map(&reader, 0).expect("the map is in the file");
428 resolves(&keys, &map);
429
430 fs::remove_file(&path).expect("clean up");
431 }
432
433 #[test]
434 fn a_null_in_the_key_column_does_not_shift_the_rows_after_it() {
435 let mut keys = (1..=1200_i64).map(Some).collect::<Vec<_>>();
440 keys[3] = None;
441 keys[900] = None;
442 let path = table_of("nulls", &keys);
443 let built = build_key_maps(&path, "parent", &[0]).expect("build");
444 assert_eq!(built[0].rows, 1198, "a null is not a key");
445
446 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
447 let map = key_map(&reader, 0).expect("the map is in the file");
448 resolves(&keys, &map);
449
450 fs::remove_file(&path).expect("clean up");
451 }
452
453 #[test]
454 fn a_column_with_a_repeat_in_it_is_mapped_and_reported_as_no_parent() {
455 let mut keys = (1..=500_i64).map(Some).collect::<Vec<_>>();
460 keys[200] = Some(7);
461 let path = table_of("repeat", &keys);
462 let built = build_key_maps(&path, "parent", &[0]).expect("build");
463 assert!(!built[0].distinct, "a repeat is observed rather than declared away");
464 assert!(!built[0].built, "and a map no rid can be resolved through is not kept");
465 assert!(built[0].bytes > 0, "what it would have cost is still reported");
466
467 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
468 assert!(key_map(&reader, 0).is_none(), "nothing was written to read back");
469 assert!(reader.table().sections().is_empty());
470
471 fs::remove_file(&path).expect("clean up");
472 }
473
474 #[test]
475 fn a_table_with_no_key_map_answers_with_none_rather_than_an_error() {
476 let path = table_of("absent", &[Some(1), Some(2)]);
479 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
480 assert!(key_map(&reader, 0).is_none());
481 assert!(key_map(&reader, 99).is_none(), "a column that does not exist is not a panic");
482 fs::remove_file(&path).expect("clean up");
483 }
484
485 #[test]
486 fn a_stale_key_map_is_ignored_and_the_table_still_reads() {
487 let path = table_of("stale", &(1..=100_i64).map(Some).collect::<Vec<_>>());
488 build_key_maps(&path, "parent", &[0]).expect("build");
489
490 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
493 assert!(key_map(&reader, 0).is_some());
494 let generation = reader.table().generation();
495 drop(reader);
496
497 let mut entry = Catalog::open(&path)
499 .expect("reopen")
500 .table("parent")
501 .expect("the table")
502 .table()
503 .sections()[0];
504 assert!(entry.usable(generation));
505 entry.generation = generation + 1;
506 assert!(!entry.usable(generation), "a rewrite invalidates rather than corrupts");
507
508 fs::remove_file(&path).expect("clean up");
509 }
510
511 #[test]
512 fn a_torn_key_map_costs_the_shortcut_and_not_the_query() {
513 let keys = (1..=200_i64).map(Some).collect::<Vec<_>>();
514 let path = table_of("torn", &keys);
515 build_key_maps(&path, "parent", &[0]).expect("build");
516
517 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
518 let extent = reader
519 .extents(&reader.table().sections()[0])
520 .expect("extent table")
521 .first()
522 .copied()
523 .expect("one extent");
524 drop(reader);
525 let file = fs::OpenOptions::new().write(true).open(&path).expect("reopen to corrupt");
526 crate::write_at(&file, extent.offset, &[0xff; 8]).expect("flip the header");
527 drop(file);
528
529 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
530 assert!(key_map(&reader, 0).is_none(), "a payload that does not checksum is not a map");
531 assert_eq!(reader.table().rows(), 200, "and the table is untouched");
532
533 fs::remove_file(&path).expect("clean up");
534 }
535
536 #[test]
537 fn a_column_with_no_integer_key_form_is_refused_by_name() {
538 let path = path("varchar");
539 let mut writer =
540 Writer::create(&path, "parent", vec![Field::new("name", LogicalType::Varchar)])
541 .expect("new file");
542 let chunk = Chunk::new(vec![
543 Vector::from_values(LogicalType::Varchar, &[Value::Varchar("a".into())])
544 .expect("one name"),
545 ])
546 .expect("one column");
547 writer.append(&chunk).expect("a part");
548 writer.finish().expect("commit");
549
550 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
551 let error = KeyColumn::new(&reader, 0).expect_err("a string key needs its codes");
552 assert!(error.to_string().contains("integer key form"), "{error}");
553
554 fs::remove_file(&path).expect("clean up");
555 }
556
557 #[test]
558 fn several_columns_are_mapped_in_one_commit() {
559 let path = path("two_columns");
560 let mut writer = Writer::create(
561 &path,
562 "parent",
563 vec![
564 Field::required("id", LogicalType::BigInt),
565 Field::required("code", LogicalType::Integer),
566 ],
567 )
568 .expect("new file");
569 let ids = (1..=400_i64).map(Value::BigInt).collect::<Vec<_>>();
570 let codes = (1..=400_i32).map(|code| Value::Integer(code * 3)).collect::<Vec<_>>();
571 let chunk = Chunk::new(vec![
572 Vector::from_values(LogicalType::BigInt, &ids).expect("ids"),
573 Vector::from_values(LogicalType::Integer, &codes).expect("codes"),
574 ])
575 .expect("two columns");
576 writer.append(&chunk).expect("a part");
577 writer.finish().expect("commit");
578
579 let built = build_key_maps(&path, "parent", &[0, 1]).expect("build both");
580 assert_eq!(built.len(), 2);
581 assert_eq!(built[0].form, Form::Identity);
582 assert_eq!(built[1].form, Form::Dense);
583
584 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
585 assert_eq!(reader.table().sections().len(), 2, "one commit and two entries");
586 assert_eq!(key_map(&reader, 0).expect("the id map").form(), Form::Identity);
587 assert_eq!(key_map(&reader, 1).expect("the code map").form(), Form::Dense);
588 assert_eq!(
589 key_map(&reader, 1).expect("the code map").lookup(9).expect("lookup"),
590 Some(2),
591 "the third code is the third row"
592 );
593
594 fs::remove_file(&path).expect("clean up");
595 }
596
597 #[test]
598 fn a_map_that_does_not_fit_the_budget_is_measured_and_not_written() {
599 let keys = (0..100_000_i64).map(|value| Some(value * 8)).collect::<Vec<_>>();
606 let path = table_of("budget", &keys);
607 let built = build_key_maps(&path, "parent", &[0]).expect("build");
608 assert_eq!(built[0].form, Form::Dense);
609 assert!(!built[0].built, "a map ten times its column does not fit a tenth of it");
610 assert!(built[0].bytes as u64 > built[0].column_bytes, "{built:?}");
611
612 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
613 assert!(reader.table().sections().is_empty(), "and nothing was written");
614 assert!(key_map(&reader, 0).is_none());
615 drop(reader);
616
617 let built = build_key_maps_within(&path, "parent", &[0], 100_000).expect("build");
620 assert!(built[0].built);
621 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
622 let map = key_map(&reader, 0).expect("the map is in the file");
623 resolves(&keys, &map);
624
625 fs::remove_file(&path).expect("clean up");
626 }
627
628 #[test]
629 fn the_budget_admits_the_cheapest_maps_it_can_fit() {
630 let path = path("budget_order");
636 let mut writer = Writer::create(
637 &path,
638 "parent",
639 vec![
640 Field::required("id", LogicalType::BigInt),
641 Field::required("code", LogicalType::BigInt),
642 ],
643 )
644 .expect("new file");
645 let ids = (1..=100_000_i64).map(Value::BigInt).collect::<Vec<_>>();
646 let codes = (1..=100_000_i64)
647 .map(|code| Value::BigInt((code * 2_147_483_647) % 999_999_937))
648 .collect::<Vec<_>>();
649 for part in 0..100 {
650 let at = part * 1000;
651 let chunk = Chunk::new(vec![
652 Vector::from_values(LogicalType::BigInt, &ids[at..at + 1000]).expect("ids"),
653 Vector::from_values(LogicalType::BigInt, &codes[at..at + 1000]).expect("codes"),
654 ])
655 .expect("two columns");
656 writer.append(&chunk).expect("a part");
657 }
658 writer.finish().expect("commit");
659
660 let built = build_key_maps(&path, "parent", &[1, 0]).expect("build");
661 assert_eq!(built[0].column, 1, "the report is in the order it was asked in");
662 assert_eq!(built[0].form, Form::Sorted);
663 assert!(!built[0].built, "the sorted map did not fit: {built:?}");
664 assert!(built[1].built, "the identity map did, and was reached second: {built:?}");
665
666 let reader = Catalog::open(&path).expect("reopen").table("parent").expect("the table");
667 assert!(key_map(&reader, 0).is_some());
668 assert!(key_map(&reader, 1).is_none());
669
670 fs::remove_file(&path).expect("clean up");
671 }
672}