Skip to main content

rudb_native/
prepare.rs

1//! A stripe encoded before the writer is asked for it.
2//!
3//! A load fed by many pipeline instances has one [`Writer`] behind one lock, and until this every
4//! instance encoded its stripe while holding that lock. On the 32 core box, loading the ClickBench
5//! 10m sample spent 68% of all its processor time in the stripe encode, all of it under the lock,
6//! and the instances waited 146 seconds between them for a load whose wall clock was 20.6 seconds.
7//! The machine was one encode at a time with thirty one readers queued behind it.
8//!
9//! Almost none of that work needs the writer. A plain column's pages, its sieves, its ranges and
10//! its statistics depend on the stripe's own rows and nothing else. The one thing a stripe shares
11//! with the rest of the table is a varchar column's global dictionary, because a code has to mean
12//! the same value in every page of the column. So a stripe is taken in four steps:
13//!
14//! 1. [`Preparer::prepare`], with no lock. Every column without a global dictionary is encoded to
15//!    its pages, and every column with one is coded against a dictionary of the stripe's own, which
16//!    holds each distinct value of the stripe once, in the order the rows first held it. The
17//!    statistics of every column are folded into a gather of the stripe's own.
18//! 2. [`Writer::merge`], under the lock. The stripe's dictionaries go into the global ones a
19//!    distinct value at a time, which gives back what each local code is globally, and the gathers
20//!    are absorbed. This is the only step that has to see the stripes one at a time.
21//! 3. [`Merged::pages`], with no lock. The codes are turned into global ones and built into pages.
22//! 4. [`Writer::write`], under the lock. The pages go into the file.
23//!
24//! A value merged in the order the stripe first held it gets the code it would have got had the
25//! stripe been coded against the global dictionary row by row, because the rows before its first
26//! appearance hold only values that were already merged. So a writer taking the four steps one
27//! after the other writes the same bytes as one that coded every row against the global dictionary,
28//! which is what [`Writer::flush_pending`] does.
29//!
30//! A stripe lets go of its rows at the end of the first step. What it carries from there on is its
31//! pages, its stripe dictionaries and codes, and a few numbers a part, so the stripes queued for the
32//! lock are a fraction of the size of the rows they came from. A column that loses its dictionary
33//! after it was coded against one is rebuilt from the stripe dictionary, which holds every value
34//! the rows did. Keeping the rows until the write instead took the 10m ClickBench load on the 32
35//! core box from 3.5 GB resident to 9.9 GB, with thirty two stripes waiting at a time.
36
37use 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
53/// How many stripes are being prepared or paged right now, across every writer in the process.
54///
55/// A stripe's columns are spread over threads of their own, which is what a writer being fed by
56/// one caller needs, because that caller is the only one encoding. Thirty two callers each doing
57/// that at once would be a thousand threads on a machine with thirty two cores. So each one takes
58/// its share of the machine: the cores over however many stripes are being worked on right now.
59static BUSY: AtomicUsize = AtomicUsize::new(0);
60
61/// One stripe's share of the machine, held for as long as the stripe is being worked on.
62struct 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        // A stripe of one part is one small page a column, which is less than a thread is worth.
70        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/// Encodes stripes for one [`Writer`] without the writer.
82///
83/// Handed out by [`Writer::preparer`] and cheap to hold. It shares with its writer which varchar
84/// columns still have a global dictionary, so a stripe prepared after the first one decided a
85/// column should not have one is encoded plainly from the start.
86#[derive(Debug, Clone)]
87pub struct Preparer {
88    types: Vec<LogicalType>,
89    coded: Arc<[AtomicBool]>,
90    profile: Option<Arc<LoadProfile>>,
91}
92
93/// A stripe that has been through [`Preparer::prepare`] and is waiting for [`Writer::merge`].
94#[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/// A stripe that has been through [`Writer::merge`] and is waiting for [`Merged::pages`].
104#[derive(Debug)]
105pub struct Merged {
106    parts: Vec<Part>,
107    columns: Vec<Merge>,
108    profile: Option<Arc<LoadProfile>>,
109}
110
111/// A stripe that has been through [`Merged::pages`] and is waiting for [`Writer::write`].
112#[derive(Debug)]
113pub struct Paged {
114    parts: Vec<Part>,
115    columns: Vec<ColumnStripe>,
116}
117
118/// One column of a prepared stripe.
119#[derive(Debug)]
120enum Column {
121    /// Finished, because the column has no global dictionary.
122    Pages(ColumnStripe),
123    /// Coded against the stripe's own dictionary, waiting to be merged into the global one.
124    Coded(Local),
125}
126
127/// One column of a merged stripe.
128#[derive(Debug)]
129enum Merge {
130    Pages(ColumnStripe),
131    /// The local codes of every part, and the global code of every local one.
132    Codes {
133        parts: Vec<LocalPart>,
134        global: Vec<u32>,
135    },
136    /// A column that was prepared against a dictionary it no longer has, which is every column
137    /// prepared before the first stripe decided it should not have one. Encoded again, plainly,
138    /// from the values its stripe dictionary holds.
139    Plain(Local),
140}
141
142/// No value after this one has its hash.
143const END: u32 = u32::MAX;
144
145/// A dictionary of one column of one stripe.
146///
147/// The values are compared by their bytes rather than by a second hash, because they are all here
148/// to compare. The global dictionary has two hashes to go on because its values are mostly in the
149/// file by now. Both hashes are taken here, once a distinct value, so that merging it takes none.
150#[derive(Debug, Default)]
151struct Local {
152    /// The first value holding each hash.
153    first: HashMap<u64, u32, Spread>,
154    /// The next value holding the same hash as this one, or [`END`].
155    next: Vec<u32>,
156    hashes: Vec<u64>,
157    checks: Vec<u64>,
158    /// The values back to back, and where each one ends.
159    bytes: Vec<u8>,
160    ends: Vec<usize>,
161    /// How many rows that are not null hold each value, and how many are null.
162    counts: Vec<u64>,
163    nulls: u64,
164    parts: Vec<LocalPart>,
165}
166
167/// One part of one column coded against its stripe's dictionary.
168#[derive(Debug)]
169struct LocalPart {
170    codes: Vec<u32>,
171    /// What [`push_validity`] wrote for the part, which is the page's second field onwards.
172    validity: Vec<u8>,
173    range: Range,
174}
175
176impl Local {
177    /// One column of a stripe, coded.
178    ///
179    /// A null row is coded as the empty string and counted as a null rather than against it, which
180    /// is what the writer has always done with one. The code is never read, since the page's
181    /// validity says the row is null, and giving it one keeps the page one code a row.
182    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            // flatten: the page is one code a row whatever form the rows came in.
187            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                // A repeat of the row before is common enough on a sorted table to be worth a
193                // comparison before a hash, and the comparison fails on its first bytes when not.
194                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        // Only the coding needs to find a value by its bytes, and on a column of URLs the table
211        // that does it is as large as the codes.
212        local.first = HashMap::default();
213        local.next = Vec::new();
214        Ok(local)
215    }
216
217    /// The column's parts as the rows they were coded from, for a column that lost its global
218    /// dictionary after this stripe was coded against one.
219    ///
220    /// A null row comes back as a null over the empty string, which is what it was coded as, and
221    /// each part gets back the same form of validity it had, since the page records which it was.
222    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    /// Puts every value into `dictionary` in the order this stripe first held it, and says what
296    /// each one's code is there.
297    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
314/// Whether a column's first stripe says it should not have a global dictionary.
315///
316/// Every varchar column starts with one, because the writer cannot know what is in a column before
317/// it has seen some of it. A global dictionary is the right shape for a column of a few dozen
318/// values repeated down the table: the pages become small integers, a filter against a literal is
319/// one search of the sorted order rather than a comparison a row, and a group by is on the codes.
320/// It is the wrong shape for a column whose values are nearly all different. There the codes are
321/// as wide as row numbers, nothing is saved on the pages, and the membership index of a stripe is
322/// a list of very nearly every code in the column. On TPC-H the orders table written on its own
323/// goes from 52.3 MB to 41.4 MB, the load from 6.9 s to 5.8 s, and `select o_comment from orders`
324/// from 1.810 G instructions to 1.213 G, which is what the rudb parquet reader takes over the same
325/// values.
326///
327/// So the first stripe of a column is the sample and the decision is made once on it. Once, rather
328/// than per stripe, because the codes of one column have to mean the same thing in every page of
329/// it, and a column that changed its mind halfway would need its earlier stripes rewritten. The
330/// first stripe is encoded again when the answer comes out against the dictionary, which is the one
331/// stripe that pays for the decision, along with any stripe that was prepared before it was made.
332///
333/// The threshold is deliberately near the top. [`DICTIONARY_DISTINCT_IN_TEN`] of the sample has to
334/// be values never seen before, which is a column with essentially no repeats. Everything with real
335/// repetition keeps its dictionary and keeps every property that hangs off it, and nothing is
336/// claimed here about where between the two the crossover really sits.
337fn 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
342/// Runs `work` on every one of `jobs`, spread over `workers` threads, and hands back each job with
343/// what it came to, in no particular order.
344///
345/// The jobs are handed out through a queue rather than dealt in equal piles, because they are
346/// nothing like equal: `URL` on ClickBench is a string column of sixty one million distinct values
347/// and `IsMobile` is a byte. A pile that happened to hold the four large string columns would be
348/// the whole stripe and the other workers would be waiting on it. The caller hands the jobs over
349/// cheapest first and they are taken from the back, so the expensive ones go first, which is the
350/// classic answer to a last job that runs longer than everything before it.
351fn 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
390/// One column of every part of a stripe.
391fn column_of(held: &[PendingChunk], index: usize) -> Result<Vec<&Vector>> {
392    held.iter().map(|pending| pending.chunk.column(index)).collect()
393}
394
395/// Puts what [`fan_out`] handed back in column order.
396fn 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    /// Encodes a run of chunks as one stripe, as far as it can be without the writer.
409    ///
410    /// The run is what [`Writer::append_stripe`] takes, and the rules are the same: it is a stripe
411    /// of its own, and the orders have to come out in source order once the stripes are sorted. An
412    /// empty chunk is dropped.
413    ///
414    /// # Errors
415    ///
416    /// If the run is longer than [`STRIPE_PARTS`], a chunk's columns are not the table's, or one
417    /// cannot be encoded.
418    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    /// The check [`Writer::admit`] makes, here because the rows are gone by the merge.
434    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            // The statistics on the thread that is already walking the column, and in the same
454            // step, because the rows are in memory once and this is the moment they are.
455            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    /// Builds the pages the merge left to build, which is every column coded against a global
487    /// dictionary and every column that lost one after the stripe was prepared.
488    ///
489    /// # Errors
490    ///
491    /// If a column cannot be encoded or a page comes out larger than a page may be.
492    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        // A column encoded again from its rows costs more than one whose codes only need building.
499        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
523/// One column's parts as pages of global codes.
524fn 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        // None, because the codes already give the stripe an exact membership index, and an
545        // approximate one beside it would cost a hash of every string to answer a question that
546        // is already answered.
547        stripe.sieves.push(None);
548        stripe.ranges.push(part.range.clone());
549    }
550    Ok(stripe)
551}
552
553impl Writer {
554    /// Something that encodes stripes for this writer without holding it. See [`Preparer::prepare`].
555    ///
556    /// It carries the profile the writer has when it is asked for, so a writer that is going to be
557    /// given one with [`Writer::with_profile`] should be given it first.
558    #[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    /// Takes a prepared stripe into the table's dictionaries and statistics and counts its rows in.
568    ///
569    /// This is the step that has to see the stripes one at a time, and it is a hash a distinct
570    /// value of each varchar column rather than two a row. Whatever [`Writer::append_at`] left
571    /// behind is written first as its own stripe, the same rule [`Writer::append_stripe`] has.
572    ///
573    /// # Errors
574    ///
575    /// If the stripe was prepared for a table of other columns, or the buffered stripe cannot be
576    /// written.
577    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    /// [`Writer::merge`] for a stripe whose rows are already counted in.
593    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                    // Empty means nothing has been merged into it yet, so this is the column's
615                    // first stripe and the only one the decision is allowed to be made on.
616                    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    /// Writes a stripe whose pages are built.
632    ///
633    /// # Errors
634    ///
635    /// If the stripe was built for a table of another width or cannot be written.
636    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    /// All four steps one after the other, for a caller with nobody to share the writer with.
648    ///
649    /// # Errors
650    ///
651    /// The same as [`Writer::merge`], [`Merged::pages`] and [`Writer::write`].
652    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    /// The value every row holds, so a test can check a row it reads back without keeping the rows.
688    ///
689    /// `city` repeats a handful of values and has a null every so often, which keeps its dictionary.
690    /// `note` is different on every row but its nulls, which loses it on the first stripe.
691    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    /// A run of `parts` chunks starting at part `first`, as a caller hands them to the writer.
702    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    /// The runs the tests hand over, out of source order so that the stripes are sorted at commit.
717    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    /// Every stripe prepared before any of them is merged writes the file that handing the same
736    /// runs to the writer one at a time writes, byte for byte.
737    ///
738    /// That is the claim the whole split rests on. The second and third stripes here are coded
739    /// against a dictionary for `note`, which the first stripe to be merged then decides the column
740    /// should not have, so they are encoded again without it. `city` keeps its dictionary and the
741    /// later stripes' values go into it in the order the merges happen.
742    #[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    /// Two stripes merged in one order and written in the other read back as the rows they held,
772    /// which is what two instances sharing a writer do whenever the second one's pages are built
773    /// first.
774    #[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    /// A chunk that is not the table's is refused when it reaches the writer, and the writer is not
794    /// left counting its rows.
795    #[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}