1use std::collections::HashMap;
38use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering as Atomic};
39use std::sync::{Arc, Mutex};
40
41use rudb_common::{Error, LogicalType, Result};
42use rudb_metrics::{LoadProfile, Stage};
43use rudb_storage::Range;
44use rudb_vector::{Bitmap, Chunk, Data, StringColumn, Validity, Vector};
45
46use super::{
47 ColumnStripe, DICTIONARY_CHECK_SEED, DICTIONARY_DECIDE_ROWS, DICTIONARY_DISTINCT_IN_TEN,
48 GlobalDictionary, MAX_ENCODE_WORKERS, MAX_PAGE, Part, PendingChunk, STRIPE_PARTS, Spread,
49 Writer, checksum, coded_page, invalid, push_validity, seeded_checksum, stats, unique_codes,
50 weight,
51};
52
53static BUSY: AtomicUsize = AtomicUsize::new(0);
60
61struct Share(usize);
63
64impl Share {
65 fn take(columns: usize, parts: usize) -> Self {
66 let busy = BUSY.fetch_add(1, Atomic::Relaxed) + 1;
67 let cores =
68 std::thread::available_parallelism().map_or(1, usize::from).min(MAX_ENCODE_WORKERS);
69 let workers = if parts <= 1 { 1 } else { (cores / busy).clamp(1, columns.max(1)) };
71 Self(workers)
72 }
73}
74
75impl Drop for Share {
76 fn drop(&mut self) {
77 BUSY.fetch_sub(1, Atomic::Relaxed);
78 }
79}
80
81#[derive(Debug, Clone)]
87pub struct Preparer {
88 types: Vec<LogicalType>,
89 coded: Arc<[AtomicBool]>,
90 profile: Option<Arc<LoadProfile>>,
91}
92
93#[derive(Debug)]
95pub struct Prepared {
96 parts: Vec<Part>,
97 types: Vec<LogicalType>,
98 columns: Vec<Column>,
99 gathers: Vec<Option<stats::Gather>>,
100 profile: Option<Arc<LoadProfile>>,
101}
102
103#[derive(Debug)]
105pub struct Merged {
106 parts: Vec<Part>,
107 columns: Vec<Merge>,
108 profile: Option<Arc<LoadProfile>>,
109}
110
111#[derive(Debug)]
113pub struct Paged {
114 parts: Vec<Part>,
115 columns: Vec<ColumnStripe>,
116}
117
118#[derive(Debug)]
120enum Column {
121 Pages(ColumnStripe),
123 Coded(Local),
125}
126
127#[derive(Debug)]
129enum Merge {
130 Pages(ColumnStripe),
131 Codes {
133 parts: Vec<LocalPart>,
134 global: Vec<u32>,
135 },
136 Plain(Local),
140}
141
142const END: u32 = u32::MAX;
144
145#[derive(Debug, Default)]
151struct Local {
152 first: HashMap<u64, u32, Spread>,
154 next: Vec<u32>,
156 hashes: Vec<u64>,
157 checks: Vec<u64>,
158 bytes: Vec<u8>,
160 ends: Vec<usize>,
161 counts: Vec<u64>,
163 nulls: u64,
164 parts: Vec<LocalPart>,
165}
166
167#[derive(Debug)]
169struct LocalPart {
170 codes: Vec<u32>,
171 validity: Vec<u8>,
173 range: Range,
174}
175
176impl Local {
177 fn code_column(index: usize, held: &[PendingChunk]) -> Result<Self> {
183 let mut local = Self::default();
184 for pending in held {
185 let column = pending.chunk.column(index)?;
186 let flat = column.flatten()?;
188 let mut codes = Vec::with_capacity(flat.len());
189 let mut last = None;
190 for row in 0..flat.len() {
191 let text = flat.text_at(row).unwrap_or("").as_bytes();
192 let code = match last {
195 Some(code) if local.value(code) == text => code,
196 _ => local.code(text)?,
197 };
198 last = Some(code);
199 if flat.is_null_at(row) {
200 local.nulls += 1;
201 } else {
202 local.counts[code as usize] += 1;
203 }
204 codes.push(code);
205 }
206 let mut validity = Vec::new();
207 push_validity(&mut validity, &flat);
208 local.parts.push(LocalPart { codes, validity, range: Range::of(column) });
209 }
210 local.first = HashMap::default();
213 local.next = Vec::new();
214 Ok(local)
215 }
216
217 fn rows(&self) -> Result<Vec<Vector>> {
223 self.parts
224 .iter()
225 .map(|part| {
226 let len = part.codes.len();
227 let mut column = StringColumn::with_capacity(len);
228 for &code in &part.codes {
229 column.push_bytes(self.value(code));
230 }
231 let validity = match part.validity.split_first() {
232 Some((0, _)) => Validity::AllValid,
233 Some((1, _)) => Validity::AllInvalid,
234 Some((2, bits)) => {
235 let mut mask = Bitmap::all_valid(len);
236 for row in (0..len).filter(|row| bits[row / 8] & (1 << (row % 8)) == 0) {
237 mask.set(row, false);
238 }
239 Validity::Mask(mask)
240 }
241 _ => return Err(Error::internal("a coded part has no validity")),
242 };
243 Ok(Vector::flat(LogicalType::Varchar, Data::Varlen(column))?
244 .with_validity(validity))
245 })
246 .collect()
247 }
248
249 fn values(&self) -> usize {
250 self.ends.len()
251 }
252
253 fn value(&self, code: u32) -> &[u8] {
254 let code = code as usize;
255 let from = if code == 0 { 0 } else { self.ends[code - 1] };
256 &self.bytes[from..self.ends[code]]
257 }
258
259 fn code(&mut self, text: &[u8]) -> Result<u32> {
260 let hash = checksum(text);
261 let Some(&first) = self.first.get(&hash) else {
262 let code = self.push(text, hash)?;
263 self.first.insert(hash, code);
264 return Ok(code);
265 };
266 let mut at = first;
267 loop {
268 if self.value(at) == text {
269 return Ok(at);
270 }
271 match self.next[at as usize] {
272 END => break,
273 next => at = next,
274 }
275 }
276 let code = self.push(text, hash)?;
277 self.next[at as usize] = code;
278 Ok(code)
279 }
280
281 fn push(&mut self, text: &[u8], hash: u64) -> Result<u32> {
282 let code = u32::try_from(self.ends.len())
283 .ok()
284 .filter(|&code| code != END)
285 .ok_or_else(|| invalid("a stripe has too many values in one column"))?;
286 self.bytes.extend_from_slice(text);
287 self.ends.push(self.bytes.len());
288 self.next.push(END);
289 self.hashes.push(hash);
290 self.checks.push(seeded_checksum(text, DICTIONARY_CHECK_SEED));
291 self.counts.push(0);
292 Ok(code)
293 }
294
295 fn merge_into(&self, dictionary: &mut GlobalDictionary) -> Result<Vec<u32>> {
298 let mut global = Vec::with_capacity(self.values());
299 for (code, (&hash, &check)) in self.hashes.iter().zip(&self.checks).enumerate() {
300 let text = self.value(code as u32);
301 let at = dictionary.code_hashed(text, hash, check)?;
302 let count = dictionary
303 .counts
304 .get_mut(at as usize)
305 .ok_or_else(|| invalid("global dictionary count code is out of range"))?;
306 *count = count.saturating_add(self.counts[code]);
307 global.push(at);
308 }
309 dictionary.nulls = dictionary.nulls.saturating_add(self.nulls);
310 Ok(global)
311 }
312}
313
314fn drops_dictionary(rows: usize, distinct: usize) -> bool {
338 rows >= DICTIONARY_DECIDE_ROWS
339 && distinct.saturating_mul(10) > rows.saturating_mul(DICTIONARY_DISTINCT_IN_TEN)
340}
341
342fn fan_out<T: Send>(
352 jobs: Vec<usize>,
353 workers: usize,
354 profile: Option<&LoadProfile>,
355 work: impl Fn(usize) -> Result<T> + Sync,
356) -> Result<Vec<(usize, T)>> {
357 if workers <= 1 || jobs.len() <= 1 {
358 let _span = profile.map(|profile| profile.span(Stage::Pages));
359 return jobs.into_iter().map(|index| Ok((index, work(index)?))).collect();
360 }
361 let workers = workers.min(jobs.len());
362 let queue = Mutex::new(jobs);
363 let pieces = std::thread::scope(|scope| {
364 (0..workers)
365 .map(|_| {
366 scope.spawn(|| {
367 let _span = profile.map(|profile| profile.span(Stage::Pages));
368 let mut mine = Vec::new();
369 loop {
370 let taken = queue
371 .lock()
372 .map_err(|_| Error::internal("a native encode worker panicked"))?
373 .pop();
374 let Some(index) = taken else { break };
375 mine.push((index, work(index)?));
376 }
377 Ok(mine)
378 })
379 })
380 .collect::<Vec<_>>()
381 .into_iter()
382 .map(|handle| {
383 handle.join().map_err(|_| Error::internal("a native encode worker panicked"))?
384 })
385 .collect::<Result<Vec<Vec<_>>>>()
386 })?;
387 Ok(pieces.into_iter().flatten().collect())
388}
389
390fn column_of(held: &[PendingChunk], index: usize) -> Result<Vec<&Vector>> {
392 held.iter().map(|pending| pending.chunk.column(index)).collect()
393}
394
395fn in_order<T>(width: usize, done: Vec<(usize, T)>) -> Result<Vec<T>> {
397 let mut slots: Vec<Option<T>> = (0..width).map(|_| None).collect();
398 for (index, one) in done {
399 slots[index] = Some(one);
400 }
401 slots
402 .into_iter()
403 .map(|slot| slot.ok_or_else(|| Error::internal("a column was never encoded")))
404 .collect()
405}
406
407impl Preparer {
408 pub fn prepare(&self, parts: Vec<((u64, u64), Chunk)>) -> Result<Prepared> {
419 if parts.len() > STRIPE_PARTS {
420 return Err(invalid("a stripe was handed more parts than it holds"));
421 }
422 let held = parts
423 .into_iter()
424 .filter(|(_, chunk)| !chunk.is_empty())
425 .map(|(order, chunk)| PendingChunk { order, chunk })
426 .collect::<Vec<_>>();
427 for pending in &held {
428 self.fits(&pending.chunk)?;
429 }
430 self.prepare_held(held)
431 }
432
433 fn fits(&self, chunk: &Chunk) -> Result<()> {
435 if chunk.width() != self.types.len() {
436 return Err(invalid("chunk width differs from table schema"));
437 }
438 for (index, ty) in self.types.iter().enumerate() {
439 if chunk.column(index)?.logical_type() != ty {
440 return Err(invalid("chunk type differs from table schema"));
441 }
442 }
443 Ok(())
444 }
445
446 pub(crate) fn prepare_held(&self, held: Vec<PendingChunk>) -> Result<Prepared> {
447 let width = self.types.len();
448 let key = held.first().map_or((0, 0), |pending| pending.order);
449 let share = Share::take(width, held.len());
450 let mut jobs = (0..width).collect::<Vec<_>>();
451 jobs.sort_by_key(|&index| weight(&self.types[index]));
452 let done = fan_out(jobs, share.0, self.profile.as_deref(), |index| {
453 let gather = stats::Gather::new(&self.types[index], 0)
456 .filter(|_| !held.is_empty())
457 .map(|mut gather| {
458 gather.stripe(
459 key,
460 held.iter().filter_map(|pending| pending.chunk.column(index).ok()),
461 );
462 gather
463 });
464 let column = if self.coded[index].load(Atomic::Relaxed) {
465 Column::Coded(Local::code_column(index, &held)?)
466 } else {
467 Column::Pages(Writer::encode_pages(&column_of(&held, index)?)?)
468 };
469 Ok((column, gather))
470 })?;
471 drop(share);
472 let (columns, gathers) = in_order(width, done)?.into_iter().unzip();
473 let parts = held.iter().map(Part::of).collect();
474 drop(held);
475 Ok(Prepared {
476 parts,
477 types: self.types.clone(),
478 columns,
479 gathers,
480 profile: self.profile.clone(),
481 })
482 }
483}
484
485impl Merged {
486 pub fn pages(self) -> Result<Paged> {
493 let Self { parts, columns, profile } = self;
494 let width = columns.len();
495 let mut jobs = (0..width)
496 .filter(|&index| !matches!(columns[index], Merge::Pages(_)))
497 .collect::<Vec<_>>();
498 jobs.sort_by_key(|&index| matches!(columns[index], Merge::Plain(_)));
500 let share = Share::take(jobs.len(), parts.len());
501 let built = fan_out(jobs, share.0, profile.as_deref(), |index| match &columns[index] {
502 Merge::Codes { parts, global } => code_pages(parts, global),
503 Merge::Plain(local) => Writer::encode_pages(&local.rows()?.iter().collect::<Vec<_>>()),
504 Merge::Pages(_) => Err(Error::internal("a finished column was queued to be built")),
505 })?;
506 drop(share);
507 let mut slots: Vec<Option<ColumnStripe>> = (0..width).map(|_| None).collect();
508 for (index, stripe) in built {
509 slots[index] = Some(stripe);
510 }
511 let columns = columns
512 .into_iter()
513 .zip(slots)
514 .map(|(column, slot)| match (column, slot) {
515 (Merge::Pages(stripe), _) | (_, Some(stripe)) => Ok(stripe),
516 _ => Err(Error::internal("a column was never encoded")),
517 })
518 .collect::<Result<Vec<_>>>()?;
519 Ok(Paged { parts, columns })
520 }
521}
522
523fn code_pages(parts: &[LocalPart], global: &[u32]) -> Result<ColumnStripe> {
525 let mut stripe = ColumnStripe {
526 pages: Vec::with_capacity(parts.len()),
527 codes: Vec::with_capacity(parts.len()),
528 sieves: Vec::with_capacity(parts.len()),
529 ranges: Vec::with_capacity(parts.len()),
530 };
531 for part in parts {
532 let codes = part
533 .codes
534 .iter()
535 .map(|&code| global.get(code as usize).copied())
536 .collect::<Option<Vec<_>>>()
537 .ok_or_else(|| Error::internal("a stripe's code has no global code"))?;
538 let bytes = coded_page(&codes, &part.validity)?;
539 if bytes.len() > MAX_PAGE {
540 return Err(invalid("column page exceeds the configured bound"));
541 }
542 stripe.pages.push(bytes);
543 stripe.codes.push(Some(unique_codes(&codes)));
544 stripe.sieves.push(None);
548 stripe.ranges.push(part.range.clone());
549 }
550 Ok(stripe)
551}
552
553impl Writer {
554 #[must_use]
559 pub fn preparer(&self) -> Preparer {
560 Preparer {
561 types: self.table.fields.iter().map(|field| field.ty.clone()).collect(),
562 coded: Arc::clone(&self.coded),
563 profile: self.profile.clone(),
564 }
565 }
566
567 pub fn merge(&mut self, prepared: Prepared) -> Result<Merged> {
578 self.flush_pending()?;
579 if prepared.columns.len() != self.table.fields.len()
580 || prepared.types.iter().ne(self.table.fields.iter().map(|field| &field.ty))
581 {
582 return Err(invalid("a stripe was prepared for a table of other columns"));
583 }
584 self.table.rows = prepared
585 .parts
586 .iter()
587 .try_fold(self.table.rows, |rows, part| rows.checked_add(part.rows))
588 .ok_or_else(|| invalid("row count overflow"))?;
589 self.merge_held(prepared)
590 }
591
592 pub(crate) fn merge_held(&mut self, prepared: Prepared) -> Result<Merged> {
594 let Prepared { parts, columns, gathers, profile, .. } = prepared;
595 let timing = profile.as_deref().map(|profile| profile.span(Stage::Dictionary));
596 for (mine, stripe) in self.gathers.iter_mut().zip(gathers) {
597 if let (Some(mine), Some(stripe)) = (mine, stripe) {
598 mine.absorb(stripe);
599 }
600 }
601 let rows: usize = parts.iter().map(|part| part.rows).sum();
602 let mut merged = Vec::with_capacity(columns.len());
603 for (index, column) in columns.into_iter().enumerate() {
604 let dictionary = &mut self.dictionaries[index];
605 merged.push(match (column, dictionary.as_mut()) {
606 (Column::Pages(stripe), None) => Merge::Pages(stripe),
607 (Column::Pages(_), Some(_)) => {
608 return Err(Error::internal(
609 "a column with a global dictionary was prepared without one",
610 ));
611 }
612 (Column::Coded(local), None) => Merge::Plain(local),
613 (Column::Coded(local), Some(global)) => {
614 if global.values() == 0 && drops_dictionary(rows, local.values()) {
617 *dictionary = None;
618 self.coded[index].store(false, Atomic::Relaxed);
619 Merge::Plain(local)
620 } else {
621 let global = local.merge_into(global)?;
622 Merge::Codes { parts: local.parts, global }
623 }
624 }
625 });
626 }
627 drop(timing);
628 Ok(Merged { parts, columns: merged, profile })
629 }
630
631 pub fn write(&mut self, paged: Paged) -> Result<()> {
637 self.write_paged(paged)
638 }
639
640 pub(crate) fn write_paged(&mut self, paged: Paged) -> Result<()> {
641 if paged.parts.is_empty() {
642 return Ok(());
643 }
644 self.write_stripe(&paged.parts, paged.columns)
645 }
646
647 pub fn append_prepared(&mut self, prepared: Prepared) -> Result<()> {
653 let merged = self.merge(prepared)?;
654 let paged = merged.pages()?;
655 self.write(paged)
656 }
657}
658
659#[cfg(test)]
660mod tests {
661 use std::fs;
662 use std::path::PathBuf;
663 use std::time::{SystemTime, UNIX_EPOCH};
664
665 use rudb_common::{Field, Value};
666 use rudb_vector::Vector;
667
668 use super::*;
669 use crate::Reader;
670
671 const PART: usize = 1_000;
672
673 fn path(label: &str) -> PathBuf {
674 let stamp = SystemTime::now().duration_since(UNIX_EPOCH).expect("time advances").as_nanos();
675 std::env::temp_dir()
676 .join(format!("rudb-prepare-{label}-{}-{stamp}.rdb", std::process::id()))
677 }
678
679 fn fields() -> Vec<Field> {
680 vec![
681 Field::required("id", LogicalType::BigInt),
682 Field::new("city", LogicalType::Varchar),
683 Field::new("note", LogicalType::Varchar),
684 ]
685 }
686
687 fn row(id: usize) -> [Value; 3] {
692 let city = if id % 11 == 0 {
693 Value::Null
694 } else {
695 Value::Varchar(format!("city {}", (id / 7) % 13))
696 };
697 let note = if id % 17 == 0 { Value::Null } else { Value::Varchar(format!("note {id}")) };
698 [Value::BigInt(id as i64), city, note]
699 }
700
701 fn stripe(first: usize, parts: usize) -> Vec<((u64, u64), Chunk)> {
703 (first..first + parts)
704 .map(|part| {
705 let rows = (part * PART..(part + 1) * PART).map(row).collect::<Vec<_>>();
706 let column = |at: usize| {
707 let values = rows.iter().map(|row| row[at].clone()).collect::<Vec<_>>();
708 Vector::from_values(fields()[at].ty.clone(), &values).expect("a column")
709 };
710 let chunk = Chunk::new(vec![column(0), column(1), column(2)]).expect("a chunk");
711 ((part as u64, 0), chunk)
712 })
713 .collect()
714 }
715
716 fn runs() -> Vec<Vec<((u64, u64), Chunk)>> {
718 vec![stripe(5, 5), stripe(0, 5), stripe(10, 3)]
719 }
720
721 fn check(path: &PathBuf) {
722 let reader = Reader::open(path).expect("reopen");
723 assert_eq!(reader.parts(), 13);
724 for part in 0..13 {
725 let chunk = reader.read(part, &[0, 1, 2]).expect("a part");
726 for at in [0, 17, PART - 1] {
727 let want = row(part * PART + at);
728 for (column, value) in want.iter().enumerate() {
729 assert_eq!(&chunk.value_at(at, column), value, "part {part} row {at}");
730 }
731 }
732 }
733 }
734
735 #[test]
743 fn stripes_prepared_before_any_is_merged_write_the_same_bytes_as_one_at_a_time() {
744 let alone = path("alone");
745 let mut writer = Writer::create(&alone, "t", fields()).expect("a file");
746 for run in runs() {
747 writer.append_stripe(run).expect("a stripe");
748 }
749 writer.finish().expect("commit");
750
751 let split = path("split");
752 let mut writer = Writer::create(&split, "t", fields()).expect("a file");
753 let preparer = writer.preparer();
754 let prepared = runs()
755 .into_iter()
756 .map(|run| preparer.prepare(run).expect("prepared"))
757 .collect::<Vec<_>>();
758 for one in prepared {
759 writer.append_prepared(one).expect("a stripe");
760 }
761 assert!(!preparer.coded[2].load(Atomic::Relaxed), "note lost its dictionary");
762 assert!(preparer.coded[1].load(Atomic::Relaxed), "city kept its dictionary");
763 writer.finish().expect("commit");
764
765 assert_eq!(fs::read(&alone).expect("read"), fs::read(&split).expect("read"));
766 check(&split);
767 fs::remove_file(alone).expect("remove");
768 fs::remove_file(split).expect("remove");
769 }
770
771 #[test]
775 fn stripes_written_in_another_order_than_they_were_merged_read_back() {
776 let path = path("crossed");
777 let mut writer = Writer::create(&path, "t", fields()).expect("a file");
778 let preparer = writer.preparer();
779 let mut merged = runs()
780 .into_iter()
781 .map(|run| writer.merge(preparer.prepare(run).expect("prepared")).expect("merged"))
782 .map(|merged| merged.pages().expect("paged"))
783 .collect::<Vec<_>>();
784 merged.reverse();
785 for paged in merged {
786 writer.write(paged).expect("written");
787 }
788 writer.finish().expect("commit");
789 check(&path);
790 fs::remove_file(path).expect("remove");
791 }
792
793 #[test]
796 fn a_stripe_of_another_table_is_refused_at_the_merge() {
797 let path = path("refused");
798 let mut writer = Writer::create(&path, "t", fields()).expect("a file");
799 let other = Writer::create(path.with_extension("other"), "u", vec![fields().remove(0)])
800 .expect("a file");
801 let prepared = other.preparer().prepare(vec![]).expect("nothing to prepare");
802 assert!(writer.merge(prepared).is_err());
803 assert_eq!(writer.table.rows, 0);
804 drop(other);
805 fs::remove_file(path.with_extension("other")).expect("remove");
806 fs::remove_file(path).expect("remove");
807 }
808}