1use std::collections::HashMap;
45use std::ops::Deref;
46use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering as Atomic};
47use std::sync::{Arc, Mutex};
48
49use rudb_common::{Error, LogicalType, Result};
50use rudb_metrics::{LoadProfile, Stage};
51use rudb_storage::Range;
52use rudb_storage::sieve::Sieve;
53use rudb_vector::{Bitmap, Chunk, Data, StringColumn, Validity, Vector};
54
55use super::{
56 ColumnStripe, DICTIONARY_CHECK_SEED, DICTIONARY_DECIDE_ROWS, DICTIONARY_DISTINCT_IN_TEN,
57 EncodedBlock, GlobalDictionary, MAX_ENCODE_WORKERS, MAX_PAGE, Part, PendingChunk, STRIPE_PARTS,
58 Settling, Spread, Unencoded, Writer, checksum, coded_page, invalid, push_validity,
59 seeded_checksum, stats, unique_codes, weight,
60};
61
62static BUSY: AtomicUsize = AtomicUsize::new(0);
69
70pub const DICTIONARY_CAP_BYTES: u64 = 512 * 1024 * 1024;
78
79#[derive(Debug)]
86pub(crate) struct Coding {
87 flags: Box<[AtomicBool]>,
88 growth: Box<[AtomicU64]>,
90 held: AtomicU64,
92 cap: AtomicU64,
93 ceilings: Box<[AtomicU64]>,
102}
103
104impl Coding {
105 pub(crate) fn new(flags: impl IntoIterator<Item = bool>) -> Self {
106 let flags = flags.into_iter().map(AtomicBool::new).collect::<Box<[_]>>();
107 let growth = flags.iter().map(|_| AtomicU64::new(0)).collect();
108 let ceilings = flags.iter().map(|_| AtomicU64::new(u64::MAX)).collect();
109 Self {
110 flags,
111 growth,
112 held: AtomicU64::new(0),
113 cap: AtomicU64::new(DICTIONARY_CAP_BYTES),
114 ceilings,
115 }
116 }
117
118 pub(crate) fn cap(&self, bytes: u64) {
120 self.cap.store(bytes, Atomic::Relaxed);
121 }
122
123 fn recount(&self, before: u64, now: u64) -> u64 {
126 if now >= before {
127 self.held.fetch_add(now - before, Atomic::Relaxed) + (now - before)
128 } else {
129 self.held.fetch_sub(before - now, Atomic::Relaxed).saturating_sub(before - now)
130 }
131 }
132
133 fn grew_most(&self, index: usize) -> bool {
140 let mine = self.growth[index].load(Atomic::Relaxed);
141 mine > 0 && self.growth.iter().all(|other| other.load(Atomic::Relaxed) <= mine)
142 }
143}
144
145impl Deref for Coding {
146 type Target = [AtomicBool];
147
148 fn deref(&self) -> &[AtomicBool] {
149 &self.flags
150 }
151}
152
153struct Share(usize);
155
156impl Share {
157 fn take(columns: usize, parts: usize) -> Self {
158 let busy = BUSY.fetch_add(1, Atomic::Relaxed) + 1;
159 let cores =
160 std::thread::available_parallelism().map_or(1, usize::from).min(MAX_ENCODE_WORKERS);
161 let workers = if parts <= 1 { 1 } else { (cores / busy).clamp(1, columns.max(1)) };
163 Self(workers)
164 }
165}
166
167impl Drop for Share {
168 fn drop(&mut self) {
169 BUSY.fetch_sub(1, Atomic::Relaxed);
170 }
171}
172
173#[derive(Debug, Clone)]
179pub struct Preparer {
180 types: Vec<LogicalType>,
181 coded: Arc<Coding>,
182 profile: Option<Arc<LoadProfile>>,
183}
184
185#[derive(Debug)]
187pub struct Prepared {
188 parts: Vec<Part>,
189 types: Vec<LogicalType>,
190 columns: Vec<Column>,
191 gathers: Vec<Option<stats::Gather>>,
192 profile: Option<Arc<LoadProfile>>,
193}
194
195#[derive(Debug)]
197pub struct Merged {
198 parts: Vec<Part>,
199 columns: Vec<Merge>,
200 blocks: Vec<Unencoded>,
201 profile: Option<Arc<LoadProfile>>,
202 counted: bool,
205}
206
207#[derive(Debug)]
209pub struct Paged {
210 parts: Vec<Part>,
211 columns: Vec<ColumnStripe>,
212 blocks: Vec<(usize, usize, EncodedBlock)>,
214 counted: bool,
215}
216
217enum Built {
219 Stripe(ColumnStripe),
220 Block(EncodedBlock),
221}
222
223#[derive(Debug)]
225enum Column {
226 Pages(ColumnStripe),
228 Coded(Local),
230}
231
232#[derive(Debug)]
234enum Merge {
235 Pages(ColumnStripe),
236 Codes {
238 parts: Vec<LocalPart>,
239 global: Vec<u32>,
240 },
241 Plain(Local),
245}
246
247const END: u32 = u32::MAX;
249
250#[derive(Debug, Default)]
256struct Local {
257 first: HashMap<u64, u32, Spread>,
259 next: Vec<u32>,
261 hashes: Vec<u64>,
262 checks: Vec<u64>,
263 bytes: Vec<u8>,
265 ends: Vec<usize>,
266 counts: Vec<u64>,
268 nulls: u64,
269 parts: Vec<LocalPart>,
270 blob: bool,
272}
273
274#[derive(Debug)]
276struct LocalPart {
277 codes: Vec<u32>,
278 validity: Vec<u8>,
280 range: Range,
281}
282
283impl Local {
284 fn held(&self) -> usize {
286 self.first.capacity() * (size_of::<u64>() + size_of::<u32>() + 1)
288 + spilled(&self.next)
289 + spilled(&self.hashes)
290 + spilled(&self.checks)
291 + spilled(&self.bytes)
292 + spilled(&self.ends)
293 + spilled(&self.counts)
294 + spilled(&self.parts)
295 + self
296 .parts
297 .iter()
298 .map(|part| spilled(&part.codes) + spilled(&part.validity))
299 .sum::<usize>()
300 }
301
302 #[cfg(test)]
304 fn code_column(index: usize, held: &[PendingChunk]) -> Result<Self> {
305 let mut local = Self::default();
306 let mut mapped = None;
307 for pending in held {
308 local.code_part(pending.chunk.column(index)?, &mut mapped)?;
309 }
310 local.done();
311 Ok(local)
312 }
313
314 fn code_part(
323 &mut self,
324 column: &Vector,
325 mapped: &mut Option<(Arc<Vector>, Vec<u32>)>,
326 ) -> Result<()> {
327 self.blob = column.logical_type() == &LogicalType::Blob;
328 if let Some(codes) = self.code_dictionary(column, mapped)? {
329 let mut validity = Vec::new();
330 push_validity(&mut validity, column);
331 self.parts.push(LocalPart { codes, validity, range: Range::of(column) });
332 return Ok(());
333 }
334 let flat = column.flatten()?;
336 let mut codes = Vec::with_capacity(flat.len());
337 let mut last = None;
338 for row in 0..flat.len() {
339 let text = flat.bytes_at(row).unwrap_or(b"");
342 let code = match last {
345 Some(code) if self.value(code) == text => code,
346 _ => self.code(text)?,
347 };
348 last = Some(code);
349 if flat.is_null_at(row) {
350 self.nulls += 1;
351 } else {
352 self.counts[code as usize] += 1;
353 }
354 codes.push(code);
355 }
356 let mut validity = Vec::new();
357 push_validity(&mut validity, &flat);
358 self.parts.push(LocalPart { codes, validity, range: Range::of(column) });
359 Ok(())
360 }
361
362 fn done(&mut self) {
364 self.first = HashMap::default();
367 self.next = Vec::new();
368 }
369
370 fn code_dictionary(
385 &mut self,
386 column: &Vector,
387 mapped: &mut Option<(Arc<Vector>, Vec<u32>)>,
388 ) -> Result<Option<Vec<u32>>> {
389 let Some((codes, values)) = column.shared_dictionary_parts() else { return Ok(None) };
390 if !matches!(values.validity(), Validity::AllValid) {
391 return Ok(None);
392 }
393 let Some(codes) = codes.get(..column.len()) else { return Ok(None) };
394 let fresh = !matches!(mapped, Some((held, _)) if Arc::ptr_eq(held, values));
395 if fresh {
396 *mapped = Some((Arc::clone(values), vec![END; values.len()]));
397 }
398 let Some((_, map)) = mapped.as_mut() else { return Ok(None) };
399 let every = matches!(column.validity(), Validity::AllValid);
400 let mut coded = Vec::with_capacity(codes.len());
401 for (row, &code) in codes.iter().enumerate() {
402 if !every && !column.validity().is_valid(row) {
403 let code = self.code(b"")?;
405 self.nulls += 1;
406 coded.push(code);
407 continue;
408 }
409 let slot = map
410 .get_mut(code as usize)
411 .ok_or_else(|| invalid("a dictionary code is out of range"))?;
412 if *slot == END {
413 *slot = self.code(values.bytes_at(code as usize).unwrap_or(b""))?;
414 }
415 self.counts[*slot as usize] += 1;
416 coded.push(*slot);
417 }
418 Ok(Some(coded))
419 }
420
421 fn rows(&self) -> Result<Vec<Vector>> {
427 self.parts
428 .iter()
429 .map(|part| {
430 let len = part.codes.len();
431 let mut column = StringColumn::with_capacity(len);
432 for &code in &part.codes {
433 column.push_bytes(self.value(code));
434 }
435 let validity = match part.validity.split_first() {
436 Some((0, _)) => Validity::AllValid,
437 Some((1, _)) => Validity::AllInvalid,
438 Some((2, bits)) => {
439 let mut mask = Bitmap::all_valid(len);
440 for row in (0..len).filter(|row| bits[row / 8] & (1 << (row % 8)) == 0) {
441 mask.set(row, false);
442 }
443 Validity::Mask(mask)
444 }
445 _ => return Err(Error::internal("a coded part has no validity")),
446 };
447 let ty = if self.blob { LogicalType::Blob } else { LogicalType::Varchar };
448 Ok(Vector::flat(ty, Data::Varlen(column))?.with_validity(validity))
449 })
450 .collect()
451 }
452
453 fn values(&self) -> usize {
454 self.ends.len()
455 }
456
457 fn value(&self, code: u32) -> &[u8] {
458 let code = code as usize;
459 let from = if code == 0 { 0 } else { self.ends[code - 1] };
460 &self.bytes[from..self.ends[code]]
461 }
462
463 fn code(&mut self, text: &[u8]) -> Result<u32> {
464 let hash = checksum(text);
465 let Some(&first) = self.first.get(&hash) else {
466 let code = self.push(text, hash)?;
467 self.first.insert(hash, code);
468 return Ok(code);
469 };
470 let mut at = first;
471 loop {
472 if self.value(at) == text {
473 return Ok(at);
474 }
475 match self.next[at as usize] {
476 END => break,
477 next => at = next,
478 }
479 }
480 let code = self.push(text, hash)?;
481 self.next[at as usize] = code;
482 Ok(code)
483 }
484
485 fn push(&mut self, text: &[u8], hash: u64) -> Result<u32> {
486 let code = u32::try_from(self.ends.len())
487 .ok()
488 .filter(|&code| code != END)
489 .ok_or_else(|| invalid("a stripe has too many values in one column"))?;
490 self.bytes.extend_from_slice(text);
491 self.ends.push(self.bytes.len());
492 self.next.push(END);
493 self.hashes.push(hash);
494 self.checks.push(seeded_checksum(text, DICTIONARY_CHECK_SEED));
495 self.counts.push(0);
496 Ok(code)
497 }
498
499 fn merge_into(&self, dictionary: &mut GlobalDictionary) -> Result<Vec<u32>> {
502 let mut global = Vec::with_capacity(self.values());
503 for (code, (&hash, &check)) in self.hashes.iter().zip(&self.checks).enumerate() {
504 let text = self.value(code as u32);
505 let at = dictionary.code_hashed(text, hash, check)?;
506 let count = dictionary
507 .counts
508 .get_mut(at as usize)
509 .ok_or_else(|| invalid("global dictionary count code is out of range"))?;
510 *count = count.saturating_add(self.counts[code]);
511 global.push(at);
512 }
513 dictionary.nulls = dictionary.nulls.saturating_add(self.nulls);
514 Ok(global)
515 }
516}
517
518fn drops_dictionary(rows: usize, distinct: usize) -> bool {
542 rows >= DICTIONARY_DECIDE_ROWS
543 && distinct.saturating_mul(10) > rows.saturating_mul(DICTIONARY_DISTINCT_IN_TEN)
544}
545
546fn demotes(rows: usize, new: usize, total: u64, coding: &Coding, index: usize) -> bool {
559 drops_dictionary(rows, new)
560 || (total > coding.cap.load(Atomic::Relaxed) && coding.grew_most(index))
561}
562
563fn fan_out<T: Send>(
573 jobs: Vec<usize>,
574 workers: usize,
575 profile: Option<&LoadProfile>,
576 work: impl Fn(usize) -> Result<T> + Sync,
577) -> Result<Vec<(usize, T)>> {
578 if workers <= 1 || jobs.len() <= 1 {
579 let _span = profile.map(|profile| profile.span(Stage::Pages));
580 return jobs.into_iter().map(|index| Ok((index, work(index)?))).collect();
581 }
582 let workers = workers.min(jobs.len());
583 let queue = Mutex::new(jobs);
584 let pieces = std::thread::scope(|scope| {
585 (0..workers)
586 .map(|_| {
587 scope.spawn(|| {
588 let _span = profile.map(|profile| profile.span(Stage::Pages));
589 let mut mine = Vec::new();
590 loop {
591 let taken = queue
592 .lock()
593 .map_err(|_| Error::internal("a native encode worker panicked"))?
594 .pop();
595 let Some(index) = taken else { break };
596 mine.push((index, work(index)?));
597 }
598 Ok(mine)
599 })
600 })
601 .collect::<Vec<_>>()
602 .into_iter()
603 .map(|handle| {
604 handle.join().map_err(|_| Error::internal("a native encode worker panicked"))?
605 })
606 .collect::<Result<Vec<Vec<_>>>>()
607 })?;
608 Ok(pieces.into_iter().flatten().collect())
609}
610
611impl Preparer {
612 pub fn prepare(&self, parts: Vec<((u64, u64), Chunk)>) -> Result<Prepared> {
623 if parts.len() > STRIPE_PARTS {
624 return Err(invalid("a stripe was handed more parts than it holds"));
625 }
626 let held = parts
627 .into_iter()
628 .filter(|(_, chunk)| !chunk.is_empty())
629 .map(|(order, chunk)| PendingChunk { order, chunk })
630 .collect::<Vec<_>>();
631 for pending in &held {
632 self.fits(&pending.chunk)?;
633 }
634 self.prepare_held(held)
635 }
636
637 fn fits(&self, chunk: &Chunk) -> Result<()> {
639 if chunk.width() != self.types.len() {
640 return Err(invalid("chunk width differs from table schema"));
641 }
642 for (index, ty) in self.types.iter().enumerate() {
643 if chunk.column(index)?.logical_type() != ty {
644 return Err(invalid("chunk type differs from table schema"));
645 }
646 }
647 Ok(())
648 }
649
650 pub(crate) fn prepare_held(&self, held: Vec<PendingChunk>) -> Result<Prepared> {
651 let mut building = self.start();
652 self.feed_held(&mut building, held)?;
653 self.finish(building)
654 }
655
656 #[must_use]
669 pub fn start(&self) -> Building {
670 let columns = (0..self.types.len())
671 .map(|index| {
672 let body = if self.coded[index].load(Atomic::Relaxed) {
673 Body::Coded(Local::default(), None)
674 } else {
675 Body::Pages(ColumnStripe::default(), Settling::default())
676 };
677 let mut gather = stats::Gather::new(&self.types[index], 0);
678 let ceiling = self.coded.ceilings[index].load(Atomic::Relaxed);
679 if let Some(gather) = gather.as_mut().filter(|_| ceiling < u64::MAX) {
680 gather.cap_at(ceiling);
681 }
682 Mutex::new(Growing { body, gather })
683 })
684 .collect();
685 Building { parts: Vec::new(), columns }
686 }
687
688 pub fn feed(&self, building: &mut Building, parts: Vec<((u64, u64), Chunk)>) -> Result<()> {
698 if building.parts.len().saturating_add(parts.len()) > STRIPE_PARTS {
699 return Err(invalid("a stripe was handed more parts than it holds"));
700 }
701 let held = parts
702 .into_iter()
703 .filter(|(_, chunk)| !chunk.is_empty())
704 .map(|(order, chunk)| PendingChunk { order, chunk })
705 .collect::<Vec<_>>();
706 for pending in &held {
707 self.fits(&pending.chunk)?;
708 }
709 self.feed_held(building, held)
710 }
711
712 fn feed_held(&self, building: &mut Building, held: Vec<PendingChunk>) -> Result<()> {
713 if held.is_empty() {
714 return Ok(());
715 }
716 let width = self.types.len();
717 let opening = building.parts.is_empty();
719 let key = held.first().map_or((0, 0), |pending| pending.order);
720 let share = Share::take(width, held.len());
721 let mut jobs = (0..width).collect::<Vec<_>>();
722 jobs.sort_by_key(|&index| weight(&self.types[index]));
723 let columns = &building.columns;
724 fan_out(jobs, share.0, self.profile.as_deref(), |index| {
725 let mut growing = columns[index]
726 .lock()
727 .map_err(|_| Error::internal("a native encode worker panicked"))?;
728 let Growing { body, gather } = &mut *growing;
729 if matches!(body, Body::Coded(..)) && !self.coded[index].load(Atomic::Relaxed) {
730 body.plain()?;
731 }
732 if let Some(gather) = gather.as_mut() {
733 if opening {
736 gather.open_stripe(key);
737 }
738 for pending in &held {
739 gather.part(pending.chunk.column(index)?);
740 }
741 }
742 match body {
743 Body::Coded(local, mapped) => {
744 for pending in &held {
745 local.code_part(pending.chunk.column(index)?, mapped)?;
746 }
747 }
748 Body::Pages(stripe, settling) => {
749 for pending in &held {
750 Writer::encode_page(stripe, settling, pending.chunk.column(index)?)?;
751 }
752 }
753 }
754 Ok(())
755 })?;
756 drop(share);
757 building.parts.extend(held.iter().map(Part::of));
758 Ok(())
759 }
760
761 pub fn finish(&self, building: Building) -> Result<Prepared> {
767 let empty = building.parts.is_empty();
768 let (columns, gathers) = building
769 .columns
770 .into_iter()
771 .enumerate()
772 .map(|(index, growing)| {
773 let Growing { body, gather } = growing
774 .into_inner()
775 .map_err(|_| Error::internal("a native encode worker panicked"))?;
776 let column = match body {
777 Body::Coded(mut local, _) => {
778 local.done();
779 Column::Coded(local)
780 }
781 Body::Pages(stripe, _) => Column::Pages(stripe),
782 };
783 let gather = gather.filter(|_| !empty).map(|mut gather| {
785 gather.close_stripe();
786 if let Some(ceiling) = gather.ceiling() {
787 self.coded.ceilings[index].fetch_min(ceiling, Atomic::Relaxed);
788 }
789 gather
790 });
791 Ok((column, gather))
792 })
793 .collect::<Result<Vec<_>>>()?
794 .into_iter()
795 .unzip();
796 Ok(Prepared {
797 parts: building.parts,
798 types: self.types.clone(),
799 columns,
800 gathers,
801 profile: self.profile.clone(),
802 })
803 }
804}
805
806#[derive(Debug)]
808pub struct Building {
809 parts: Vec<Part>,
810 columns: Vec<Mutex<Growing>>,
813}
814
815impl Building {
816 #[must_use]
818 pub fn parts(&self) -> usize {
819 self.parts.len()
820 }
821
822 #[must_use]
829 pub fn held(&self) -> u64 {
830 let bytes: usize = self
831 .columns
832 .iter()
833 .map(|growing| {
834 growing.lock().map_or(0, |growing| match &growing.body {
835 Body::Pages(stripe, _) => stripe_bytes(stripe),
836 Body::Coded(local, mapped) => {
837 local.held() + mapped.as_ref().map_or(0, |(_, codes)| spilled(codes))
838 }
839 })
840 })
841 .sum();
842 bytes as u64
843 }
844}
845
846fn stripe_bytes(stripe: &ColumnStripe) -> usize {
848 stripe.pages.iter().map(Vec::capacity).sum::<usize>()
849 + spilled(&stripe.pages)
850 + spilled(&stripe.sums)
851 + stripe.codes.iter().flatten().map(spilled).sum::<usize>()
852 + spilled(&stripe.codes)
853 + stripe.sieves.iter().flatten().map(Sieve::len).sum::<usize>()
854 + spilled(&stripe.sieves)
855 + spilled(&stripe.ranges)
856}
857
858fn spilled<T>(values: &Vec<T>) -> usize {
860 values.capacity() * size_of::<T>()
861}
862
863#[derive(Debug)]
865struct Growing {
866 body: Body,
867 gather: Option<stats::Gather>,
868}
869
870#[derive(Debug)]
872enum Body {
873 Pages(ColumnStripe, Settling),
875 Coded(Local, Option<(Arc<Vector>, Vec<u32>)>),
878}
879
880impl Body {
881 fn plain(&mut self) -> Result<()> {
888 let Self::Coded(local, _) = self else { return Ok(()) };
889 let mut stripe = ColumnStripe::default();
890 let mut settling = Settling::default();
891 for rows in local.rows()? {
892 Writer::encode_page(&mut stripe, &mut settling, &rows)?;
893 }
894 *self = Self::Pages(stripe, settling);
895 Ok(())
896 }
897}
898
899enum Slot<'a> {
901 Owned(&'a mut Option<GlobalDictionary>, &'a mut Option<stats::Gather>),
903 Lent(&'a Mutex<LentColumn>, &'a Lent),
905}
906
907struct Step<'a> {
909 index: usize,
910 column: Column,
911 slot: Slot<'a>,
912 gather: Option<stats::Gather>,
914}
915
916impl Step<'_> {
917 fn cost(&self, coded: &Coding) -> usize {
921 match &self.column {
922 Column::Coded(local) if coded[self.index].load(Atomic::Relaxed) => {
923 local.values().saturating_add(1)
924 }
925 _ => 0,
926 }
927 }
928
929 fn run(
931 self,
932 rows: usize,
933 coded: &Coding,
934 profile: Option<&LoadProfile>,
935 ) -> Result<(usize, Merge, Vec<Unencoded>)> {
936 let Self { index, column, slot, gather } = self;
937 match slot {
938 Slot::Owned(dictionary, mine) => {
939 merge_column(index, column, gather, dictionary, mine, rows, coded, profile)
940 }
941 Slot::Lent(held, lent) => {
942 let mut held = held.lock().map_err(|_| Error::internal("a merge panicked"))?;
943 if lent.reclaimed.load(Atomic::Acquire) {
946 return Err(Error::internal("a stripe was merged after its table was closed"));
947 }
948 let LentColumn { dictionary, gather: mine } = &mut *held;
949 merge_column(index, column, gather, dictionary, mine, rows, coded, profile)
950 }
951 }
952 }
953}
954
955#[expect(clippy::too_many_arguments, reason = "one column's share of the stripe's merge state")]
957fn merge_column(
958 index: usize,
959 column: Column,
960 stripe: Option<stats::Gather>,
961 dictionary: &mut Option<GlobalDictionary>,
962 gather: &mut Option<stats::Gather>,
963 rows: usize,
964 coded: &Coding,
965 profile: Option<&LoadProfile>,
966) -> Result<(usize, Merge, Vec<Unencoded>)> {
967 if let (Some(mine), Some(stripe)) = (gather.as_mut(), stripe) {
968 mine.absorb(stripe);
969 }
970 let mut new = None;
971 let merge = match (column, dictionary.as_mut()) {
972 (Column::Pages(stripe), None) => Merge::Pages(stripe),
973 (Column::Pages(stripe), Some(global)) if global.demoted => Merge::Pages(stripe),
974 (Column::Pages(_), Some(_)) => {
975 return Err(Error::internal(
976 "a column with a global dictionary was prepared without one",
977 ));
978 }
979 (Column::Coded(local), None) => Merge::Plain(local),
980 (Column::Coded(local), Some(global)) if global.demoted => Merge::Plain(local),
982 (Column::Coded(local), Some(global)) => {
983 if global.values() == 0 && drops_dictionary(rows, local.values()) {
986 if let Some(profile) = profile {
987 profile.release(global.charged);
988 }
989 coded.recount(global.charged, 0);
990 *dictionary = None;
991 coded[index].store(false, Atomic::Relaxed);
992 Merge::Plain(local)
993 } else {
994 let before = global.values();
995 let codes = local.merge_into(global)?;
996 new = Some(global.values() - before);
997 Merge::Codes { parts: local.parts, global: codes }
998 }
999 }
1000 };
1001 if let (Some(new), Some(global)) = (new, dictionary.as_mut()) {
1003 let now = global.held_bytes();
1004 coded.growth[index].store(now.saturating_sub(global.charged), Atomic::Relaxed);
1005 let total =
1006 coded.held.load(Atomic::Relaxed).saturating_add(now).saturating_sub(global.charged);
1007 if demotes(rows, new, total, coded, index) {
1008 global.demote();
1009 coded[index].store(false, Atomic::Relaxed);
1010 coded.growth[index].store(0, Atomic::Relaxed);
1011 }
1012 }
1013 let blocks = match dictionary {
1018 Some(dictionary) => {
1019 dictionary.settle()?;
1020 let blocks = dictionary.hand_out(index);
1021 let (before, now) = dictionary.recharge(profile);
1022 coded.recount(before, now);
1023 blocks
1024 }
1025 None => Vec::new(),
1026 };
1027 Ok((index, merge, blocks))
1028}
1029
1030fn merge_columns(prepared: Prepared, slots: Vec<Slot<'_>>, coded: &Coding) -> Result<Merged> {
1038 let Prepared { parts, columns, gathers, profile, .. } = prepared;
1039 let timing = profile.as_deref().map(|profile| profile.span(Stage::Dictionary));
1040 let rows: usize = parts.iter().map(|part| part.rows).sum();
1041 let width = columns.len();
1042 if slots.len() != width || gathers.len() != width {
1043 return Err(Error::internal("a stripe was merged into a table of another width"));
1044 }
1045 let mut steps = columns
1046 .into_iter()
1047 .zip(gathers)
1048 .zip(slots)
1049 .enumerate()
1050 .map(|(index, ((column, gather), slot))| Step { index, column, slot, gather })
1051 .collect::<Vec<_>>();
1052 steps.sort_by_key(|step| step.cost(coded));
1055 let workers = std::thread::available_parallelism()
1056 .map_or(1, usize::from)
1057 .min(MAX_ENCODE_WORKERS)
1058 .min(steps.iter().filter(|step| step.cost(coded) > 0).count())
1059 .max(1);
1060 let done = if workers <= 1 {
1061 steps
1062 .into_iter()
1063 .map(|step| step.run(rows, coded, profile.as_deref()))
1064 .collect::<Result<Vec<_>>>()?
1065 } else {
1066 let queue = Mutex::new(steps);
1067 let pieces = std::thread::scope(|scope| {
1068 (0..workers)
1069 .map(|_| {
1070 scope.spawn(|| {
1071 let mut mine = Vec::new();
1072 loop {
1073 let taken = queue
1074 .lock()
1075 .map_err(|_| Error::internal("a merge worker panicked"))?
1076 .pop();
1077 let Some(step) = taken else { break };
1078 mine.push(step.run(rows, coded, profile.as_deref())?);
1079 }
1080 Ok(mine)
1081 })
1082 })
1083 .collect::<Vec<_>>()
1084 .into_iter()
1085 .map(|handle| {
1086 handle.join().map_err(|_| Error::internal("a merge worker panicked"))?
1087 })
1088 .collect::<Result<Vec<Vec<_>>>>()
1089 })?;
1090 pieces.into_iter().flatten().collect()
1091 };
1092 let mut slots: Vec<Option<(Merge, Vec<Unencoded>)>> = (0..width).map(|_| None).collect();
1093 for (index, merge, blocks) in done {
1094 slots[index] = Some((merge, blocks));
1095 }
1096 let mut merged = Vec::with_capacity(width);
1097 let mut blocks = Vec::new();
1098 for slot in slots {
1099 let (merge, handed) = slot.ok_or_else(|| Error::internal("a column was never merged"))?;
1100 merged.push(merge);
1101 blocks.extend(handed);
1102 }
1103 drop(timing);
1104 Ok(Merged { parts, columns: merged, blocks, profile, counted: false })
1105}
1106
1107#[derive(Debug)]
1109pub(crate) struct Lent {
1110 columns: Box<[Mutex<LentColumn>]>,
1111 reclaimed: AtomicBool,
1113}
1114
1115#[derive(Debug)]
1117pub(crate) struct LentColumn {
1118 pub(crate) dictionary: Option<GlobalDictionary>,
1119 gather: Option<stats::Gather>,
1120}
1121
1122impl Lent {
1123 pub(crate) fn columns(&self) -> &[Mutex<LentColumn>] {
1124 &self.columns
1125 }
1126
1127 fn take_back(&self, blocks: Vec<(usize, usize, EncodedBlock)>) -> Result<()> {
1129 for (column, at, block) in blocks {
1130 self.columns
1131 .get(column)
1132 .ok_or_else(|| Error::internal("a dictionary block came back to no column"))?
1133 .lock()
1134 .map_err(|_| Error::internal("a merge panicked"))?
1135 .dictionary
1136 .as_mut()
1137 .ok_or_else(|| Error::internal("a dictionary block came back to no dictionary"))?
1138 .take_back(at, block)?;
1139 }
1140 Ok(())
1141 }
1142
1143 #[allow(clippy::type_complexity)]
1145 pub(crate) fn reclaim(
1146 &self,
1147 ) -> Result<(Vec<Option<GlobalDictionary>>, Vec<Option<stats::Gather>>)> {
1148 self.reclaimed.store(true, Atomic::Release);
1149 let mut dictionaries = Vec::with_capacity(self.columns.len());
1150 let mut gathers = Vec::with_capacity(self.columns.len());
1151 for column in &self.columns {
1152 let mut held = column.lock().map_err(|_| Error::internal("a merge panicked"))?;
1153 dictionaries.push(held.dictionary.take());
1154 gathers.push(held.gather.take());
1155 }
1156 Ok((dictionaries, gathers))
1157 }
1158}
1159
1160#[derive(Debug, Clone)]
1167pub struct Merger {
1168 lent: Arc<Lent>,
1169 types: Vec<LogicalType>,
1170 coded: Arc<Coding>,
1171}
1172
1173impl Merger {
1174 pub fn merge(&self, prepared: Prepared) -> Result<Merged> {
1180 if prepared.types != self.types {
1181 return Err(invalid("a stripe was prepared for a table of other columns"));
1182 }
1183 let slots = self.lent.columns.iter().map(|column| Slot::Lent(column, &self.lent)).collect();
1184 merge_columns(prepared, slots, &self.coded)
1185 }
1186
1187 pub fn give_back(&self, paged: &mut Paged) -> Result<()> {
1194 self.lent.take_back(std::mem::take(&mut paged.blocks))
1195 }
1196}
1197
1198impl Merged {
1199 pub fn pages(self) -> Result<Paged> {
1207 let Self { parts, columns, blocks, profile, counted } = self;
1208 let width = columns.len();
1209 let mut jobs = (width..width + blocks.len())
1213 .chain((0..width).filter(|&index| !matches!(columns[index], Merge::Pages(_))))
1214 .collect::<Vec<_>>();
1215 jobs.sort_by_key(|&index| index < width && matches!(columns[index], Merge::Plain(_)));
1217 let share = Share::take(jobs.len(), parts.len());
1218 let built = fan_out(jobs, share.0, profile.as_deref(), |index| {
1219 let Some(column) = columns.get(index) else {
1220 return Ok(Built::Block(blocks[index - width].encode()?));
1221 };
1222 Ok(Built::Stripe(match column {
1223 Merge::Codes { parts, global } => code_pages(parts, global)?,
1224 Merge::Plain(local) => {
1225 Writer::encode_pages(&local.rows()?.iter().collect::<Vec<_>>())?
1226 }
1227 Merge::Pages(_) => {
1228 return Err(Error::internal("a finished column was queued to be built"));
1229 }
1230 }))
1231 })?;
1232 drop(share);
1233 let mut slots: Vec<Option<ColumnStripe>> = (0..width).map(|_| None).collect();
1234 let mut encoded = Vec::with_capacity(blocks.len());
1235 for (index, one) in built {
1236 match one {
1237 Built::Stripe(stripe) => slots[index] = Some(stripe),
1238 Built::Block(block) => {
1239 let (column, at) = blocks[index - width].place();
1240 encoded.push((column, at, block));
1241 }
1242 }
1243 }
1244 let columns = columns
1245 .into_iter()
1246 .zip(slots)
1247 .map(|(column, slot)| match (column, slot) {
1248 (Merge::Pages(stripe), _) | (_, Some(stripe)) => Ok(stripe),
1249 _ => Err(Error::internal("a column was never encoded")),
1250 })
1251 .collect::<Result<Vec<_>>>()?;
1252 Ok(Paged { parts, columns, blocks: encoded, counted })
1253 }
1254}
1255
1256fn code_pages(parts: &[LocalPart], global: &[u32]) -> Result<ColumnStripe> {
1258 let mut stripe = ColumnStripe {
1259 pages: Vec::with_capacity(parts.len()),
1260 sums: Vec::with_capacity(parts.len()),
1261 codes: Vec::with_capacity(parts.len()),
1262 sieves: Vec::with_capacity(parts.len()),
1263 ranges: Vec::with_capacity(parts.len()),
1264 };
1265 for part in parts {
1266 let codes = part
1267 .codes
1268 .iter()
1269 .map(|&code| global.get(code as usize).copied())
1270 .collect::<Option<Vec<_>>>()
1271 .ok_or_else(|| Error::internal("a stripe's code has no global code"))?;
1272 let bytes = coded_page(&codes, &part.validity)?;
1273 if bytes.len() > MAX_PAGE {
1274 return Err(invalid("column page exceeds the configured bound"));
1275 }
1276 stripe.sums.push(checksum(&bytes));
1277 stripe.pages.push(bytes);
1278 stripe.codes.push(Some(unique_codes(&codes)));
1279 stripe.sieves.push(None);
1283 stripe.ranges.push(part.range.clone());
1284 }
1285 Ok(stripe)
1286}
1287
1288impl Writer {
1289 #[must_use]
1294 pub fn preparer(&self) -> Preparer {
1295 Preparer {
1296 types: self.table.fields.iter().map(|field| field.ty.clone()).collect(),
1297 coded: Arc::clone(&self.coded),
1298 profile: self.profile.clone(),
1299 }
1300 }
1301
1302 pub fn merge(&mut self, prepared: Prepared) -> Result<Merged> {
1313 self.flush_pending()?;
1314 if prepared.columns.len() != self.table.fields.len()
1315 || prepared.types.iter().ne(self.table.fields.iter().map(|field| &field.ty))
1316 {
1317 return Err(invalid("a stripe was prepared for a table of other columns"));
1318 }
1319 self.table.rows = prepared
1320 .parts
1321 .iter()
1322 .try_fold(self.table.rows, |rows, part| rows.checked_add(part.rows))
1323 .ok_or_else(|| invalid("row count overflow"))?;
1324 self.merge_held(prepared)
1325 }
1326
1327 pub(crate) fn merge_held(&mut self, prepared: Prepared) -> Result<Merged> {
1336 let slots = match &self.lent {
1337 Some(lent) => lent.columns.iter().map(|column| Slot::Lent(column, lent)).collect(),
1338 None => self
1339 .dictionaries
1340 .iter_mut()
1341 .zip(self.gathers.iter_mut())
1342 .map(|(dictionary, gather)| Slot::Owned(dictionary, gather))
1343 .collect::<Vec<_>>(),
1344 };
1345 let mut merged = merge_columns(prepared, slots, &self.coded)?;
1346 merged.counted = true;
1347 Ok(merged)
1348 }
1349
1350 pub fn merger(&mut self) -> Result<Merger> {
1360 self.flush_pending()?;
1361 let lent = match &self.lent {
1362 Some(lent) => Arc::clone(lent),
1363 None => {
1364 let lent = Arc::new(Lent {
1365 columns: std::mem::take(&mut self.dictionaries)
1366 .into_iter()
1367 .zip(std::mem::take(&mut self.gathers))
1368 .map(|(dictionary, gather)| Mutex::new(LentColumn { dictionary, gather }))
1369 .collect(),
1370 reclaimed: AtomicBool::new(false),
1371 });
1372 self.lent = Some(Arc::clone(&lent));
1373 lent
1374 }
1375 };
1376 Ok(Merger {
1377 lent,
1378 types: self.table.fields.iter().map(|field| field.ty.clone()).collect(),
1379 coded: Arc::clone(&self.coded),
1380 })
1381 }
1382
1383 pub fn write(&mut self, paged: Paged) -> Result<()> {
1389 self.write_paged(paged)
1390 }
1391
1392 pub(crate) fn write_paged(&mut self, paged: Paged) -> Result<()> {
1393 let Paged { parts, columns, blocks, counted } = paged;
1394 if !counted {
1395 self.table.rows = parts
1396 .iter()
1397 .try_fold(self.table.rows, |rows, part| rows.checked_add(part.rows))
1398 .ok_or_else(|| invalid("row count overflow"))?;
1399 }
1400 if let Some(lent) = &self.lent {
1401 lent.take_back(blocks)?;
1402 } else {
1403 for (column, at, block) in blocks {
1404 self.dictionaries
1405 .get_mut(column)
1406 .and_then(Option::as_mut)
1407 .ok_or_else(|| {
1408 Error::internal("a dictionary block came back to no dictionary")
1409 })?
1410 .take_back(at, block)?;
1411 }
1412 }
1413 if parts.is_empty() {
1414 return self.place_blocks();
1415 }
1416 self.write_stripe(&parts, columns)
1417 }
1418
1419 pub fn append_prepared(&mut self, prepared: Prepared) -> Result<()> {
1425 let merged = self.merge(prepared)?;
1426 let paged = merged.pages()?;
1427 self.write(paged)
1428 }
1429}
1430
1431#[cfg(test)]
1432mod tests {
1433 use std::fs;
1434 use std::path::PathBuf;
1435 use std::time::{SystemTime, UNIX_EPOCH};
1436
1437 use rudb_common::{Field, Value};
1438 use rudb_vector::Vector;
1439
1440 use super::*;
1441 use crate::Reader;
1442
1443 const PART: usize = 1_000;
1444
1445 #[test]
1451 fn dictionary_parts_code_as_their_rows_would() {
1452 let texts = |values: &[&str]| {
1453 Arc::new(
1454 Vector::from_values(
1455 LogicalType::Varchar,
1456 &values
1457 .iter()
1458 .map(|text| Value::Varchar((*text).to_string()))
1459 .collect::<Vec<_>>(),
1460 )
1461 .expect("a dictionary"),
1462 )
1463 };
1464 let shared = texts(&["b", "a", "", "c", "unused"]);
1465 let other = texts(&["c", "d", "a"]);
1466 let mut nulls = Bitmap::all_valid(6);
1467 nulls.set(1, false);
1468 nulls.set(4, false);
1469 let parts = [
1470 Vector::dictionary_over(vec![3, 3, 1, 0, 2, 1], Arc::clone(&shared)).expect("codes"),
1471 Vector::dictionary_over(vec![0, 3, 1, 1, 3, 2], Arc::clone(&shared))
1472 .expect("codes")
1473 .with_validity(Validity::Mask(nulls)),
1474 Vector::dictionary_over(vec![1, 2, 0, 1], other).expect("codes"),
1475 ];
1476 let held = |flat: bool| {
1477 parts
1478 .iter()
1479 .enumerate()
1480 .map(|(at, part)| PendingChunk {
1481 order: (at as u64, 0),
1482 chunk: Chunk::new(vec![if flat {
1483 part.flatten().expect("flat")
1484 } else {
1485 part.clone()
1486 }])
1487 .expect("a chunk"),
1488 })
1489 .collect::<Vec<_>>()
1490 };
1491 let parquet = held(false);
1492 let before = rudb_common::slow::here();
1493 let coded = Local::code_column(0, &parquet).expect("coded");
1494 assert_eq!(
1495 rudb_common::slow::here().since(before).get(rudb_common::slow::Cause::Flatten),
1496 0,
1497 "a part that came in as codes was flattened",
1498 );
1499 let flat = Local::code_column(0, &held(true)).expect("coded");
1500 assert_eq!(coded.values(), flat.values());
1501 for code in 0..flat.values() as u32 {
1502 assert_eq!(coded.value(code), flat.value(code), "value {code}");
1503 }
1504 assert_eq!(coded.counts, flat.counts);
1505 assert_eq!(coded.nulls, flat.nulls);
1506 assert_eq!(coded.nulls, 2);
1507 assert_eq!(coded.hashes, flat.hashes);
1508 assert_eq!(coded.checks, flat.checks);
1509 assert_eq!(coded.parts.len(), flat.parts.len());
1510 for (coded, flat) in coded.parts.iter().zip(&flat.parts) {
1511 assert_eq!(coded.codes, flat.codes);
1512 assert_eq!(coded.validity, flat.validity);
1513 }
1514 }
1515
1516 fn path(label: &str) -> PathBuf {
1517 let stamp = SystemTime::now().duration_since(UNIX_EPOCH).expect("time advances").as_nanos();
1518 std::env::temp_dir()
1519 .join(format!("rudb-prepare-{label}-{}-{stamp}.rdb", std::process::id()))
1520 }
1521
1522 fn fields() -> Vec<Field> {
1523 vec![
1524 Field::required("id", LogicalType::BigInt),
1525 Field::new("city", LogicalType::Varchar),
1526 Field::new("note", LogicalType::Varchar),
1527 ]
1528 }
1529
1530 fn row(id: usize) -> [Value; 3] {
1535 let city = if id.is_multiple_of(11) {
1536 Value::Null
1537 } else {
1538 Value::Varchar(format!("city {}", (id / 7) % 13))
1539 };
1540 let note =
1541 if id.is_multiple_of(17) { Value::Null } else { Value::Varchar(format!("note {id}")) };
1542 [Value::BigInt(id as i64), city, note]
1543 }
1544
1545 fn stripe(first: usize, parts: usize) -> Vec<((u64, u64), Chunk)> {
1547 (first..first + parts)
1548 .map(|part| {
1549 let rows = (part * PART..(part + 1) * PART).map(row).collect::<Vec<_>>();
1550 let column = |at: usize| {
1551 let values = rows.iter().map(|row| row[at].clone()).collect::<Vec<_>>();
1552 Vector::from_values(fields()[at].ty.clone(), &values).expect("a column")
1553 };
1554 let chunk = Chunk::new(vec![column(0), column(1), column(2)]).expect("a chunk");
1555 ((part as u64, 0), chunk)
1556 })
1557 .collect()
1558 }
1559
1560 fn runs() -> Vec<Vec<((u64, u64), Chunk)>> {
1562 vec![stripe(5, 5), stripe(0, 5), stripe(10, 3)]
1563 }
1564
1565 fn check(path: &PathBuf) {
1566 let reader = Reader::open(path).expect("reopen");
1567 assert_eq!(reader.parts(), 13);
1568 for part in 0..13 {
1569 let chunk = reader.read(part, &[0, 1, 2]).expect("a part");
1570 for at in [0, 17, PART - 1] {
1571 let want = row(part * PART + at);
1572 for (column, value) in want.iter().enumerate() {
1573 assert_eq!(&chunk.value_at(at, column), value, "part {part} row {at}");
1574 }
1575 }
1576 }
1577 }
1578
1579 #[test]
1587 fn stripes_prepared_before_any_is_merged_write_the_same_bytes_as_one_at_a_time() {
1588 let alone = path("alone");
1589 let mut writer = Writer::create(&alone, "t", fields()).expect("a file");
1590 for run in runs() {
1591 writer.append_stripe(run).expect("a stripe");
1592 }
1593 writer.finish().expect("commit");
1594
1595 let split = path("split");
1596 let mut writer = Writer::create(&split, "t", fields()).expect("a file");
1597 let preparer = writer.preparer();
1598 let prepared = runs()
1599 .into_iter()
1600 .map(|run| preparer.prepare(run).expect("prepared"))
1601 .collect::<Vec<_>>();
1602 for one in prepared {
1603 writer.append_prepared(one).expect("a stripe");
1604 }
1605 assert!(!preparer.coded[2].load(Atomic::Relaxed), "note lost its dictionary");
1606 assert!(preparer.coded[1].load(Atomic::Relaxed), "city kept its dictionary");
1607 writer.finish().expect("commit");
1608
1609 assert_eq!(fs::read(&alone).expect("read"), fs::read(&split).expect("read"));
1610 check(&split);
1611 fs::remove_file(alone).expect("remove");
1612 fs::remove_file(split).expect("remove");
1613 }
1614
1615 #[test]
1621 fn stripes_counted_under_an_earlier_ceiling_count_what_they_would_have() {
1622 let estimates = |writer: &Writer| {
1623 writer
1624 .gathers
1625 .iter()
1626 .map(|gather| gather.as_ref().and_then(stats::Gather::distinct))
1627 .collect::<Vec<_>>()
1628 };
1629 let want = fields()
1631 .iter()
1632 .enumerate()
1633 .map(|(column, field)| {
1634 let mut gather = stats::Gather::new(&field.ty, 0)?;
1635 for (_, chunk) in runs().into_iter().flatten() {
1636 gather.part(chunk.column(column).expect("a column"));
1637 }
1638 gather.distinct()
1639 })
1640 .collect::<Vec<_>>();
1641
1642 for reversed in [false, true] {
1643 let split = path("capped");
1644 let mut writer = Writer::create(&split, "t", fields()).expect("a file");
1645 let preparer = writer.preparer();
1646 let mut prepared = runs()
1647 .into_iter()
1648 .map(|run| preparer.prepare(run).expect("prepared"))
1649 .collect::<Vec<_>>();
1650 for column in [0, 2] {
1651 assert!(preparer.coded.ceilings[column].load(Atomic::Relaxed) < u64::MAX);
1652 }
1653 if reversed {
1654 prepared.reverse();
1655 }
1656 for one in prepared {
1657 writer.append_prepared(one).expect("a stripe");
1658 }
1659 assert_eq!(estimates(&writer), want, "reversed {reversed}");
1660 writer.finish().expect("commit");
1661 fs::remove_file(split).expect("remove");
1662 }
1663 }
1664
1665 #[test]
1669 fn stripes_written_in_another_order_than_they_were_merged_read_back() {
1670 let path = path("crossed");
1671 let mut writer = Writer::create(&path, "t", fields()).expect("a file");
1672 let preparer = writer.preparer();
1673 let mut merged = runs()
1674 .into_iter()
1675 .map(|run| writer.merge(preparer.prepare(run).expect("prepared")).expect("merged"))
1676 .map(|merged| merged.pages().expect("paged"))
1677 .collect::<Vec<_>>();
1678 merged.reverse();
1679 for paged in merged {
1680 writer.write(paged).expect("written");
1681 }
1682 writer.finish().expect("commit");
1683 check(&path);
1684 fs::remove_file(path).expect("remove");
1685 }
1686
1687 #[test]
1690 fn stripes_merged_through_a_merger_write_the_same_bytes_as_the_writer() {
1691 let alone = path("alone-merger");
1692 let mut writer = Writer::create(&alone, "t", fields()).expect("a file");
1693 for run in runs() {
1694 writer.append_stripe(run).expect("a stripe");
1695 }
1696 writer.finish().expect("commit");
1697
1698 let lent = path("lent");
1699 let mut writer = Writer::create(&lent, "t", fields()).expect("a file");
1700 let preparer = writer.preparer();
1701 let merger = writer.merger().expect("a merger");
1702 for run in runs() {
1703 let merged = merger.merge(preparer.prepare(run).expect("prepared")).expect("merged");
1704 let mut paged = merged.pages().expect("paged");
1705 merger.give_back(&mut paged).expect("given back");
1706 writer.write(paged).expect("written");
1707 }
1708 assert_eq!(writer.table.rows, 13 * PART);
1709 writer.finish().expect("commit");
1710
1711 assert_eq!(fs::read(&alone).expect("read"), fs::read(&lent).expect("read"));
1712 check(&lent);
1713 fs::remove_file(alone).expect("remove");
1714 fs::remove_file(lent).expect("remove");
1715 }
1716
1717 #[test]
1720 fn stripes_fed_in_batches_write_the_same_bytes_as_prepared_whole() {
1721 let whole = path("whole");
1722 let mut writer = Writer::create(&whole, "t", fields()).expect("a file");
1723 for run in runs() {
1724 writer.append_stripe(run).expect("a stripe");
1725 }
1726 writer.finish().expect("commit");
1727
1728 let fed = path("fed");
1729 let mut writer = Writer::create(&fed, "t", fields()).expect("a file");
1730 let preparer = writer.preparer();
1731 let merger = writer.merger().expect("a merger");
1732 let nothing = Chunk::new(
1733 fields()
1734 .iter()
1735 .map(|field| Vector::from_values(field.ty.clone(), &[]).expect("a column"))
1736 .collect(),
1737 )
1738 .expect("a chunk");
1739 for mut run in runs() {
1740 let mut building = preparer.start();
1741 preparer.feed(&mut building, Vec::new()).expect("fed nothing");
1742 while !run.is_empty() {
1743 let rest = run.split_off(2.min(run.len()));
1744 let mut batch = std::mem::replace(&mut run, rest);
1745 batch.push(((u64::MAX, 0), nothing.clone()));
1746 preparer.feed(&mut building, batch).expect("fed");
1747 }
1748 let merged =
1749 merger.merge(preparer.finish(building).expect("finished")).expect("merged");
1750 let mut paged = merged.pages().expect("paged");
1751 merger.give_back(&mut paged).expect("given back");
1752 writer.write(paged).expect("written");
1753 }
1754 writer.finish().expect("commit");
1755
1756 assert_eq!(fs::read(&whole).expect("read"), fs::read(&fed).expect("read"));
1757 check(&fed);
1758 fs::remove_file(whole).expect("remove");
1759 fs::remove_file(fed).expect("remove");
1760 }
1761
1762 #[test]
1765 fn a_column_dropped_while_its_stripe_is_built_turns_to_pages() {
1766 let path = path("dropped-while-built");
1767 let mut writer = Writer::create(&path, "t", fields()).expect("a file");
1768 let preparer = writer.preparer();
1769 let merger = writer.merger().expect("a merger");
1770 let write = |writer: &mut Writer, building: Building| {
1771 let merged =
1772 merger.merge(preparer.finish(building).expect("finished")).expect("merged");
1773 let mut paged = merged.pages().expect("paged");
1774 merger.give_back(&mut paged).expect("given back");
1775 writer.write(paged).expect("written");
1776 };
1777 let mut runs = runs().into_iter();
1778 let mut first = preparer.start();
1779 let mut second = preparer.start();
1780 let mut later = runs.next().expect("a run");
1781 preparer.feed(&mut second, later.drain(..2).collect()).expect("fed");
1782 preparer.feed(&mut first, runs.next().expect("a run")).expect("fed");
1783 write(&mut writer, first);
1784 let note = |building: &Building| {
1785 matches!(building.columns[2].lock().expect("unpoisoned").body, Body::Pages(..))
1786 };
1787 assert!(!note(&second), "still coded until it is fed again");
1788 preparer.feed(&mut second, later).expect("fed");
1789 assert!(note(&second), "turned to pages once fed after the drop");
1790 write(&mut writer, second);
1791 let mut last = preparer.start();
1792 preparer.feed(&mut last, runs.next().expect("a run")).expect("fed");
1793 write(&mut writer, last);
1794 writer.finish().expect("commit");
1795 check(&path);
1796 fs::remove_file(path).expect("remove");
1797 }
1798
1799 #[test]
1801 fn a_stripe_fed_more_parts_than_it_holds_is_refused() {
1802 let path = path("overfed");
1803 let writer = Writer::create(&path, "t", fields()).expect("a file");
1804 let preparer = writer.preparer();
1805 let mut building = preparer.start();
1806 preparer.feed(&mut building, stripe(0, STRIPE_PARTS - 1)).expect("fed");
1807 assert_eq!(building.parts(), STRIPE_PARTS - 1);
1808 assert!(preparer.feed(&mut building, stripe(STRIPE_PARTS, 2)).is_err());
1809 drop(writer);
1810 let _ = fs::remove_file(path);
1811 }
1812
1813 #[test]
1816 fn stripes_merged_on_several_threads_at_once_read_back() {
1817 let path = path("merged-at-once");
1818 let mut writer = Writer::create(&path, "t", fields()).expect("a file");
1819 let preparer = writer.preparer();
1820 let merger = writer.merger().expect("a merger");
1821 let writer = Mutex::new(writer);
1822 std::thread::scope(|scope| {
1823 for run in runs() {
1824 let (preparer, merger, writer) = (&preparer, &merger, &writer);
1825 scope.spawn(move || {
1826 let merged =
1827 merger.merge(preparer.prepare(run).expect("prepared")).expect("merged");
1828 let mut paged = merged.pages().expect("paged");
1829 merger.give_back(&mut paged).expect("given back");
1830 writer.lock().expect("the writer").write(paged).expect("written");
1831 });
1832 }
1833 });
1834 writer.into_inner().expect("the writer").finish().expect("commit");
1835 check(&path);
1836 fs::remove_file(path).expect("remove");
1837 }
1838
1839 #[test]
1842 fn a_merge_after_the_table_is_closed_is_refused() {
1843 let path = path("late");
1844 let mut writer = Writer::create(&path, "t", fields()).expect("a file");
1845 let preparer = writer.preparer();
1846 let merger = writer.merger().expect("a merger");
1847 writer.finish().expect("commit");
1848 let prepared = preparer.prepare(stripe(0, 2)).expect("prepared");
1849 assert!(merger.merge(prepared).is_err());
1850 fs::remove_file(path).expect("remove");
1851 }
1852
1853 #[test]
1856 fn a_stripe_of_another_table_is_refused_at_the_merge() {
1857 let path = path("refused");
1858 let mut writer = Writer::create(&path, "t", fields()).expect("a file");
1859 let other = Writer::create(path.with_extension("other"), "u", vec![fields().remove(0)])
1860 .expect("a file");
1861 let prepared = other.preparer().prepare(vec![]).expect("nothing to prepare");
1862 assert!(writer.merge(prepared).is_err());
1863 assert_eq!(writer.table.rows, 0);
1864 drop(other);
1865 fs::remove_file(path.with_extension("other")).expect("remove");
1866 fs::remove_file(path).expect("remove");
1867 }
1868
1869 fn turning(id: usize) -> [Value; 3] {
1873 let url = match id {
1874 _ if id.is_multiple_of(13) => Value::Null,
1875 _ if id < 5 * PART => Value::Varchar(format!("https://example.com/{}", id % 20)),
1876 _ => Value::Varchar(format!("https://example.com/page/{id}")),
1877 };
1878 [Value::BigInt(id as i64), Value::Varchar(format!("city {}", id % 13)), url]
1879 }
1880
1881 fn turning_fields() -> Vec<Field> {
1882 vec![
1883 Field::required("id", LogicalType::BigInt),
1884 Field::new("city", LogicalType::Varchar),
1885 Field::new("url", LogicalType::Varchar),
1886 ]
1887 }
1888
1889 fn turning_stripe(
1890 rows: fn(usize) -> [Value; 3],
1891 first: usize,
1892 parts: usize,
1893 ) -> Vec<((u64, u64), Chunk)> {
1894 (first..first + parts)
1895 .map(|part| {
1896 let rows = (part * PART..(part + 1) * PART).map(rows).collect::<Vec<_>>();
1897 let column = |at: usize| {
1898 let values = rows.iter().map(|row| row[at].clone()).collect::<Vec<_>>();
1899 Vector::from_values(turning_fields()[at].ty.clone(), &values).expect("a column")
1900 };
1901 let chunk = Chunk::new(vec![column(0), column(1), column(2)]).expect("a chunk");
1902 ((part as u64, 0), chunk)
1903 })
1904 .collect()
1905 }
1906
1907 fn turning_runs() -> Vec<Vec<((u64, u64), Chunk)>> {
1908 vec![
1909 turning_stripe(turning, 0, 5),
1910 turning_stripe(turning, 5, 5),
1911 turning_stripe(turning, 10, 3),
1912 ]
1913 }
1914
1915 fn check_turned(path: &PathBuf) {
1917 let reader = Reader::open(path).expect("reopen");
1918 assert_eq!(reader.parts(), 13);
1919 assert_eq!(reader.table().demoted, [false, false, true], "only url is demoted");
1920 for part in 0..13 {
1921 let chunk = reader.read(part, &[0, 1, 2]).expect("a part");
1922 let url = chunk.column(2).expect("url");
1923 assert!(url.stable_dictionary_parts().is_none(), "part {part} hands out no codes");
1924 for at in 0..PART {
1925 let want = turning(part * PART + at);
1926 for (column, value) in want.iter().enumerate() {
1927 assert_eq!(&chunk.value_at(at, column), value, "part {part} row {at}");
1928 }
1929 }
1930 }
1931 assert_eq!(reader.distinct_values(2).expect("asked"), None);
1933 assert_eq!(reader.text_extremes(2).expect("asked"), None);
1934 assert_eq!(reader.exact_frequencies(2).expect("asked"), None);
1935 assert_eq!(reader.top_frequencies(2, 5).expect("asked"), None);
1936 assert!(!reader.skips_codes(0, 2, &[0]).expect("asked"), "no code proves a value absent");
1937 assert_eq!(reader.distinct_values(1).expect("asked"), Some(13));
1939 assert!(reader.text_extremes(1).expect("asked").is_some());
1940 }
1941
1942 #[test]
1946 fn a_column_that_turns_unique_is_demoted_and_reads_back() {
1947 let alone = path("demoted-alone");
1948 let mut writer = Writer::create(&alone, "t", turning_fields()).expect("a file");
1949 let preparer = writer.preparer();
1950 let mut runs = turning_runs().into_iter();
1951 writer.append_stripe(runs.next().expect("a run")).expect("a stripe");
1952 assert!(preparer.coded[2].load(Atomic::Relaxed), "url repeats in its first stripe");
1953 for run in runs {
1954 writer.append_stripe(run).expect("a stripe");
1955 }
1956 assert!(!preparer.coded[2].load(Atomic::Relaxed), "url was demoted");
1957 assert!(preparer.coded[1].load(Atomic::Relaxed), "city kept its dictionary");
1958 writer.finish().expect("commit");
1959 check_turned(&alone);
1960
1961 let split = path("demoted-split");
1962 let mut writer = Writer::create(&split, "t", turning_fields()).expect("a file");
1963 let preparer = writer.preparer();
1964 let prepared = turning_runs()
1965 .into_iter()
1966 .map(|run| preparer.prepare(run).expect("prepared"))
1967 .collect::<Vec<_>>();
1968 for one in prepared {
1969 writer.append_prepared(one).expect("a stripe");
1970 }
1971 writer.finish().expect("commit");
1972 check_turned(&split);
1973
1974 fs::remove_file(alone).expect("remove");
1975 fs::remove_file(split).expect("remove");
1976 }
1977
1978 fn growing(id: usize) -> [Value; 3] {
1981 let url = Value::Varchar(format!("https://example.com/{}", id / 4));
1982 [Value::BigInt(id as i64), Value::Varchar(format!("city {}", id % 13)), url]
1983 }
1984
1985 #[test]
1988 fn the_dictionary_cap_demotes_the_column_that_grew_most() {
1989 let path = path("capped");
1990 let mut writer = Writer::create(&path, "t", turning_fields())
1991 .expect("a file")
1992 .with_dictionary_cap(1 << 30);
1993 let preparer = writer.preparer();
1994 writer.append_stripe(turning_stripe(growing, 0, 5)).expect("a stripe");
1995 assert!(preparer.coded[2].load(Atomic::Relaxed), "url is under the cap");
1996 assert!(preparer.coded[1].load(Atomic::Relaxed), "city is under the cap");
1997
1998 writer.coded.cap(1);
1999 writer.append_stripe(turning_stripe(growing, 5, 5)).expect("a stripe");
2000 assert!(!preparer.coded[2].load(Atomic::Relaxed), "url grew most and was demoted");
2001 assert!(preparer.coded[1].load(Atomic::Relaxed), "city grew nothing and keeps it");
2002 writer.append_stripe(turning_stripe(growing, 10, 3)).expect("a stripe");
2003 assert!(preparer.coded[1].load(Atomic::Relaxed), "city still grows nothing");
2004 writer.finish().expect("commit");
2005
2006 let reader = Reader::open(&path).expect("reopen");
2007 assert_eq!(reader.table().demoted, [false, false, true]);
2008 for part in 0..13 {
2009 let chunk = reader.read(part, &[0, 1, 2]).expect("a part");
2010 for at in 0..PART {
2011 let want = growing(part * PART + at);
2012 for (column, value) in want.iter().enumerate() {
2013 assert_eq!(&chunk.value_at(at, column), value, "part {part} row {at}");
2014 }
2015 }
2016 }
2017 assert_eq!(reader.distinct_values(1).expect("asked"), Some(13));
2018 assert_eq!(reader.distinct_values(2).expect("asked"), None);
2019 fs::remove_file(path).expect("remove");
2020 }
2021}