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. Every
21//!    dictionary block the merge filled is taken out with the stripe.
22//! 3. [`Merged::pages`], with no lock. The codes are turned into global ones and built into pages,
23//!    and the dictionary blocks the merge took out are encoded.
24//! 4. [`Writer::write`], under the lock. The dictionary blocks go back in order, and they and the
25//!    pages go into the file.
26//!
27//! The dictionary blocks were encoded in the fourth step, under the lock, until the 10m ClickBench
28//! load on the 32 core box was measured spending 2.7 of its 14 seconds there, on thirty two threads
29//! spawned for it every stripe, with every instance queued behind them.
30//!
31//! A value merged in the order the stripe first held it gets the code it would have got had the
32//! stripe been coded against the global dictionary row by row, because the rows before its first
33//! appearance hold only values that were already merged. So a writer taking the four steps one
34//! after the other writes the same bytes as one that coded every row against the global dictionary,
35//! which is what [`Writer::flush_pending`] does.
36//!
37//! A stripe lets go of its rows at the end of the first step. What it carries from there on is its
38//! pages, its stripe dictionaries and codes, and a few numbers a part, so the stripes queued for the
39//! lock are a fraction of the size of the rows they came from. A column that loses its dictionary
40//! after it was coded against one is rebuilt from the stripe dictionary, which holds every value
41//! the rows did. Keeping the rows until the write instead took the 10m ClickBench load on the 32
42//! core box from 3.5 GB resident to 9.9 GB, with thirty two stripes waiting at a time.
43
44use 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
62/// How many stripes are being prepared or paged right now, across every writer in the process.
63///
64/// A stripe's columns are spread over threads of their own, which is what a writer being fed by
65/// one caller needs, because that caller is the only one encoding. Thirty two callers each doing
66/// that at once would be a thousand threads on a machine with thirty two cores. So each one takes
67/// its share of the machine: the cores over however many stripes are being worked on right now.
68static BUSY: AtomicUsize = AtomicUsize::new(0);
69
70/// What all of a table's global dictionaries may hold at once before the fastest growing one is
71/// demoted, from section 5.5 of the encoding spec.
72///
73/// Dictionaries are the one thing a load holds that grows with the table rather than with the
74/// stripe. `hits` has tens of millions of distinct `URL`, `Title`, `Referer` and `SearchPhrase`
75/// values, at forty five to seventy bytes each for the lookup alone, which is past the whole two
76/// gigabyte bound on its own.
77pub const DICTIONARY_CAP_BYTES: u64 = 512 * 1024 * 1024;
78
79/// Which varchar columns still code against a global dictionary, and what their dictionaries hold
80/// between them.
81///
82/// Shared by a writer with every [`Preparer`] and [`Merger`] it hands out. The flags are what a
83/// stripe prepared later reads to decide whether to code a column at all, and the rest is what a
84/// merge reads to decide whether its column should stop, see [`demotes`].
85#[derive(Debug)]
86pub(crate) struct Coding {
87    flags: Box<[AtomicBool]>,
88    /// What each column's dictionary grew by in the last stripe merged into it.
89    growth: Box<[AtomicU64]>,
90    /// What every dictionary held the last time it was merged into, added up.
91    held: AtomicU64,
92    cap: AtomicU64,
93    /// The lowest distinct count ceiling a stripe of each column has reported, `u64::MAX` until one
94    /// has.
95    ///
96    /// Every stripe prepared for a writer is merged into the writer's statistics, so a stripe's full
97    /// sketch is one of the parts of the union, and no hash above its ceiling can be in the union's
98    /// bottom k. A stripe started later drops those hashes rather than hashing them into a table
99    /// that the union would trim them from anyway. Without it, every stripe filled a sketch from
100    /// empty, and `Sketch::insert` was about 2% of a `lineitem` load's samples.
101    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    /// Sets what the dictionaries may hold between them before one is demoted.
119    pub(crate) fn cap(&self, bytes: u64) {
120        self.cap.store(bytes, Atomic::Relaxed);
121    }
122
123    /// Moves what one dictionary is counted for from `before` to `now`, and hands back the new
124    /// total.
125    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    /// Whether the column grew the most in the stripes last merged into each column.
134    ///
135    /// Read without a lock across columns that may be merging at the same moment, so it is about
136    /// the last stripe or the one before it. That is close enough for choosing which column to
137    /// stop: a column that grows fastest keeps doing so, and it is asked again next stripe. A column
138    /// that did not grow at all is never the one, since stopping it would free nothing.
139    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
153/// One stripe's share of the machine, held for as long as the stripe is being worked on.
154struct 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        // A stripe of one part is one small page a column, which is less than a thread is worth.
162        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/// Encodes stripes for one [`Writer`] without the writer.
174///
175/// Handed out by [`Writer::preparer`] and cheap to hold. It shares with its writer which varchar
176/// columns still have a global dictionary, so a stripe prepared after the first one decided a
177/// column should not have one is encoded plainly from the start.
178#[derive(Debug, Clone)]
179pub struct Preparer {
180    types: Vec<LogicalType>,
181    coded: Arc<Coding>,
182    profile: Option<Arc<LoadProfile>>,
183}
184
185/// A stripe that has been through [`Preparer::prepare`] and is waiting for [`Writer::merge`].
186#[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/// A stripe that has been through [`Writer::merge`] and is waiting for [`Merged::pages`].
196#[derive(Debug)]
197pub struct Merged {
198    parts: Vec<Part>,
199    columns: Vec<Merge>,
200    blocks: Vec<Unencoded>,
201    profile: Option<Arc<LoadProfile>>,
202    /// Whether the rows are counted into the table yet. [`Writer::merge`] counts them, and a
203    /// [`Merger`] leaves them for [`Writer::write`], since it has no table to count them into.
204    counted: bool,
205}
206
207/// A stripe that has been through [`Merged::pages`] and is waiting for [`Writer::write`].
208#[derive(Debug)]
209pub struct Paged {
210    parts: Vec<Part>,
211    columns: Vec<ColumnStripe>,
212    /// Encoded dictionary blocks, each with its column and block number.
213    blocks: Vec<(usize, usize, EncodedBlock)>,
214    counted: bool,
215}
216
217/// What one job of [`Merged::pages`] built.
218enum Built {
219    Stripe(ColumnStripe),
220    Block(EncodedBlock),
221}
222
223/// One column of a prepared stripe.
224#[derive(Debug)]
225enum Column {
226    /// Finished, because the column has no global dictionary.
227    Pages(ColumnStripe),
228    /// Coded against the stripe's own dictionary, waiting to be merged into the global one.
229    Coded(Local),
230}
231
232/// One column of a merged stripe.
233#[derive(Debug)]
234enum Merge {
235    Pages(ColumnStripe),
236    /// The local codes of every part, and the global code of every local one.
237    Codes {
238        parts: Vec<LocalPart>,
239        global: Vec<u32>,
240    },
241    /// A column that was prepared against a dictionary it no longer has, which is every column
242    /// prepared before the first stripe decided it should not have one. Encoded again, plainly,
243    /// from the values its stripe dictionary holds.
244    Plain(Local),
245}
246
247/// No value after this one has its hash.
248const END: u32 = u32::MAX;
249
250/// A dictionary of one column of one stripe.
251///
252/// The values are compared by their bytes rather than by a second hash, because they are all here
253/// to compare. The global dictionary has two hashes to go on because its values are mostly in the
254/// file by now. Both hashes are taken here, once a distinct value, so that merging it takes none.
255#[derive(Debug, Default)]
256struct Local {
257    /// The first value holding each hash.
258    first: HashMap<u64, u32, Spread>,
259    /// The next value holding the same hash as this one, or [`END`].
260    next: Vec<u32>,
261    hashes: Vec<u64>,
262    checks: Vec<u64>,
263    /// The values back to back, and where each one ends.
264    bytes: Vec<u8>,
265    ends: Vec<usize>,
266    /// How many rows that are not null hold each value, and how many are null.
267    counts: Vec<u64>,
268    nulls: u64,
269    parts: Vec<LocalPart>,
270    /// Whether the column is a blob rather than a varchar, for building its rows back.
271    blob: bool,
272}
273
274/// One part of one column coded against its stripe's dictionary.
275#[derive(Debug)]
276struct LocalPart {
277    codes: Vec<u32>,
278    /// What [`push_validity`] wrote for the part, which is the page's second field onwards.
279    validity: Vec<u8>,
280    range: Range,
281}
282
283impl Local {
284    /// The bytes the dictionary and the codes of the parts so far take up.
285    fn held(&self) -> usize {
286        // A hash table slot is its key, its value and one control byte.
287        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    /// One column of a stripe, coded in one go.
303    #[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    /// One more part of a column of a stripe, coded.
315    ///
316    /// A null row is coded as the empty string and counted as a null rather than against it, which
317    /// is what the writer has always done with one. The code is never read, since the page's
318    /// validity says the row is null, and giving it one keeps the page one code a row.
319    ///
320    /// The parts come to [`Local::code_part`] one at a time, in order, and [`Local::done`] ends the
321    /// column, so a stripe can be coded as its parts arrive rather than once they are all held.
322    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        // flatten: the page is one code a row whatever form the rows came in.
335        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            // bytes_at rather than text_at: the rows were checked for UTF-8 when they came in,
340            // and checking every one again cost more than coding it.
341            let text = flat.bytes_at(row).unwrap_or(b"");
342            // A repeat of the row before is common enough on a sorted table to be worth a
343            // comparison before a hash, and the comparison fails on its first bytes when not.
344            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    /// Ends a column once its last part is coded.
363    fn done(&mut self) {
364        // Only the coding needs to find a value by its bytes, and on a column of URLs the table
365        // that does it is as large as the codes.
366        self.first = HashMap::default();
367        self.next = Vec::new();
368    }
369
370    /// Codes a part that came in as codes into a dictionary of its own, which is how a Parquet page
371    /// written with a dictionary arrives, by coding each value of that dictionary once rather than
372    /// each row.
373    ///
374    /// Every row then costs a lookup in `mapped`, which holds the local code of each value of the
375    /// last dictionary seen by its position in it, and a value is coded the first time a row holds
376    /// it. So the codes come out in the order the rows first held each value, the same as coding the
377    /// rows one at a time, and the stripe is the same stripe either way. The parts of one Parquet
378    /// column chunk share their dictionary, so `mapped` carries over from one part to the next
379    /// while it is the same one, found by the pointer and kept alive by holding it.
380    ///
381    /// `None` for anything else, and for a dictionary with a null in it, since a null row there is
382    /// found through the value it points at and not through the part's own validity, which is the
383    /// one [`push_validity`] writes.
384    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                // Coded as the empty string and counted as a null, as `code_part` does.
404                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    /// The column's parts as the rows they were coded from, for a column that lost its global
422    /// dictionary after this stripe was coded against one.
423    ///
424    /// A null row comes back as a null over the empty string, which is what it was coded as, and
425    /// each part gets back the same form of validity it had, since the page records which it was.
426    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    /// Puts every value into `dictionary` in the order this stripe first held it, and says what
500    /// each one's code is there.
501    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
518/// Whether a column's first stripe says it should not have a global dictionary.
519///
520/// Every varchar column starts with one, because the writer cannot know what is in a column before
521/// it has seen some of it. A global dictionary is the right shape for a column of a few dozen
522/// values repeated down the table: the pages become small integers, a filter against a literal is
523/// one search of the sorted order rather than a comparison a row, and a group by is on the codes.
524/// It is the wrong shape for a column whose values are nearly all different. There the codes are
525/// as wide as row numbers, nothing is saved on the pages, and the membership index of a stripe is
526/// a list of very nearly every code in the column. On TPC-H the orders table written on its own
527/// 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`
528/// from 1.810 G instructions to 1.213 G, which is what the rudb parquet reader takes over the same
529/// values.
530///
531/// So the first stripe of a column is the sample and the decision is made once on it. Once, rather
532/// than per stripe, because the codes of one column have to mean the same thing in every page of
533/// it, and a column that changed its mind halfway would need its earlier stripes rewritten. The
534/// first stripe is encoded again when the answer comes out against the dictionary, which is the one
535/// stripe that pays for the decision, along with any stripe that was prepared before it was made.
536///
537/// The threshold is deliberately near the top. [`DICTIONARY_DISTINCT_IN_TEN`] of the sample has to
538/// be values never seen before, which is a column with essentially no repeats. Everything with real
539/// repetition keeps its dictionary and keeps every property that hangs off it, and nothing is
540/// claimed here about where between the two the crossover really sits.
541fn 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
546/// Whether a column's dictionary should stop taking values after a stripe that added `new` of them
547/// in `rows` rows, with the dictionaries holding `total` bytes between them.
548///
549/// Section 5.5 of the encoding spec, which asks at every stripe what [`drops_dictionary`] asks at
550/// the first. A column that turns into a column of new values partway through, which is what a
551/// URL column of a log does once its first hours are past, stops growing its dictionary one stripe
552/// after it turns rather than at the end of the load. The stripes already coded keep their codes,
553/// so unlike the first stripe's decision this one costs nothing to make late.
554///
555/// The second reason is the cap. Once the dictionaries together hold more than it, the column that
556/// grew the most in its last stripe is the one that stops, because it is the one that would have
557/// taken the most of what is left.
558fn 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
563/// Runs `work` on every one of `jobs`, spread over `workers` threads, and hands back each job with
564/// what it came to, in no particular order.
565///
566/// The jobs are handed out through a queue rather than dealt in equal piles, because they are
567/// nothing like equal: `URL` on ClickBench is a string column of sixty one million distinct values
568/// and `IsMobile` is a byte. A pile that happened to hold the four large string columns would be
569/// the whole stripe and the other workers would be waiting on it. The caller hands the jobs over
570/// cheapest first and they are taken from the back, so the expensive ones go first, which is the
571/// classic answer to a last job that runs longer than everything before it.
572fn 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    /// Encodes a run of chunks as one stripe, as far as it can be without the writer.
613    ///
614    /// The run is what [`Writer::append_stripe`] takes, and the rules are the same: it is a stripe
615    /// of its own, and the orders have to come out in source order once the stripes are sorted. An
616    /// empty chunk is dropped.
617    ///
618    /// # Errors
619    ///
620    /// If the run is longer than [`STRIPE_PARTS`], a chunk's columns are not the table's, or one
621    /// cannot be encoded.
622    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    /// The check [`Writer::admit`] makes, here because the rows are gone by the merge.
638    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    /// Starts a stripe that will be handed over a few parts at a time.
657    ///
658    /// [`Preparer::prepare`] takes a stripe whole, which means a caller holds every part of it
659    /// decoded until the last one arrives, and every load instance holds one. Here each batch is
660    /// encoded as it comes and let go, so what an instance holds is one batch of rows and what it
661    /// has built from the batches before. The stripe comes out the same: the parts of a column go through
662    /// the same steps in the same order, only with the rows of later parts not yet in memory.
663    ///
664    /// Whether a column is coded against its global dictionary is read here, once, so every part of
665    /// the stripe is encoded the same way even if the column is demoted while it is being built.
666    /// A stripe that ends up coded against a dictionary its column no longer has is encoded again
667    /// plainly at the merge, which is what happens to a stripe prepared whole at the same moment.
668    #[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    /// Encodes the next few parts of a stripe that was started with [`Preparer::start`].
689    ///
690    /// The same rules as [`Preparer::prepare`]: the orders come in source order, an empty chunk is
691    /// dropped, and the stripe holds no more than [`STRIPE_PARTS`] in all.
692    ///
693    /// # Errors
694    ///
695    /// If the stripe would hold too many parts, a chunk's columns are not the table's, or one cannot
696    /// be encoded.
697    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        // The stripe's key is its first part's order, and its statistics open with it.
718        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                // The statistics on the thread that is already walking the column, and in the same
734                // step, because the rows are in memory once and this is the moment they are.
735                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    /// Ends a stripe that was started with [`Preparer::start`], ready for the merge.
762    ///
763    /// # Errors
764    ///
765    /// If a column's worker panicked.
766    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                // A stripe of no parts has no statistics, rather than an empty stripe of them.
784                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/// A stripe that [`Preparer::start`] began and [`Preparer::feed`] is adding parts to.
807#[derive(Debug)]
808pub struct Building {
809    parts: Vec<Part>,
810    /// Behind a lock each so a column can be handed to whichever worker takes it. Only one does
811    /// at a time, so the locks are never waited on.
812    columns: Vec<Mutex<Growing>>,
813}
814
815impl Building {
816    /// How many parts the stripe holds so far.
817    #[must_use]
818    pub fn parts(&self) -> usize {
819        self.parts.len()
820    }
821
822    /// The bytes the stripe holds so far: every column's codes, stripe dictionary and pages.
823    ///
824    /// The rows a load hands in are charged until they are encoded, and this is what they turn
825    /// into. It is not the small fraction of them it sounds like. A text column coded against its
826    /// stripe dictionary keeps four bytes of code a row and every distinct value with its hashes,
827    /// and each worker of a load has a stripe of its own going.
828    #[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
846/// The bytes a column's finished pages and what goes beside them take up.
847fn 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
858/// The bytes a vector's buffer takes up, whatever is in it.
859fn spilled<T>(values: &Vec<T>) -> usize {
860    values.capacity() * size_of::<T>()
861}
862
863/// One column of a stripe that is being built.
864#[derive(Debug)]
865struct Growing {
866    body: Body,
867    gather: Option<stats::Gather>,
868}
869
870/// What one column of a stripe being built has come to so far.
871#[derive(Debug)]
872enum Body {
873    /// Pages, with what the parts so far have settled on.
874    Pages(ColumnStripe, Settling),
875    /// Codes against the stripe's own dictionary, with the Parquet dictionary the last part came
876    /// in, as [`Local::code_dictionary`] keeps it.
877    Coded(Local, Option<(Arc<Vector>, Vec<u32>)>),
878}
879
880impl Body {
881    /// Turns a column coded against its dictionary into pages, for a column that lost its
882    /// dictionary while the stripe was being built.
883    ///
884    /// The merge would encode the whole stripe again plainly once it saw the column had none. Doing
885    /// it here, on the parts coded so far, means the parts still to come are encoded once. A load
886    /// starts every stripe it has in flight before the first one reaches the merge and decides.
887    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
899/// Where one column's dictionary and statistics are while a stripe is merged into them.
900enum Slot<'a> {
901    /// In the writer, which the caller holds.
902    Owned(&'a mut Option<GlobalDictionary>, &'a mut Option<stats::Gather>),
903    /// Lent to a [`Merger`], behind the column's own lock.
904    Lent(&'a Mutex<LentColumn>, &'a Lent),
905}
906
907/// One column of a stripe on its way through [`merge_columns`].
908struct Step<'a> {
909    index: usize,
910    column: Column,
911    slot: Slot<'a>,
912    /// The stripe's statistics for the column, when the column keeps them.
913    gather: Option<stats::Gather>,
914}
915
916impl Step<'_> {
917    /// Roughly what the merge costs: a hash a distinct value when there is a global dictionary to
918    /// merge into, and next to nothing otherwise. Read off `coded` rather than the dictionary, so a
919    /// lent column does not have to be locked to be sorted.
920    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    /// Merges the column, settles its dictionary's shape and hands out the blocks it filled.
930    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                // Checked with the column locked, so a merge either finishes before the writer
944                // takes this column back or is refused.
945                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/// One column of [`merge_columns`].
956#[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        // Prepared before the column was demoted and merged after.
981        (Column::Coded(local), Some(global)) if global.demoted => Merge::Plain(local),
982        (Column::Coded(local), Some(global)) => {
983            // Empty means nothing has been merged into it yet, so this is the column's first
984            // stripe and the only one the decision is allowed to be made on.
985            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    // Asked before the blocks go out, so that a demotion's sealed part block goes out with them.
1002    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    // Settled here rather than when the stripe is written, so that the blocks this merge filled go
1014    // out with it already knowing their shape. A column still too small to settle one keeps its
1015    // blocks until it can, which is at most `PAYLOAD_SAMPLE_BLOCKS` of them, because encoding them
1016    // now would be encoding them without having looked at the column.
1017    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
1030/// Merges every column of a stripe into the dictionaries and statistics in `slots`.
1031///
1032/// Every column is merged on its own, because nothing one column's merge reads or writes belongs to
1033/// another: its statistics, its global dictionary and its flag in `coded`. So the columns are
1034/// spread over threads, and a stripe takes as long as its slowest column rather than all of them.
1035/// The answer is the same in any order, because a column's merge only depends on the stripes
1036/// merged into that column before it.
1037fn 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    // Taken from the back, so the biggest merges start first and the last one to finish is
1053    // small, the same reason `fan_out` hands its jobs over cheapest first.
1054    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/// The dictionaries and statistics of a table while a [`Merger`] has them, one lock a column.
1108#[derive(Debug)]
1109pub(crate) struct Lent {
1110    columns: Box<[Mutex<LentColumn>]>,
1111    /// Set when the writer takes them back, after which a merge is refused.
1112    reclaimed: AtomicBool,
1113}
1114
1115/// One column of [`Lent`].
1116#[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    /// Puts encoded blocks back into their dictionaries, each under its own column's lock.
1128    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    /// Everything lent, handed back to the writer.
1144    #[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/// Merges prepared stripes into a writer's dictionaries and statistics without the writer.
1161///
1162/// Handed out by [`Writer::merger`]. With it, a load that shares one writer between many threads
1163/// holds the writer's lock only to write, and two stripes merge at once as long as they are on
1164/// different columns. A stripe merged here is written with [`Writer::write`] as usual, and that is
1165/// where its rows are counted in.
1166#[derive(Debug, Clone)]
1167pub struct Merger {
1168    lent: Arc<Lent>,
1169    types: Vec<LogicalType>,
1170    coded: Arc<Coding>,
1171}
1172
1173impl Merger {
1174    /// [`Writer::merge`], one column lock at a time instead of the writer.
1175    ///
1176    /// # Errors
1177    ///
1178    /// If the stripe was prepared for a table of other columns, or the table was closed.
1179    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    /// Puts a stripe's encoded dictionary blocks back, so that [`Writer::write`] does not wait on a
1188    /// column's lock while it holds its own.
1189    ///
1190    /// # Errors
1191    ///
1192    /// If a block comes back to a column without a dictionary, or comes back twice.
1193    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    /// Builds the pages the merge left to build, which is every column coded against a global
1200    /// dictionary and every column that lost one after the stripe was prepared, and encodes the
1201    /// dictionary blocks the merge filled.
1202    ///
1203    /// # Errors
1204    ///
1205    /// If a column or a block cannot be encoded or a page comes out larger than a page may be.
1206    pub fn pages(self) -> Result<Paged> {
1207        let Self { parts, columns, blocks, profile, counted } = self;
1208        let width = columns.len();
1209        // The blocks go first so that they are taken last. One block is a thousand values, which is
1210        // less than any column of a stripe, and small jobs at the end are what keeps the last
1211        // worker from finishing long after the others.
1212        let mut jobs = (width..width + blocks.len())
1213            .chain((0..width).filter(|&index| !matches!(columns[index], Merge::Pages(_))))
1214            .collect::<Vec<_>>();
1215        // A column encoded again from its rows costs more than one whose codes only need building.
1216        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
1256/// One column's parts as pages of global codes.
1257fn 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        // None, because the codes already give the stripe an exact membership index, and an
1280        // approximate one beside it would cost a hash of every string to answer a question that
1281        // is already answered.
1282        stripe.sieves.push(None);
1283        stripe.ranges.push(part.range.clone());
1284    }
1285    Ok(stripe)
1286}
1287
1288impl Writer {
1289    /// Something that encodes stripes for this writer without holding it. See [`Preparer::prepare`].
1290    ///
1291    /// It carries the profile the writer has when it is asked for, so a writer that is going to be
1292    /// given one with [`Writer::with_profile`] should be given it first.
1293    #[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    /// Takes a prepared stripe into the table's dictionaries and statistics and counts its rows in.
1303    ///
1304    /// This is the step that has to see the stripes one at a time, and it is a hash a distinct
1305    /// value of each varchar column rather than two a row. Whatever [`Writer::append_at`] left
1306    /// behind is written first as its own stripe, the same rule [`Writer::append_stripe`] has.
1307    ///
1308    /// # Errors
1309    ///
1310    /// If the stripe was prepared for a table of other columns, or the buffered stripe cannot be
1311    /// written.
1312    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    /// [`Writer::merge`] for a stripe whose rows are already counted in.
1328    ///
1329    /// Every column is merged on its own, because nothing one column's merge reads or writes
1330    /// belongs to another: its statistics, its global dictionary and its flag in `coded`. So the
1331    /// columns are spread over threads, and the lock is held for the slowest column rather than for
1332    /// all of them. On ClickBench `hits` the lock was busy 98% of a load and the merge was four
1333    /// fifths of that, while two thirds of the machine waited for it. The answer is the same in any
1334    /// order, because a column's merge only depends on the stripes merged into it before.
1335    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    /// Hands the dictionaries and the statistics to a [`Merger`], so that stripes can be merged
1351    /// without this writer's lock.
1352    ///
1353    /// Whatever [`Writer::append_at`] left behind is written first, the same rule
1354    /// [`Writer::merge`] has. The writer takes them back when the table is closed.
1355    ///
1356    /// # Errors
1357    ///
1358    /// If the buffered stripe cannot be written.
1359    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    /// Writes a stripe whose pages are built.
1384    ///
1385    /// # Errors
1386    ///
1387    /// If the stripe was built for a table of another width or cannot be written.
1388    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    /// All four steps one after the other, for a caller with nobody to share the writer with.
1420    ///
1421    /// # Errors
1422    ///
1423    /// The same as [`Writer::merge`], [`Merged::pages`] and [`Writer::write`].
1424    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    /// A stripe that came in as codes into Parquet style dictionaries codes to the same stripe as
1446    /// the same rows would flat: the same values in the same order, the same codes, counts, nulls
1447    /// and validity. Two of the parts share a dictionary and one has one of its own, and one of the
1448    /// shared ones has nulls pointing at a value that is not the empty string. None of them is
1449    /// flattened on the way.
1450    #[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    /// The value every row holds, so a test can check a row it reads back without keeping the rows.
1531    ///
1532    /// `city` repeats a handful of values and has a null every so often, which keeps its dictionary.
1533    /// `note` is different on every row but its nulls, which loses it on the first stripe.
1534    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    /// A run of `parts` chunks starting at part `first`, as a caller hands them to the writer.
1546    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    /// The runs the tests hand over, out of source order so that the stripes are sorted at commit.
1561    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    /// Every stripe prepared before any of them is merged writes the file that handing the same
1580    /// runs to the writer one at a time writes, byte for byte.
1581    ///
1582    /// That is the claim the whole split rests on. The second and third stripes here are coded
1583    /// against a dictionary for `note`, which the first stripe to be merged then decides the column
1584    /// should not have, so they are encoded again without it. `city` keeps its dictionary and the
1585    /// later stripes' values go into it in the order the merges happen.
1586    #[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    /// Stripes counted under the ceiling an earlier stripe reported come to the same distinct count
1616    /// as stripes counted from nothing, whatever order they are merged in.
1617    ///
1618    /// The file cannot say this on its own, because a table this small may not be given the budget
1619    /// for its sketches, so the writer's statistics are compared before it closes.
1620    #[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        // What one gather makes of every row, with no stripe and so no ceiling anywhere.
1630        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    /// Two stripes merged in one order and written in the other read back as the rows they held,
1666    /// which is what two instances sharing a writer do whenever the second one's pages are built
1667    /// first.
1668    #[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    /// Stripes merged through a [`Merger`] write the same bytes as the writer merging them itself,
1688    /// and their rows are counted in when they are written.
1689    #[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    /// A stripe fed a few parts at a time, with an empty batch and an empty chunk among them,
1718    /// writes the same bytes as the same stripe prepared whole.
1719    #[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    /// A stripe started while `note` still had its dictionary, which the first stripe's merge then
1763    /// dropped, encodes the rest of `note` as pages and reads back.
1764    #[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    /// A stripe is held to [`STRIPE_PARTS`] across all its batches, not only within one.
1800    #[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    /// Stripes merged on several threads at once through one [`Merger`] and written in whatever
1814    /// order they finish read back as the rows they held.
1815    #[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    /// A merge that comes after the table is closed is refused rather than merged into
1840    /// dictionaries nothing will write.
1841    #[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    /// A chunk that is not the table's is refused when it reaches the writer, and the writer is not
1854    /// left counting its rows.
1855    #[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    /// A table whose `url` repeats twenty values for its first five parts and never repeats after,
1870    /// the way a log's URLs look once its first hours are past, next to a `city` that repeats
1871    /// throughout.
1872    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    /// Reads every row of the turned table back and checks what the reader says about `url`.
1916    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        // Everything the dictionary would have vouched for covers only the first stripes.
1932        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        // `city` keeps its dictionary and all of it.
1938        assert_eq!(reader.distinct_values(1).expect("asked"), Some(13));
1939        assert!(reader.text_extremes(1).expect("asked").is_some());
1940    }
1941
1942    /// A column whose second stripe is nearly all new values stops growing its dictionary there,
1943    /// whether the later stripes were prepared before that decision or after it, and every row
1944    /// reads back.
1945    #[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    /// A table where `url` takes a quarter of its rows as new values every stripe, which keeps its
1979    /// dictionary under the per-stripe rule, and `city` stops growing after its first stripe.
1980    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    /// Once the dictionaries together pass the cap, the column that grew the most stops and the
1986    /// one that did not grow keeps its dictionary.
1987    #[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}