1#![forbid(unsafe_code)]
8
9use std::cmp::Ordering;
10use std::collections::HashMap;
11use std::fs::{File, OpenOptions};
12use std::io::{Read, Seek, SeekFrom, Write};
13use std::mem::size_of;
14use std::path::Path;
15use std::sync::Mutex;
16use std::sync::{Arc, OnceLock};
17
18use rudb_common::bounds::Bound;
19use rudb_common::{Error, Field, LogicalType, Result, Value};
20use rudb_storage::{Probe, Range, Zone};
21use rudb_vector::string::StringColumn;
22use rudb_vector::validity::Validity;
23use rudb_vector::{Buffer, Chunk, Data, TextSource, Vector};
24
25const MAGIC_V7: &[u8; 8] = b"RUDBNV7\0";
26const MAGIC_V8: &[u8; 8] = b"RUDBNV8\0";
27const MAGIC: &[u8; 8] = b"RUDBNV9\0";
28const DIRECTORY_V7: &[u8; 8] = b"RUDBDIR7";
29const DIRECTORY_V8: &[u8; 8] = b"RUDBDIR8";
30const DIRECTORY: &[u8; 8] = b"RUDBDIR9";
31const FORMAT: u32 = 9;
32const HEADER: u64 = 80;
33const SLOT_BYTES: usize = 28;
34const MAX_PAGE: usize = 256 * 1024 * 1024;
35const MAX_DIRECTORY: usize = 128 * 1024 * 1024;
36const FREQUENCIES_V1: &[u8; 8] = b"RUDBFQ1\0";
37const FREQUENCIES: &[u8; 8] = b"RUDBFQ2\0";
38const FREQUENCY_CANDIDATES: usize = 32_768;
39const FREQUENCY_ENTRIES: usize = 512;
40const FREQUENCY_BUILD_RANK: usize = 10;
41const FREQUENCY_ORDINALS: usize = 65_536;
42const MAX_FREQUENCY_WORKERS: usize = 16;
43
44fn io(error: std::io::Error) -> Error {
45 Error::io(error.to_string())
46}
47
48fn invalid(message: &str) -> Error {
49 Error::invalid_input(format!("invalid rudb native file: {message}"))
50}
51
52fn checksum(bytes: &[u8]) -> u64 {
53 const P1: u64 = 11_400_714_785_074_694_791;
54 const P2: u64 = 14_029_467_366_897_019_727;
55 const P3: u64 = 1_609_587_929_392_839_161;
56 const P4: u64 = 9_650_029_242_287_828_579;
57 const P5: u64 = 2_870_177_450_012_600_261;
58 let round = |state: u64, word: u64| {
59 state.wrapping_add(word.wrapping_mul(P2)).rotate_left(31).wrapping_mul(P1)
60 };
61 let merge = |state: u64, lane: u64| (state ^ round(0, lane)).wrapping_mul(P1).wrapping_add(P4);
62 let word =
63 |at: usize| u64::from_le_bytes(bytes[at..at + 8].try_into().expect("eight checksum bytes"));
64
65 let mut at = 0;
66 let mut hash = if bytes.len() >= 32 {
67 let mut one = P1.wrapping_add(P2);
68 let mut two = P2;
69 let mut three = 0;
70 let mut four = 0_u64.wrapping_sub(P1);
71 while at + 32 <= bytes.len() {
72 one = round(one, word(at));
73 two = round(two, word(at + 8));
74 three = round(three, word(at + 16));
75 four = round(four, word(at + 24));
76 at += 32;
77 }
78 let combined = one
79 .rotate_left(1)
80 .wrapping_add(two.rotate_left(7))
81 .wrapping_add(three.rotate_left(12))
82 .wrapping_add(four.rotate_left(18));
83 merge(merge(merge(merge(combined, one), two), three), four)
84 } else {
85 P5
86 };
87 hash = hash.wrapping_add(bytes.len() as u64);
88 while at + 8 <= bytes.len() {
89 hash ^= round(0, word(at));
90 hash = hash.rotate_left(27).wrapping_mul(P1).wrapping_add(P4);
91 at += 8;
92 }
93 if at + 4 <= bytes.len() {
94 let tail = u32::from_le_bytes(bytes[at..at + 4].try_into().expect("four checksum bytes"));
95 hash ^= u64::from(tail).wrapping_mul(P1);
96 hash = hash.rotate_left(23).wrapping_mul(P2).wrapping_add(P3);
97 at += 4;
98 }
99 while at < bytes.len() {
100 hash ^= u64::from(bytes[at]).wrapping_mul(P5);
101 hash = hash.rotate_left(11).wrapping_mul(P1);
102 at += 1;
103 }
104 hash ^= hash >> 33;
105 hash = hash.wrapping_mul(P2);
106 hash ^= hash >> 29;
107 hash = hash.wrapping_mul(P3);
108 hash ^ (hash >> 32)
109}
110
111#[derive(Debug, Clone, Copy)]
112struct Slot {
113 offset: u64,
114 length: u32,
115 generation: u64,
116 hash: u64,
117}
118
119impl Slot {
120 fn bytes(self) -> [u8; SLOT_BYTES] {
121 let mut result = [0; SLOT_BYTES];
122 result[..8].copy_from_slice(&self.offset.to_le_bytes());
123 result[8..12].copy_from_slice(&self.length.to_le_bytes());
124 result[12..20].copy_from_slice(&self.generation.to_le_bytes());
125 result[20..28].copy_from_slice(&self.hash.to_le_bytes());
126 result
127 }
128
129 fn read(bytes: &[u8]) -> Self {
130 Self {
131 offset: u64::from_le_bytes(bytes[..8].try_into().expect("eight bytes")),
132 length: u32::from_le_bytes(bytes[8..12].try_into().expect("four bytes")),
133 generation: u64::from_le_bytes(bytes[12..20].try_into().expect("eight bytes")),
134 hash: u64::from_le_bytes(bytes[20..28].try_into().expect("eight bytes")),
135 }
136 }
137}
138
139#[derive(Debug, Clone, Copy)]
140struct Page {
141 offset: u64,
142 length: u32,
143 hash: u64,
144}
145
146#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
147enum FrequencyValue {
148 Null,
149 Integer(i128),
150 Code(u32),
151}
152
153#[derive(Debug, Clone)]
154struct FrequencyEntry {
155 value: FrequencyValue,
156 count: u64,
157}
158
159#[derive(Debug, Clone)]
164struct FrequencySummary {
165 entries: Vec<FrequencyEntry>,
166 omitted_max: u64,
167 ordinals: Vec<u64>,
168}
169
170#[derive(Debug, Clone, PartialEq, Eq)]
172pub struct FrequencyOccurrences {
173 pub omitted_max: u64,
175 pub ordinals: Vec<u64>,
177}
178
179#[derive(Debug, Clone)]
181pub struct Stripe {
182 rows: usize,
183 pages: Vec<Page>,
184 memberships: Vec<Option<Page>>,
185 zone: Zone,
186}
187
188impl Stripe {
189 #[must_use]
191 pub fn rows(&self) -> usize {
192 self.rows
193 }
194}
195
196#[derive(Debug, Clone)]
198pub struct Table {
199 name: String,
200 fields: Vec<Field>,
201 stripes: Vec<Stripe>,
202 rows: usize,
203 dictionaries: Vec<Option<Page>>,
204 frequencies: Vec<Option<FrequencySummary>>,
205 version: u32,
209}
210
211impl Table {
212 #[must_use]
214 pub fn name(&self) -> &str {
215 &self.name
216 }
217
218 #[must_use]
220 pub fn fields(&self) -> &[Field] {
221 &self.fields
222 }
223
224 #[must_use]
226 pub fn rows(&self) -> usize {
227 self.rows
228 }
229
230 #[must_use]
232 pub fn stripes(&self) -> &[Stripe] {
233 &self.stripes
234 }
235}
236
237#[derive(Debug)]
239struct GlobalDictionary {
240 primary: HashMap<u64, u32>,
241 collisions: HashMap<u64, Vec<u32>>,
242 offsets: Vec<u32>,
243 payload: Vec<u8>,
244 counts: Vec<u64>,
245 nulls: u64,
246}
247
248impl GlobalDictionary {
249 fn new() -> Self {
250 Self {
251 primary: HashMap::new(),
252 collisions: HashMap::new(),
253 offsets: vec![0],
254 payload: Vec::new(),
255 counts: Vec::new(),
256 nulls: 0,
257 }
258 }
259
260 fn bytes(&self, code: u32) -> Option<&[u8]> {
261 let start = *self.offsets.get(code as usize)? as usize;
262 let end = *self.offsets.get(code as usize + 1)? as usize;
263 self.payload.get(start..end)
264 }
265
266 fn code(&mut self, text: &str) -> Result<u32> {
267 let hash = checksum(text.as_bytes());
268 if let Some(&code) = self.primary.get(&hash) {
269 if self.bytes(code) == Some(text.as_bytes()) {
270 return Ok(code);
271 }
272 if let Some(codes) = self.collisions.get(&hash) {
273 if let Some(code) =
274 codes.iter().copied().find(|&code| self.bytes(code) == Some(text.as_bytes()))
275 {
276 return Ok(code);
277 }
278 }
279 let code = self.insert(text)?;
280 self.collisions.entry(hash).or_default().push(code);
281 return Ok(code);
282 }
283 let code = self.insert(text)?;
284 self.primary.insert(hash, code);
285 Ok(code)
286 }
287
288 fn insert(&mut self, text: &str) -> Result<u32> {
289 let code = u32::try_from(self.offsets.len() - 1)
290 .map_err(|_| invalid("global dictionary has too many values"))?;
291 self.payload.extend_from_slice(text.as_bytes());
292 self.offsets.push(
293 u32::try_from(self.payload.len())
294 .map_err(|_| invalid("global dictionary payload exceeds 4 GiB"))?,
295 );
296 self.counts.push(0);
297 Ok(code)
298 }
299
300 fn ranked(&self) -> Vec<(u64, u32)> {
320 let count = self.offsets.len() - 1;
321 let mut ranked = (0..count)
322 .map(|code| {
323 let code = code as u32;
324 (head(self.bytes(code).unwrap_or_default()), code)
325 })
326 .collect::<Vec<_>>();
327 ranked.sort_unstable_by(|left, right| {
328 left.0.cmp(&right.0).then_with(|| self.bytes(left.1).cmp(&self.bytes(right.1)))
329 });
330 ranked
331 }
332
333 fn observe(&mut self, code: u32, null: bool) -> Result<()> {
334 if null {
335 self.nulls = self.nulls.saturating_add(1);
336 return Ok(());
337 }
338 let count = self
339 .counts
340 .get_mut(code as usize)
341 .ok_or_else(|| invalid("global dictionary count code is out of range"))?;
342 *count = count.saturating_add(1);
343 Ok(())
344 }
345}
346
347#[derive(Debug)]
349pub struct Writer {
350 file: File,
351 table: Table,
352 generation: u64,
353 order: Vec<(u64, u64)>,
354 next_order: u64,
355 dictionaries: Vec<Option<GlobalDictionary>>,
356 pending: Vec<PendingStripe>,
357}
358
359#[derive(Debug)]
360struct PendingStripe {
361 order: (u64, u64),
362 rows: usize,
363 pages: Vec<Vec<u8>>,
364 memberships: Vec<Option<Vec<u8>>>,
365 zone: Zone,
366}
367
368const EXTENT_STRIPES: usize = 32;
369
370impl Writer {
371 pub fn create(
377 path: impl AsRef<Path>,
378 name: impl Into<String>,
379 fields: Vec<Field>,
380 ) -> Result<Self> {
381 for field in &fields {
382 type_tag(&field.ty)?;
383 }
384 let mut file =
385 OpenOptions::new().write(true).read(true).create_new(true).open(path).map_err(io)?;
386 let mut header = [0; HEADER as usize];
387 header[..8].copy_from_slice(MAGIC);
388 header[8..12].copy_from_slice(&FORMAT.to_le_bytes());
389 file.write_all(&header).map_err(io)?;
390 Ok(Self {
391 file,
392 dictionaries: fields
393 .iter()
394 .map(|field| (field.ty == LogicalType::Varchar).then(GlobalDictionary::new))
395 .collect(),
396 table: Table {
397 name: name.into(),
398 dictionaries: vec![None; fields.len()],
399 fields,
400 stripes: Vec::new(),
401 rows: 0,
402 frequencies: Vec::new(),
403 version: FORMAT,
404 },
405 generation: 1,
406 order: Vec::new(),
407 next_order: 0,
408 pending: Vec::with_capacity(EXTENT_STRIPES),
409 })
410 }
411
412 pub fn append(&mut self, chunk: &Chunk) -> Result<()> {
418 let order = (self.next_order, 0);
419 self.next_order = self.next_order.saturating_add(1);
420 self.append_at(order, chunk)
421 }
422
423 pub fn append_at(&mut self, order: (u64, u64), chunk: &Chunk) -> Result<()> {
433 if chunk.is_empty() {
434 return Ok(());
435 }
436 if chunk.width() != self.table.fields.len() {
437 return Err(invalid("chunk width differs from table schema"));
438 }
439 let mut pages = Vec::with_capacity(chunk.width());
440 let mut memberships = Vec::with_capacity(chunk.width());
441 for (index, field) in self.table.fields.iter().enumerate() {
442 let column = chunk.column(index)?;
443 if column.logical_type() != &field.ty {
444 return Err(invalid("chunk type differs from table schema"));
445 }
446 let (bytes, membership) = encode(column, self.dictionaries[index].as_mut())?;
447 if bytes.len() > MAX_PAGE {
448 return Err(invalid("column page exceeds the configured bound"));
449 }
450 pages.push(bytes);
451 memberships.push(membership);
452 }
453 self.table.rows = self
454 .table
455 .rows
456 .checked_add(chunk.len())
457 .ok_or_else(|| invalid("row count overflow"))?;
458 self.pending.push(PendingStripe {
459 order,
460 rows: chunk.len(),
461 pages,
462 memberships,
463 zone: Zone::of(chunk),
464 });
465 if self.pending.len() == EXTENT_STRIPES {
466 self.flush_pending()?;
467 }
468 Ok(())
469 }
470
471 fn flush_pending(&mut self) -> Result<()> {
473 if self.pending.is_empty() {
474 return Ok(());
475 }
476 let width = self.table.fields.len();
477 let mut pages = vec![Vec::with_capacity(width); self.pending.len()];
478 let mut memberships = vec![vec![None; width]; self.pending.len()];
479 for column in 0..width {
480 for (stripe, pending) in self.pending.iter().enumerate() {
481 let bytes = &pending.pages[column];
482 let offset = self.file.stream_position().map_err(io)?;
483 self.file.write_all(bytes).map_err(io)?;
484 pages[stripe].push(Page {
485 offset,
486 length: u32::try_from(bytes.len())
487 .map_err(|_| invalid("page length overflow"))?,
488 hash: checksum(bytes),
489 });
490 }
491 for (stripe, pending) in self.pending.iter().enumerate() {
492 let Some(bytes) = &pending.memberships[column] else { continue };
493 let offset = self.file.stream_position().map_err(io)?;
494 self.file.write_all(bytes).map_err(io)?;
495 *memberships[stripe]
496 .get_mut(column)
497 .ok_or_else(|| invalid("membership column is missing"))? = Some(Page {
498 offset,
499 length: u32::try_from(bytes.len())
500 .map_err(|_| invalid("membership page length overflow"))?,
501 hash: checksum(bytes),
502 });
503 }
504 }
505 for ((pending, pages), memberships) in self.pending.drain(..).zip(pages).zip(memberships) {
506 self.table.stripes.push(Stripe {
507 rows: pending.rows,
508 pages,
509 memberships,
510 zone: pending.zone,
511 });
512 self.order.push(pending.order);
513 }
514 Ok(())
515 }
516
517 fn numeric_frequency(&self, column: usize) -> Result<Option<FrequencySummary>> {
521 let ty = &self.table.fields[column].ty;
522 if !matches!(
523 ty,
524 LogicalType::TinyInt
525 | LogicalType::SmallInt
526 | LogicalType::Integer
527 | LogicalType::BigInt
528 | LogicalType::UTinyInt
529 | LogicalType::USmallInt
530 | LogicalType::UInteger
531 | LogicalType::UBigInt
532 | LogicalType::Date
533 | LogicalType::Timestamp
534 ) {
535 return Ok(None);
536 }
537 let mut candidates: HashMap<FrequencyValue, u32> = HashMap::new();
538 let mut decrements = 0_u64;
539 self.visit_numeric(column, |_, value| {
540 if let Some(count) = candidates.get_mut(&value) {
541 *count = count.saturating_add(1);
542 } else if candidates.len() < FREQUENCY_CANDIDATES {
543 candidates.insert(value, 1);
544 } else {
545 candidates.retain(|_, count| {
546 *count -= 1;
547 *count != 0
548 });
549 decrements = decrements.saturating_add(1);
550 }
551 })?;
552 let (exact, ordinals) = if decrements == 0 {
553 (
554 candidates
555 .into_iter()
556 .map(|(value, count)| (value, u64::from(count)))
557 .collect::<HashMap<_, _>>(),
558 Vec::new(),
559 )
560 } else {
561 let mut lower = candidates.values().copied().collect::<Vec<_>>();
562 lower.sort_unstable_by(|left, right| right.cmp(left));
563 if lower.len() < FREQUENCY_BUILD_RANK
564 || u64::from(lower[FREQUENCY_BUILD_RANK - 1]) <= decrements
565 {
566 return Ok(None);
567 }
568 let mut exact =
569 candidates.into_keys().map(|value| (value, 0_u64)).collect::<HashMap<_, _>>();
570 let mut ordinals = Vec::new();
571 let mut exceeded = false;
572 self.visit_numeric(column, |ordinal, value| {
573 if let Some(count) = exact.get_mut(&value) {
574 *count = count.saturating_add(1);
575 if !exceeded {
576 if ordinals.len() < FREQUENCY_ORDINALS {
577 ordinals.push(ordinal);
578 } else {
579 ordinals.clear();
580 exceeded = true;
581 }
582 }
583 }
584 })?;
585 (exact, ordinals)
586 };
587 let mut entries = exact
588 .into_iter()
589 .map(|(value, count)| FrequencyEntry { value, count })
590 .collect::<Vec<_>>();
591 entries.sort_unstable_by(|left, right| {
592 right.count.cmp(&left.count).then_with(|| frequency_order(left.value, right.value))
593 });
594 let omitted_max =
595 entries.get(FREQUENCY_ENTRIES).map_or(decrements, |entry| decrements.max(entry.count));
596 entries.truncate(FREQUENCY_ENTRIES);
597 Ok(Some(FrequencySummary { entries, omitted_max, ordinals }))
598 }
599
600 fn visit_numeric(
601 &self,
602 column: usize,
603 mut visit: impl FnMut(u64, FrequencyValue),
604 ) -> Result<()> {
605 let ty = &self.table.fields[column].ty;
606 let mut start = 0_u64;
607 for stripe in &self.table.stripes {
608 let page = stripe.pages[column];
609 let mut bytes = vec![0; page.length as usize];
610 read_at(&self.file, page.offset, &mut bytes)?;
611 if checksum(&bytes) != page.hash {
612 return Err(invalid("column page checksum differs while building frequencies"));
613 }
614 let vector = decode(ty, stripe.rows, &bytes, None)?;
615 for row in 0..stripe.rows {
617 let value = if vector.is_null_at(row) {
618 FrequencyValue::Null
619 } else {
620 let widened = match vector.signed_at(row) {
624 Some(value) => Some(value),
625 None => match vector.value_at(row) {
626 Value::UTinyInt(value) => Some(i128::from(value)),
627 Value::USmallInt(value) => Some(i128::from(value)),
628 Value::UInteger(value) => Some(i128::from(value)),
629 Value::UBigInt(value) => Some(i128::from(value)),
630 _ => None,
631 },
632 };
633 FrequencyValue::Integer(widened.ok_or_else(|| {
634 invalid("numeric frequency page did not contain an integer value")
635 })?)
636 };
637 visit(start.saturating_add(row as u64), value);
638 }
639 start = start.saturating_add(stripe.rows as u64);
640 }
641 Ok(())
642 }
643
644 fn numeric_frequencies(&self) -> Result<Vec<Option<FrequencySummary>>> {
646 let columns = self
647 .table
648 .fields
649 .iter()
650 .enumerate()
651 .filter_map(|(column, field)| {
652 matches!(
653 field.ty,
654 LogicalType::TinyInt
655 | LogicalType::SmallInt
656 | LogicalType::Integer
657 | LogicalType::BigInt
658 | LogicalType::UTinyInt
659 | LogicalType::USmallInt
660 | LogicalType::UInteger
661 | LogicalType::UBigInt
662 | LogicalType::Date
663 | LogicalType::Timestamp
664 )
665 .then_some(column)
666 })
667 .collect::<Vec<_>>();
668 let workers = std::thread::available_parallelism()
669 .map_or(1, usize::from)
670 .min(MAX_FREQUENCY_WORKERS)
671 .min(columns.len());
672 if workers <= 1 {
673 let mut frequencies = vec![None; self.table.fields.len()];
674 for column in columns {
675 frequencies[column] = self.numeric_frequency(column)?;
676 }
677 return Ok(frequencies);
678 }
679 let width = columns.len().div_ceil(workers);
680 let pieces = std::thread::scope(|scope| {
681 columns
682 .chunks(width)
683 .map(|columns| {
684 scope.spawn(|| {
685 columns
686 .iter()
687 .map(|&column| Ok((column, self.numeric_frequency(column)?)))
688 .collect::<Result<Vec<_>>>()
689 })
690 })
691 .collect::<Vec<_>>()
692 .into_iter()
693 .map(|handle| {
694 handle
695 .join()
696 .map_err(|_| Error::internal("a native frequency worker panicked"))?
697 })
698 .collect::<Result<Vec<_>>>()
699 })?;
700 let mut frequencies = vec![None; self.table.fields.len()];
701 for piece in pieces {
702 for (column, summary) in piece {
703 frequencies[column] = summary;
704 }
705 }
706 Ok(frequencies)
707 }
708
709 pub fn finish(mut self) -> Result<Table> {
715 self.flush_pending()?;
716 let mut stripes = std::mem::take(&mut self.order)
717 .into_iter()
718 .zip(std::mem::take(&mut self.table.stripes))
719 .collect::<Vec<_>>();
720 stripes.sort_by_key(|(order, _)| *order);
721 self.table.stripes = stripes.into_iter().map(|(_, stripe)| stripe).collect();
722 self.table.frequencies = self.numeric_frequencies()?;
723 let dictionaries = std::mem::take(&mut self.dictionaries);
724 let orders = rankings(&dictionaries)?;
725 for (index, (dictionary, order)) in dictionaries.into_iter().zip(orders).enumerate() {
726 let Some(dictionary) = dictionary else { continue };
727 self.table.frequencies[index] = Some(code_frequency(&dictionary));
728 let encoded = encode_global_dictionary(dictionary, &order)?;
729 let offset = self.file.stream_position().map_err(io)?;
730 self.file.write_all(&encoded.index).map_err(io)?;
731 self.file.write_all(&encoded.ranks).map_err(io)?;
732 self.file.write_all(&encoded.payload).map_err(io)?;
733 let length = encoded
734 .index
735 .len()
736 .checked_add(encoded.ranks.len())
737 .and_then(|len| len.checked_add(encoded.payload.len()))
738 .ok_or_else(|| invalid("dictionary page length overflow"))?;
739 self.table.dictionaries[index] = Some(Page {
740 offset,
741 length: u32::try_from(length)
742 .map_err(|_| invalid("dictionary page length overflow"))?,
743 hash: checksum(&encoded.index),
744 });
745 }
746 let directory = encode_directory(&self.table)?;
747 if directory.len() > MAX_DIRECTORY {
748 return Err(invalid("directory exceeds the configured bound"));
749 }
750 let offset = self.file.stream_position().map_err(io)?;
751 self.file.write_all(&directory).map_err(io)?;
752 self.file.sync_all().map_err(io)?;
753 let slot = Slot {
754 offset,
755 length: u32::try_from(directory.len())
756 .map_err(|_| invalid("directory length overflow"))?,
757 generation: self.generation,
758 hash: checksum(&directory),
759 };
760 self.file.seek(SeekFrom::Start(16)).map_err(io)?;
761 self.file.write_all(&slot.bytes()).map_err(io)?;
762 self.file.sync_all().map_err(io)?;
763 Ok(self.table)
764 }
765}
766
767#[derive(Debug, Clone)]
769pub struct Reader {
770 file: Arc<File>,
771 table: Arc<Table>,
772 dictionaries: Arc<Vec<OnceLock<Arc<Vector>>>>,
773 extents: Arc<Vec<Vec<ExtentPart>>>,
774 extent_cache: Arc<Vec<Mutex<Vec<CachedExtent>>>>,
775}
776
777#[derive(Debug, Clone, Copy, Default)]
778struct ExtentPart {
779 offset: u64,
780 length: usize,
781 page_start: usize,
782}
783
784#[derive(Debug)]
785struct CachedExtent {
786 offset: u64,
787 bytes: Arc<Vec<u8>>,
788}
789
790const CACHED_EXTENTS_PER_COLUMN: usize = 8;
791
792type CrossingCache = OnceLock<Box<[OnceLock<Result<Vec<u8>>>]>>;
793
794#[derive(Debug)]
795struct NativeText {
796 file: Arc<File>,
797 offsets: Vec<u32>,
798 ranks: usize,
801 rank_at: u64,
805 rank_hashes: Vec<u64>,
806 rank_blocks: Vec<OnceLock<Result<Vec<u8>>>>,
807 payload: u64,
808 payload_len: usize,
809 hashes: Vec<u64>,
810 payload_blocks: Vec<OnceLock<Result<Vec<u8>>>>,
811 crossing: Vec<CrossingCache>,
812}
813
814const TEXT_PAYLOAD_BLOCK: usize = 64 * 1024;
815const TEXT_CROSSING_BLOCK: usize = 1024;
816
817const TEXT_RANK_BLOCK: usize = 512;
827
828const RANK_ENTRY: usize = size_of::<u64>() + size_of::<u32>();
830
831impl NativeText {
832 fn payload_block(&self, block: usize) -> Result<Option<&[u8]>> {
833 let Some(slot) = self.payload_blocks.get(block) else { return Ok(None) };
834 slot.get_or_init(|| {
835 let start = block
836 .checked_mul(TEXT_PAYLOAD_BLOCK)
837 .ok_or_else(|| invalid("global dictionary block offset overflow"))?;
838 let len = TEXT_PAYLOAD_BLOCK.min(
839 self.payload_len
840 .checked_sub(start)
841 .ok_or_else(|| invalid("global dictionary block starts past payload"))?,
842 );
843 let mut bytes = vec![0; len];
844 read_at(&self.file, self.payload + start as u64, &mut bytes)?;
845 if checksum(&bytes)
846 != *self
847 .hashes
848 .get(block)
849 .ok_or_else(|| invalid("global dictionary block has no checksum"))?
850 {
851 return Err(invalid("global dictionary payload checksum differs"));
852 }
853 Ok(bytes)
854 })
855 .as_ref()
856 .map(|bytes| Some(bytes.as_slice()))
857 .map_err(Clone::clone)
858 }
859
860 fn rank_parts(&self, rank: usize) -> Result<(&[u8], usize)> {
867 let slot = self
868 .rank_blocks
869 .get(rank / TEXT_RANK_BLOCK)
870 .ok_or_else(|| invalid("global dictionary rank is past the order"))?;
871 let block = slot
872 .get_or_init(|| {
873 let first = rank / TEXT_RANK_BLOCK * TEXT_RANK_BLOCK;
874 let len = TEXT_RANK_BLOCK.min(self.ranks - first) * RANK_ENTRY;
875 let mut bytes = vec![0; len];
876 read_at(&self.file, self.rank_at + (first * RANK_ENTRY) as u64, &mut bytes)?;
877 if checksum(&bytes)
878 != *self
879 .rank_hashes
880 .get(rank / TEXT_RANK_BLOCK)
881 .ok_or_else(|| invalid("global dictionary rank block has no checksum"))?
882 {
883 return Err(invalid("global dictionary rank checksum differs"));
884 }
885 Ok(bytes)
886 })
887 .as_ref()
888 .map_err(Clone::clone)?;
889 Ok((block.as_slice(), rank % TEXT_RANK_BLOCK))
890 }
891
892 fn head_at(&self, rank: usize) -> Result<u64> {
894 let (block, within) = self.rank_parts(rank)?;
895 let at = within * size_of::<u64>();
896 let bytes = block
897 .get(at..at + size_of::<u64>())
898 .ok_or_else(|| invalid("global dictionary rank block is short of heads"))?;
899 Ok(u64::from_le_bytes(bytes.try_into().expect("eight bytes")))
900 }
901}
902
903impl TextSource for NativeText {
904 fn len(&self) -> usize {
905 self.offsets.len().saturating_sub(1)
906 }
907
908 fn bytes_at(&self, index: usize) -> Result<Option<&[u8]>> {
909 let (Some(&start), Some(&end)) = (self.offsets.get(index), self.offsets.get(index + 1))
910 else {
911 return Ok(None);
912 };
913 if start == end {
914 return Ok(Some(&[]));
915 }
916 let first = start as usize / TEXT_PAYLOAD_BLOCK;
917 let last = (end as usize - 1) / TEXT_PAYLOAD_BLOCK;
918 if first == last {
919 let Some(block) = self.payload_block(first)? else { return Ok(None) };
920 let within = start as usize % TEXT_PAYLOAD_BLOCK;
921 return Ok(block.get(within..within + (end - start) as usize));
922 }
923 let Some(crossing) = self.crossing.get(index / TEXT_CROSSING_BLOCK) else {
924 return Ok(None);
925 };
926 let block = crossing.get_or_init(|| {
927 (0..TEXT_CROSSING_BLOCK).map(|_| OnceLock::new()).collect::<Vec<_>>().into_boxed_slice()
928 });
929 block[index % TEXT_CROSSING_BLOCK]
930 .get_or_init(|| {
931 let mut bytes = Vec::with_capacity((end - start) as usize);
932 for part in first..=last {
933 let source = self
934 .payload_block(part)?
935 .ok_or_else(|| invalid("global dictionary block is missing"))?;
936 let from = if part == first { start as usize % TEXT_PAYLOAD_BLOCK } else { 0 };
937 let to = if part == last {
938 (end as usize - 1) % TEXT_PAYLOAD_BLOCK + 1
939 } else {
940 source.len()
941 };
942 bytes.extend_from_slice(source.get(from..to).ok_or_else(|| {
943 invalid("global dictionary value exceeds its payload block")
944 })?);
945 }
946 Ok(bytes)
947 })
948 .as_ref()
949 .map(|bytes| Some(bytes.as_slice()))
950 .map_err(Clone::clone)
951 }
952
953 fn bytes_len_at(&self, index: usize) -> Result<Option<usize>> {
954 let (Some(&start), Some(&end)) = (self.offsets.get(index), self.offsets.get(index + 1))
955 else {
956 return Ok(None);
957 };
958 Ok(Some((end - start) as usize))
959 }
960
961 fn ranks(&self) -> Option<usize> {
962 (self.ranks > 0).then_some(self.ranks)
963 }
964
965 fn compare_rank(&self, rank: usize, wanted: &[u8]) -> Result<Ordering> {
966 let settled = self.head_at(rank)?.cmp(&head(wanted));
970 if settled != Ordering::Equal {
971 return Ok(settled);
972 }
973 let code = self.code_at_rank(rank)?;
974 let bytes = self
975 .bytes_at(code as usize)?
976 .ok_or_else(|| invalid("global dictionary order names a code it does not have"))?;
977 Ok(bytes.cmp(wanted))
978 }
979
980 fn code_at_rank(&self, rank: usize) -> Result<u32> {
981 let (block, within) = self.rank_parts(rank)?;
982 let heads = block.len() / RANK_ENTRY * size_of::<u64>();
983 let at = heads + within * size_of::<u32>();
984 let bytes = block
985 .get(at..at + size_of::<u32>())
986 .ok_or_else(|| invalid("global dictionary rank block is short of codes"))?;
987 let code = u32::from_le_bytes(bytes.try_into().expect("four bytes"));
988 if code as usize >= self.len() {
989 return Err(invalid("global dictionary order names a code it does not have"));
990 }
991 Ok(code)
992 }
993
994 fn footprint(&self) -> usize {
995 self.offsets.capacity() * size_of::<u32>()
996 + self.rank_hashes.capacity() * size_of::<u64>()
997 + self.rank_blocks.capacity() * size_of::<OnceLock<Result<Vec<u8>>>>()
998 + self
999 .rank_blocks
1000 .iter()
1001 .filter_map(OnceLock::get)
1002 .filter_map(|result| result.as_ref().ok())
1003 .map(Vec::capacity)
1004 .sum::<usize>()
1005 + self.payload_blocks.capacity() * size_of::<OnceLock<Result<Vec<u8>>>>()
1006 + self.hashes.capacity() * size_of::<u64>()
1007 + self
1008 .payload_blocks
1009 .iter()
1010 .filter_map(OnceLock::get)
1011 .filter_map(|result| result.as_ref().ok())
1012 .map(Vec::capacity)
1013 .sum::<usize>()
1014 + self.crossing.capacity() * size_of::<CrossingCache>()
1015 + self
1016 .crossing
1017 .iter()
1018 .filter_map(OnceLock::get)
1019 .map(|block| {
1020 block.len() * size_of::<OnceLock<Result<Vec<u8>>>>()
1021 + block
1022 .iter()
1023 .filter_map(OnceLock::get)
1024 .filter_map(|result| result.as_ref().ok())
1025 .map(Vec::capacity)
1026 .sum::<usize>()
1027 })
1028 .sum::<usize>()
1029 }
1030}
1031
1032fn extent_parts(table: &Table) -> Result<Vec<Vec<ExtentPart>>> {
1034 let mut refs = Vec::with_capacity(table.stripes.len().saturating_mul(table.fields.len()));
1035 for (stripe, entry) in table.stripes.iter().enumerate() {
1036 for (column, page) in entry.pages.iter().enumerate() {
1037 refs.push((page.offset, column, stripe, page.length as usize));
1038 }
1039 }
1040 refs.sort_unstable_by_key(|entry| entry.0);
1041 let mut parts = vec![vec![ExtentPart::default(); table.fields.len()]; table.stripes.len()];
1042 let mut first = 0;
1043 while first < refs.len() {
1044 let (offset, column, _, first_len) = refs[first];
1045 let mut end = offset
1046 .checked_add(first_len as u64)
1047 .ok_or_else(|| invalid("column extent range overflow"))?;
1048 let mut last = first + 1;
1049 while last < refs.len()
1050 && last - first < EXTENT_STRIPES
1051 && refs[last].1 == column
1052 && refs[last].0 == end
1053 {
1054 end = end
1055 .checked_add(refs[last].3 as u64)
1056 .ok_or_else(|| invalid("column extent range overflow"))?;
1057 last += 1;
1058 }
1059 let length = usize::try_from(end - offset)
1060 .map_err(|_| invalid("column extent length exceeds this platform"))?;
1061 for &(_, _, stripe, _) in &refs[first..last] {
1062 let page = table.stripes[stripe].pages[column];
1063 let page_start = usize::try_from(page.offset - offset)
1064 .map_err(|_| invalid("column page offset exceeds this platform"))?;
1065 parts[stripe][column] = ExtentPart { offset, length, page_start };
1066 }
1067 first = last;
1068 }
1069 Ok(parts)
1070}
1071
1072impl Reader {
1073 pub fn open(path: impl AsRef<Path>) -> Result<Self> {
1079 let mut file = File::open(path).map_err(io)?;
1080 let size = file.metadata().map_err(io)?.len();
1081 if size < HEADER {
1082 return Err(invalid("file is shorter than its header"));
1083 }
1084 let mut header = [0; HEADER as usize];
1085 file.read_exact(&mut header).map_err(io)?;
1086 let version = u32::from_le_bytes([header[8], header[9], header[10], header[11]]);
1087 let known = [(MAGIC, FORMAT), (MAGIC_V8, 8), (MAGIC_V7, 7)];
1088 if !known.iter().any(|(magic, known)| &header[..8] == *magic && version == *known) {
1089 return Err(invalid("magic or major version is unsupported"));
1090 }
1091 let mut selected = None;
1092 for start in [16, 16 + SLOT_BYTES] {
1093 let slot = Slot::read(&header[start..start + SLOT_BYTES]);
1094 if slot.generation == 0 || slot.length == 0 || slot.length as usize > MAX_DIRECTORY {
1095 continue;
1096 }
1097 let Some(end) = slot.offset.checked_add(u64::from(slot.length)) else { continue };
1098 if slot.offset < HEADER || end > size {
1099 continue;
1100 }
1101 let mut bytes = vec![0; slot.length as usize];
1102 file.seek(SeekFrom::Start(slot.offset)).map_err(io)?;
1103 file.read_exact(&mut bytes).map_err(io)?;
1104 if checksum(&bytes) == slot.hash
1105 && selected
1106 .as_ref()
1107 .is_none_or(|(old, _): &(Slot, Vec<u8>)| old.generation < slot.generation)
1108 {
1109 selected = Some((slot, bytes));
1110 }
1111 }
1112 let (_, bytes) = selected.ok_or_else(|| invalid("no committed directory slot is valid"))?;
1113 let table = decode_directory(&bytes, size, version)?;
1114 let extents = extent_parts(&table)?;
1115 let dictionaries = (0..table.fields.len()).map(|_| OnceLock::new()).collect();
1116 let extent_cache = (0..table.fields.len())
1117 .map(|_| Mutex::new(Vec::with_capacity(CACHED_EXTENTS_PER_COLUMN)))
1118 .collect::<Vec<_>>();
1119 Ok(Self {
1120 file: Arc::new(file),
1121 table: Arc::new(table),
1122 dictionaries: Arc::new(dictionaries),
1123 extents: Arc::new(extents),
1124 extent_cache: Arc::new(extent_cache),
1125 })
1126 }
1127
1128 #[must_use]
1130 pub fn table(&self) -> &Table {
1131 &self.table
1132 }
1133
1134 pub fn top_frequencies(&self, column: usize, top: usize) -> Result<Option<Vec<(Value, u64)>>> {
1143 let field = self
1144 .table
1145 .fields
1146 .get(column)
1147 .ok_or_else(|| invalid("frequency column index out of range"))?;
1148 let Some(summary) = self.table.frequencies.get(column).and_then(Option::as_ref) else {
1149 return Ok(None);
1150 };
1151 if top == 0 || summary.entries.len() < top {
1152 return Ok(None);
1153 }
1154 let boundary = summary.entries[top - 1].count;
1155 if boundary <= summary.omitted_max {
1156 return Ok(None);
1157 }
1158 let dictionary =
1159 if field.ty == LogicalType::Varchar { self.dictionary(column)? } else { None };
1160 let mut out = Vec::with_capacity(summary.entries.len());
1161 for entry in &summary.entries {
1162 let value = match entry.value {
1163 FrequencyValue::Null => Value::Null,
1164 FrequencyValue::Integer(value) => match field.ty {
1165 LogicalType::TinyInt => Value::TinyInt(
1166 i8::try_from(value)
1167 .map_err(|_| invalid("frequency TINYINT is out of range"))?,
1168 ),
1169 LogicalType::UTinyInt => Value::UTinyInt(
1170 u8::try_from(value)
1171 .map_err(|_| invalid("frequency UTINYINT is out of range"))?,
1172 ),
1173 LogicalType::USmallInt => Value::USmallInt(
1174 u16::try_from(value)
1175 .map_err(|_| invalid("frequency USMALLINT is out of range"))?,
1176 ),
1177 LogicalType::UInteger => Value::UInteger(
1178 u32::try_from(value)
1179 .map_err(|_| invalid("frequency UINTEGER is out of range"))?,
1180 ),
1181 LogicalType::UBigInt => Value::UBigInt(
1182 u64::try_from(value)
1183 .map_err(|_| invalid("frequency UBIGINT is out of range"))?,
1184 ),
1185 LogicalType::SmallInt => Value::SmallInt(
1186 i16::try_from(value)
1187 .map_err(|_| invalid("frequency SMALLINT is out of range"))?,
1188 ),
1189 LogicalType::Integer => Value::Integer(
1190 i32::try_from(value)
1191 .map_err(|_| invalid("frequency INTEGER is out of range"))?,
1192 ),
1193 LogicalType::BigInt => Value::BigInt(
1194 i64::try_from(value)
1195 .map_err(|_| invalid("frequency BIGINT is out of range"))?,
1196 ),
1197 LogicalType::Date => Value::Date(
1198 i32::try_from(value)
1199 .map_err(|_| invalid("frequency DATE is out of range"))?,
1200 ),
1201 LogicalType::Timestamp => Value::Timestamp(
1202 i64::try_from(value)
1203 .map_err(|_| invalid("frequency TIMESTAMP is out of range"))?,
1204 ),
1205 _ => return Err(invalid("integer frequency belongs to another type")),
1206 },
1207 FrequencyValue::Code(code) => dictionary
1208 .as_ref()
1209 .ok_or_else(|| invalid("frequency code has no dictionary"))?
1210 .try_value_at(code as usize)?,
1211 };
1212 out.push((value, entry.count));
1213 }
1214 Ok(Some(out))
1215 }
1216
1217 pub fn frequency_occurrences(&self, column: usize) -> Result<Option<FrequencyOccurrences>> {
1227 self.table
1228 .fields
1229 .get(column)
1230 .ok_or_else(|| invalid("frequency column index out of range"))?;
1231 let Some(summary) = self.table.frequencies.get(column).and_then(Option::as_ref) else {
1232 return Ok(None);
1233 };
1234 if summary.ordinals.is_empty() {
1235 return Ok(None);
1236 }
1237 Ok(Some(FrequencyOccurrences {
1238 omitted_max: summary.omitted_max,
1239 ordinals: summary.ordinals.clone(),
1240 }))
1241 }
1242
1243 fn dictionary(&self, column: usize) -> Result<Option<Arc<Vector>>> {
1244 let Some(page) = self.table.dictionaries[column] else { return Ok(None) };
1245 if let Some(dictionary) = self.dictionaries[column].get() {
1246 return Ok(Some(Arc::clone(dictionary)));
1247 }
1248 let dictionary = Arc::new(open_global_dictionary(
1249 Arc::clone(&self.file),
1250 page,
1251 &self.table.fields[column].ty,
1252 self.table.version >= 9,
1253 )?);
1254 let _ = self.dictionaries[column].set(Arc::clone(&dictionary));
1255 Ok(Some(self.dictionaries[column].get().map_or(dictionary, Arc::clone)))
1256 }
1257
1258 pub fn read(&self, stripe: usize, columns: &[usize]) -> Result<Chunk> {
1264 self.read_impl(stripe, columns, true)
1265 }
1266
1267 pub fn read_sparse(&self, stripe: usize, columns: &[usize]) -> Result<Chunk> {
1276 self.read_impl(stripe, columns, false)
1277 }
1278
1279 pub fn skips_codes(&self, stripe: usize, column: usize, candidates: &[u32]) -> Result<bool> {
1288 if candidates.is_empty() {
1289 return Ok(true);
1290 }
1291 if candidates.windows(2).any(|pair| pair[0] >= pair[1]) {
1292 return Err(Error::internal("native code candidates are not sorted and unique"));
1293 }
1294 let stripe =
1295 self.table.stripes.get(stripe).ok_or_else(|| invalid("stripe index out of range"))?;
1296 let Some(page) = stripe.memberships.get(column).copied().flatten() else {
1297 return Ok(false);
1298 };
1299 let mut bytes = vec![0; page.length as usize];
1300 read_at(&self.file, page.offset, &mut bytes)?;
1301 if checksum(&bytes) != page.hash {
1302 return Err(invalid("membership page checksum differs"));
1303 }
1304 let codes = decode_membership(&bytes)?;
1305 let mut left = 0;
1306 let mut right = 0;
1307 while left < codes.len() && right < candidates.len() {
1308 match codes[left].cmp(&candidates[right]) {
1309 Ordering::Less => left += 1,
1310 Ordering::Greater => right += 1,
1311 Ordering::Equal => return Ok(false),
1312 }
1313 }
1314 Ok(true)
1315 }
1316
1317 fn read_impl(&self, stripe: usize, columns: &[usize], prefetch: bool) -> Result<Chunk> {
1318 let stripe_index = stripe;
1319 let stripe =
1320 self.table.stripes.get(stripe).ok_or_else(|| invalid("stripe index out of range"))?;
1321 let mut picked = Vec::with_capacity(columns.len());
1322 for &column in columns {
1323 let field = self
1324 .table
1325 .fields
1326 .get(column)
1327 .ok_or_else(|| invalid("column index out of range"))?;
1328 let page = stripe.pages.get(column).ok_or_else(|| invalid("stripe page is missing"))?;
1329 let part = self
1330 .extents
1331 .get(stripe_index)
1332 .and_then(|parts| parts.get(column))
1333 .ok_or_else(|| invalid("column extent is missing"))?;
1334 let bytes = if !prefetch || part.length == page.length as usize {
1335 let mut bytes = vec![0; page.length as usize];
1336 read_at(&self.file, page.offset, &mut bytes)?;
1337 Arc::new(bytes)
1338 } else {
1339 let cached = self.extent_cache[column]
1340 .lock()
1341 .map_err(|_| invalid("column extent cache is poisoned"))?
1342 .iter()
1343 .find(|cached| cached.offset == part.offset)
1344 .map(|cached| Arc::clone(&cached.bytes));
1345 if let Some(bytes) = cached {
1346 bytes
1347 } else {
1348 let mut bytes = vec![0; part.length];
1349 read_at(&self.file, part.offset, &mut bytes)?;
1350 let bytes = Arc::new(bytes);
1351 let mut cache = self.extent_cache[column]
1352 .lock()
1353 .map_err(|_| invalid("column extent cache is poisoned"))?;
1354 if let Some(cached) = cache.iter().find(|cached| cached.offset == part.offset) {
1355 Arc::clone(&cached.bytes)
1356 } else {
1357 if cache.len() == CACHED_EXTENTS_PER_COLUMN {
1358 cache.remove(0);
1359 }
1360 cache.push(CachedExtent { offset: part.offset, bytes: Arc::clone(&bytes) });
1361 bytes
1362 }
1363 }
1364 };
1365 let page_start = if prefetch { part.page_start } else { 0 };
1366 let end = page_start
1367 .checked_add(page.length as usize)
1368 .ok_or_else(|| invalid("column page range overflow"))?;
1369 let page_bytes = bytes
1370 .get(page_start..end)
1371 .ok_or_else(|| invalid("column page exceeds its extent"))?;
1372 if checksum(page_bytes) != page.hash {
1373 return Err(invalid("column page checksum differs"));
1374 }
1375 let dictionary = self.dictionary(column)?;
1376 picked.push(decode(&field.ty, stripe.rows, page_bytes, dictionary)?);
1377 }
1378 Chunk::with_rows(picked, stripe.rows)
1379 }
1380
1381 #[must_use]
1383 pub fn skips(&self, stripe: usize, probes: &[Probe]) -> bool {
1384 self.table.stripes.get(stripe).is_some_and(|stripe| stripe.zone.skips(probes))
1385 }
1386}
1387
1388#[cfg(unix)]
1389fn read_at(file: &File, mut offset: u64, mut bytes: &mut [u8]) -> Result<()> {
1390 use std::os::unix::fs::FileExt;
1391 while !bytes.is_empty() {
1392 let read = file.read_at(bytes, offset).map_err(io)?;
1393 if read == 0 {
1394 return Err(invalid("column page ends before its declared length"));
1395 }
1396 offset += read as u64;
1397 bytes = &mut bytes[read..];
1398 }
1399 Ok(())
1400}
1401
1402#[cfg(not(unix))]
1403fn read_at(file: &File, offset: u64, bytes: &mut [u8]) -> Result<()> {
1404 let mut file = file.try_clone().map_err(io)?;
1405 file.seek(SeekFrom::Start(offset)).map_err(io)?;
1406 file.read_exact(bytes).map_err(io)
1407}
1408
1409fn type_tag(ty: &LogicalType) -> Result<u8> {
1410 match ty {
1411 LogicalType::SmallInt => Ok(1),
1412 LogicalType::Integer => Ok(2),
1413 LogicalType::BigInt => Ok(3),
1414 LogicalType::Varchar => Ok(4),
1415 LogicalType::Date => Ok(5),
1416 LogicalType::Timestamp => Ok(6),
1417 LogicalType::Boolean => Ok(7),
1418 LogicalType::TinyInt => Ok(8),
1419 LogicalType::UTinyInt => Ok(9),
1420 LogicalType::USmallInt => Ok(10),
1421 LogicalType::UInteger => Ok(11),
1422 LogicalType::UBigInt => Ok(12),
1423 _ => Err(Error::not_implemented(format!("native storage for {ty}"))),
1424 }
1425}
1426
1427fn tag_type(tag: u8) -> Result<LogicalType> {
1428 match tag {
1429 1 => Ok(LogicalType::SmallInt),
1430 2 => Ok(LogicalType::Integer),
1431 3 => Ok(LogicalType::BigInt),
1432 4 => Ok(LogicalType::Varchar),
1433 5 => Ok(LogicalType::Date),
1434 6 => Ok(LogicalType::Timestamp),
1435 7 => Ok(LogicalType::Boolean),
1436 8 => Ok(LogicalType::TinyInt),
1437 9 => Ok(LogicalType::UTinyInt),
1438 10 => Ok(LogicalType::USmallInt),
1439 11 => Ok(LogicalType::UInteger),
1440 12 => Ok(LogicalType::UBigInt),
1441 _ => Err(invalid("column type tag is unknown")),
1442 }
1443}
1444
1445fn put_u16(out: &mut Vec<u8>, value: u16) {
1446 out.extend_from_slice(&value.to_le_bytes());
1447}
1448fn put_u32(out: &mut Vec<u8>, value: u32) {
1449 out.extend_from_slice(&value.to_le_bytes());
1450}
1451fn put_u64(out: &mut Vec<u8>, value: u64) {
1452 out.extend_from_slice(&value.to_le_bytes());
1453}
1454fn put_var_u64(out: &mut Vec<u8>, mut value: u64) {
1455 while value >= 0x80 {
1456 out.push((value as u8 & 0x7f) | 0x80);
1457 value >>= 7;
1458 }
1459 out.push(value as u8);
1460}
1461
1462fn frequency_order(left: FrequencyValue, right: FrequencyValue) -> Ordering {
1463 match (left, right) {
1464 (FrequencyValue::Null, FrequencyValue::Null) => Ordering::Equal,
1465 (FrequencyValue::Null, _) => Ordering::Less,
1466 (_, FrequencyValue::Null) => Ordering::Greater,
1467 (FrequencyValue::Integer(left), FrequencyValue::Integer(right)) => left.cmp(&right),
1468 (FrequencyValue::Code(left), FrequencyValue::Code(right)) => left.cmp(&right),
1469 (FrequencyValue::Integer(_), FrequencyValue::Code(_)) => Ordering::Less,
1470 (FrequencyValue::Code(_), FrequencyValue::Integer(_)) => Ordering::Greater,
1471 }
1472}
1473
1474fn code_frequency(dictionary: &GlobalDictionary) -> FrequencySummary {
1475 let mut entries = dictionary
1476 .counts
1477 .iter()
1478 .enumerate()
1479 .filter(|(_, count)| **count != 0)
1480 .map(|(code, &count)| FrequencyEntry { value: FrequencyValue::Code(code as u32), count })
1481 .collect::<Vec<_>>();
1482 if dictionary.nulls != 0 {
1483 entries.push(FrequencyEntry { value: FrequencyValue::Null, count: dictionary.nulls });
1484 }
1485 entries.sort_unstable_by(|left, right| {
1486 right.count.cmp(&left.count).then_with(|| frequency_order(left.value, right.value))
1487 });
1488 let omitted_max = entries.get(FREQUENCY_ENTRIES).map_or(0, |entry| entry.count);
1489 entries.truncate(FREQUENCY_ENTRIES);
1490 FrequencySummary { entries, omitted_max, ordinals: Vec::new() }
1491}
1492
1493fn encode_directory(table: &Table) -> Result<Vec<u8>> {
1494 encode_directory_version(table, FORMAT)
1495}
1496
1497fn directory_magic(version: u32) -> &'static [u8; 8] {
1499 match version {
1500 7 => DIRECTORY_V7,
1501 8 => DIRECTORY_V8,
1502 _ => DIRECTORY,
1503 }
1504}
1505
1506fn encode_directory_version(table: &Table, version: u32) -> Result<Vec<u8>> {
1507 let mut out = directory_magic(version).to_vec();
1508 let name = table.name.as_bytes();
1509 put_u16(&mut out, u16::try_from(name.len()).map_err(|_| invalid("table name too long"))?);
1510 out.extend_from_slice(name);
1511 put_u16(&mut out, u16::try_from(table.fields.len()).map_err(|_| invalid("too many columns"))?);
1512 for field in &table.fields {
1513 let name = field.name.as_bytes();
1514 put_u16(&mut out, u16::try_from(name.len()).map_err(|_| invalid("column name too long"))?);
1515 out.extend_from_slice(name);
1516 out.push(type_tag(&field.ty)?);
1517 out.push(u8::from(field.not_null));
1518 }
1519 for dictionary in &table.dictionaries {
1520 match dictionary {
1521 None => out.push(0),
1522 Some(page) => {
1523 out.push(1);
1524 put_u64(&mut out, page.offset);
1525 put_u32(&mut out, page.length);
1526 put_u64(&mut out, page.hash);
1527 }
1528 }
1529 }
1530 put_u64(&mut out, u64::try_from(table.rows).map_err(|_| invalid("row count overflow"))?);
1531 put_u32(&mut out, u32::try_from(table.stripes.len()).map_err(|_| invalid("too many stripes"))?);
1532 for stripe in &table.stripes {
1533 put_u32(
1534 &mut out,
1535 u32::try_from(stripe.rows).map_err(|_| invalid("stripe row count overflow"))?,
1536 );
1537 for page in &stripe.pages {
1538 put_u64(&mut out, page.offset);
1539 put_u32(&mut out, page.length);
1540 put_u64(&mut out, page.hash);
1541 }
1542 if version >= 8 {
1543 for (field, membership) in table.fields.iter().zip(&stripe.memberships) {
1544 if field.ty != LogicalType::Varchar {
1545 continue;
1546 }
1547 let page = membership
1548 .ok_or_else(|| invalid("string page has no code membership index"))?;
1549 put_u64(&mut out, page.offset);
1550 put_u32(&mut out, page.length);
1551 put_u64(&mut out, page.hash);
1552 }
1553 }
1554 for range in stripe.zone.columns() {
1555 put_bound(&mut out, range.low.as_ref())?;
1556 put_bound(&mut out, range.high.as_ref())?;
1557 put_u32(
1558 &mut out,
1559 u32::try_from(range.nulls).map_err(|_| invalid("null count overflow"))?,
1560 );
1561 }
1562 }
1563 out.extend_from_slice(FREQUENCIES);
1564 put_u16(
1565 &mut out,
1566 u16::try_from(table.frequencies.len())
1567 .map_err(|_| invalid("too many frequency columns"))?,
1568 );
1569 for summary in &table.frequencies {
1570 let Some(summary) = summary else {
1571 out.push(0);
1572 continue;
1573 };
1574 out.push(1);
1575 put_u64(&mut out, summary.omitted_max);
1576 put_u32(
1577 &mut out,
1578 u32::try_from(summary.entries.len())
1579 .map_err(|_| invalid("too many frequency entries"))?,
1580 );
1581 for entry in &summary.entries {
1582 match entry.value {
1583 FrequencyValue::Null => out.push(0),
1584 FrequencyValue::Integer(value) => {
1585 out.push(1);
1586 out.extend_from_slice(&value.to_le_bytes());
1587 }
1588 FrequencyValue::Code(value) => {
1589 out.push(2);
1590 put_u32(&mut out, value);
1591 }
1592 }
1593 put_u64(&mut out, entry.count);
1594 }
1595 put_u32(
1596 &mut out,
1597 u32::try_from(summary.ordinals.len())
1598 .map_err(|_| invalid("too many frequency ordinals"))?,
1599 );
1600 let mut previous = 0_u64;
1601 for (at, &ordinal) in summary.ordinals.iter().enumerate() {
1602 let delta = if at == 0 {
1603 ordinal
1604 } else {
1605 ordinal
1606 .checked_sub(previous)
1607 .ok_or_else(|| invalid("frequency ordinals are not ordered"))?
1608 };
1609 if at != 0 && delta == 0 {
1610 return Err(invalid("frequency ordinals are not unique"));
1611 }
1612 put_var_u64(&mut out, delta);
1613 previous = ordinal;
1614 }
1615 }
1616 Ok(out)
1617}
1618
1619struct Cursor<'a> {
1620 bytes: &'a [u8],
1621 at: usize,
1622}
1623impl<'a> Cursor<'a> {
1624 fn take(&mut self, len: usize) -> Result<&'a [u8]> {
1625 let end = self.at.checked_add(len).ok_or_else(|| invalid("directory offset overflow"))?;
1626 let bytes =
1627 self.bytes.get(self.at..end).ok_or_else(|| invalid("directory is truncated"))?;
1628 self.at = end;
1629 Ok(bytes)
1630 }
1631 fn u8(&mut self) -> Result<u8> {
1632 Ok(self.take(1)?[0])
1633 }
1634 fn u16(&mut self) -> Result<u16> {
1635 Ok(u16::from_le_bytes(self.take(2)?.try_into().expect("two bytes")))
1636 }
1637 fn u32(&mut self) -> Result<u32> {
1638 Ok(u32::from_le_bytes(self.take(4)?.try_into().expect("four bytes")))
1639 }
1640 fn u64(&mut self) -> Result<u64> {
1641 Ok(u64::from_le_bytes(self.take(8)?.try_into().expect("eight bytes")))
1642 }
1643 fn var_u64(&mut self) -> Result<u64> {
1644 let mut value = 0_u64;
1645 for shift in (0..=63).step_by(7) {
1646 let byte = self.u8()?;
1647 let part = u64::from(byte & 0x7f);
1648 if shift == 63 && part > 1 {
1649 return Err(invalid("frequency ordinal varint overflows"));
1650 }
1651 value |= part << shift;
1652 if byte & 0x80 == 0 {
1653 return Ok(value);
1654 }
1655 }
1656 Err(invalid("frequency ordinal varint is too long"))
1657 }
1658 fn bound(&mut self) -> Result<Option<Bound>> {
1659 Ok(match self.u8()? {
1660 0 => None,
1661 1 => Some(Bound::Int(i128::from_le_bytes(
1662 self.take(16)?.try_into().expect("sixteen bytes"),
1663 ))),
1664 2 => Some(Bound::Real(f64::from_le_bytes(
1665 self.take(8)?.try_into().expect("eight bytes"),
1666 ))),
1667 3 => {
1668 let length = self.u32()? as usize;
1669 Some(Bound::Bytes(self.take(length)?.to_vec()))
1670 }
1671 _ => return Err(invalid("bound tag differs")),
1672 })
1673 }
1674 fn text(&mut self) -> Result<String> {
1675 let len = self.u16()? as usize;
1676 String::from_utf8(self.take(len)?.to_vec()).map_err(|_| invalid("name is not UTF-8"))
1677 }
1678}
1679
1680fn decode_directory(bytes: &[u8], size: u64, version: u32) -> Result<Table> {
1681 let mut cur = Cursor { bytes, at: 0 };
1682 if cur.take(8)? != directory_magic(version) {
1683 return Err(invalid("directory magic differs"));
1684 }
1685 let name = cur.text()?;
1686 let width = cur.u16()? as usize;
1687 let mut fields = Vec::with_capacity(width);
1688 for _ in 0..width {
1689 let name = cur.text()?;
1690 let ty = tag_type(cur.u8()?)?;
1691 let not_null = match cur.u8()? {
1692 0 => false,
1693 1 => true,
1694 _ => return Err(invalid("nullability flag differs")),
1695 };
1696 fields.push(Field { name, ty, not_null });
1697 }
1698 let mut dictionaries = Vec::with_capacity(width);
1699 for _ in 0..width {
1700 dictionaries.push(match cur.u8()? {
1701 0 => None,
1702 1 => {
1703 let page = Page { offset: cur.u64()?, length: cur.u32()?, hash: cur.u64()? };
1704 let end = page
1705 .offset
1706 .checked_add(u64::from(page.length))
1707 .ok_or_else(|| invalid("dictionary page offset overflow"))?;
1708 if page.offset < HEADER || end > size {
1713 return Err(invalid("dictionary page range is outside the file"));
1714 }
1715 Some(page)
1716 }
1717 _ => return Err(invalid("dictionary page tag differs")),
1718 });
1719 }
1720 let rows = usize::try_from(cur.u64()?).map_err(|_| invalid("row count does not fit"))?;
1721 let count = cur.u32()? as usize;
1722 let mut stripes = Vec::with_capacity(count);
1723 let mut total = 0_usize;
1724 for _ in 0..count {
1725 let stripe_rows = cur.u32()? as usize;
1726 if stripe_rows == 0 {
1727 return Err(invalid("empty stripe"));
1728 }
1729 total =
1730 total.checked_add(stripe_rows).ok_or_else(|| invalid("stripe row count overflow"))?;
1731 let mut pages = Vec::with_capacity(width);
1732 for _ in 0..width {
1733 let offset = cur.u64()?;
1734 let length = cur.u32()?;
1735 let hash = cur.u64()?;
1736 let end = offset
1737 .checked_add(u64::from(length))
1738 .ok_or_else(|| invalid("page offset overflow"))?;
1739 if offset < HEADER || end > size || length as usize > MAX_PAGE {
1740 return Err(invalid("page range is outside the file"));
1741 }
1742 pages.push(Page { offset, length, hash });
1743 }
1744 let mut memberships = vec![None; width];
1745 if version >= 8 {
1746 for (column, field) in fields.iter().enumerate() {
1747 if field.ty != LogicalType::Varchar {
1748 continue;
1749 }
1750 let page = Page { offset: cur.u64()?, length: cur.u32()?, hash: cur.u64()? };
1751 let end = page
1752 .offset
1753 .checked_add(u64::from(page.length))
1754 .ok_or_else(|| invalid("membership page offset overflow"))?;
1755 if page.offset < HEADER || end > size || page.length as usize > MAX_PAGE {
1756 return Err(invalid("membership page range is outside the file"));
1757 }
1758 memberships[column] = Some(page);
1759 }
1760 }
1761 let mut ranges = Vec::with_capacity(width);
1762 for _ in 0..width {
1763 let low = cur.bound()?;
1764 let high = cur.bound()?;
1765 let nulls = cur.u32()? as usize;
1766 if nulls > stripe_rows {
1767 return Err(invalid("null count exceeds stripe rows"));
1768 }
1769 ranges.push(Range { low, high, nulls });
1770 }
1771 stripes.push(Stripe {
1772 rows: stripe_rows,
1773 pages,
1774 memberships,
1775 zone: Zone::from_ranges(ranges),
1776 });
1777 }
1778 if total != rows {
1779 return Err(invalid("table row count differs from stripes"));
1780 }
1781 let frequencies = if cur.at == bytes.len() {
1782 vec![None; width]
1783 } else {
1784 let frequency_version = match cur.take(8)? {
1785 magic if magic == FREQUENCIES_V1 => 1,
1786 magic if magic == FREQUENCIES => 2,
1787 _ => return Err(invalid("directory extension magic differs")),
1788 };
1789 if cur.u16()? as usize != width {
1790 return Err(invalid("frequency column count differs"));
1791 }
1792 let mut frequencies = Vec::with_capacity(width);
1793 for field in &fields {
1794 let summary = match cur.u8()? {
1795 0 => None,
1796 1 => {
1797 let omitted_max = cur.u64()?;
1798 let count = cur.u32()? as usize;
1799 if count > FREQUENCY_ENTRIES {
1800 return Err(invalid("frequency entry count exceeds its bound"));
1801 }
1802 let mut entries = Vec::with_capacity(count);
1803 for _ in 0..count {
1805 let value = match cur.u8()? {
1806 0 => FrequencyValue::Null,
1807 1 => FrequencyValue::Integer(i128::from_le_bytes(
1808 cur.take(16)?.try_into().expect("sixteen bytes"),
1809 )),
1810 2 => FrequencyValue::Code(cur.u32()?),
1811 _ => return Err(invalid("frequency value tag differs")),
1812 };
1813 let valid = matches!(
1814 (&field.ty, value),
1815 (_, FrequencyValue::Null)
1816 | (LogicalType::Varchar, FrequencyValue::Code(_))
1817 | (
1818 LogicalType::TinyInt
1819 | LogicalType::SmallInt
1820 | LogicalType::Integer
1821 | LogicalType::BigInt
1822 | LogicalType::UTinyInt
1823 | LogicalType::USmallInt
1824 | LogicalType::UInteger
1825 | LogicalType::UBigInt
1826 | LogicalType::Date
1827 | LogicalType::Timestamp,
1828 FrequencyValue::Integer(_),
1829 )
1830 );
1831 if !valid {
1832 return Err(invalid("frequency value does not match its column"));
1833 }
1834 let count = cur.u64()?;
1835 if count == 0 || count > rows as u64 {
1836 return Err(invalid("frequency count is outside the table"));
1837 }
1838 entries.push(FrequencyEntry { value, count });
1839 }
1840 if entries.windows(2).any(|pair| pair[0].count < pair[1].count) {
1841 return Err(invalid("frequency entries are not descending"));
1842 }
1843 let ordinals = if frequency_version == 1 {
1844 Vec::new()
1845 } else {
1846 let ordinal_count = cur.u32()? as usize;
1847 if ordinal_count > FREQUENCY_ORDINALS || ordinal_count > rows {
1848 return Err(invalid("frequency ordinal count exceeds its bound"));
1849 }
1850 let mut ordinals = Vec::with_capacity(ordinal_count);
1851 let mut previous = 0_u64;
1852 for at in 0..ordinal_count {
1853 let delta = cur.var_u64()?;
1854 if at != 0 && delta == 0 {
1855 return Err(invalid("frequency ordinals are not increasing"));
1856 }
1857 let ordinal = if at == 0 {
1858 delta
1859 } else {
1860 previous
1861 .checked_add(delta)
1862 .ok_or_else(|| invalid("frequency ordinal overflows"))?
1863 };
1864 if ordinal >= rows as u64 {
1865 return Err(invalid("frequency ordinal is outside the table"));
1866 }
1867 ordinals.push(ordinal);
1868 previous = ordinal;
1869 }
1870 ordinals
1871 };
1872 Some(FrequencySummary { entries, omitted_max, ordinals })
1873 }
1874 _ => return Err(invalid("frequency summary tag differs")),
1875 };
1876 frequencies.push(summary);
1877 }
1878 frequencies
1879 };
1880 if cur.at != bytes.len() {
1881 return Err(invalid("directory has trailing bytes"));
1882 }
1883 Ok(Table { name, fields, stripes, rows, dictionaries, frequencies, version })
1884}
1885
1886fn put_bound(out: &mut Vec<u8>, bound: Option<&Bound>) -> Result<()> {
1887 match bound {
1888 None => out.push(0),
1889 Some(Bound::Int(value)) => {
1890 out.push(1);
1891 out.extend_from_slice(&value.to_le_bytes());
1892 }
1893 Some(Bound::Real(value)) => {
1894 out.push(2);
1895 out.extend_from_slice(&value.to_le_bytes());
1896 }
1897 Some(Bound::Bytes(value)) => {
1898 out.push(3);
1899 put_u32(out, u32::try_from(value.len()).map_err(|_| invalid("bound length overflow"))?);
1900 out.extend_from_slice(value);
1901 }
1902 }
1903 Ok(())
1904}
1905
1906fn encode(
1907 vector: &Vector,
1908 global: Option<&mut GlobalDictionary>,
1909) -> Result<(Vec<u8>, Option<Vec<u8>>)> {
1910 let ty = vector.logical_type();
1911 let flat = vector.flatten()?;
1913 let mut out = Vec::new();
1914 let mut global_codes = None;
1915 if let Some(global) = global {
1916 let mut codes = Vec::with_capacity(flat.len());
1917 for row in 0..flat.len() {
1918 let text = flat.text_at(row).unwrap_or("");
1919 let code = global.code(text)?;
1920 global.observe(code, flat.is_null_at(row))?;
1921 codes.push(code);
1922 }
1923 global_codes = Some(codes);
1924 }
1925 let membership = global_codes.as_deref().map(encode_membership);
1926 let dictionary = if global_codes.is_none() && ty == &LogicalType::Varchar {
1927 string_dictionary(&flat)?
1928 } else {
1929 None
1930 };
1931 let packed_vector = if dictionary.is_none() && global_codes.is_none() {
1932 Some(flat.bit_packed()?)
1933 } else {
1934 None
1935 };
1936 let packed = packed_vector.as_ref().and_then(Vector::packed_parts);
1937 out.push(if global_codes.is_some() {
1938 3
1939 } else if dictionary.is_some() {
1940 1
1941 } else if packed.is_some() {
1942 2
1943 } else {
1944 0
1945 });
1946 let nulls = flat.validity();
1947 let flag = match nulls {
1948 Validity::AllValid => 0,
1949 Validity::AllInvalid => 1,
1950 Validity::Mask(_) => 2,
1951 };
1952 out.push(flag);
1953 if flag == 2 {
1954 for group in (0..vector.len()).step_by(8) {
1955 let mut bits = 0_u8;
1956 for bit in 0..8 {
1957 if group + bit < vector.len() && !flat.is_null_at(group + bit) {
1958 bits |= 1 << bit;
1959 }
1960 }
1961 out.push(bits);
1962 }
1963 }
1964 if let Some(codes) = global_codes {
1965 for code in codes {
1966 put_u32(&mut out, code);
1967 }
1968 return Ok((out, membership));
1969 }
1970 if let Some(dictionary) = dictionary {
1971 out.extend_from_slice(&dictionary);
1972 return Ok((out, membership));
1973 }
1974 if let Some(packed) = packed {
1975 if packed.offset() != 0 {
1976 return Err(invalid("writer received a sliced packed vector"));
1977 }
1978 out.push(u8::try_from(packed.width()).map_err(|_| invalid("packed width overflow"))?);
1979 out.extend_from_slice(&packed.base().to_le_bytes());
1980 put_u32(
1981 &mut out,
1982 u32::try_from(packed.words().len()).map_err(|_| invalid("too many packed words"))?,
1983 );
1984 for word in packed.words() {
1985 put_u64(&mut out, *word);
1986 }
1987 return Ok((out, membership));
1988 }
1989 let data = flat.data().ok_or_else(|| invalid("scalar column did not flatten"))?;
1990 match (ty, data) {
1991 (LogicalType::TinyInt, Data::Int8(values)) => {
1992 for value in &**values {
1993 out.extend_from_slice(&value.to_le_bytes());
1994 }
1995 }
1996 (LogicalType::UTinyInt, Data::UInt8(values)) => {
1997 for value in &**values {
1998 out.extend_from_slice(&value.to_le_bytes());
1999 }
2000 }
2001 (LogicalType::SmallInt, Data::Int16(values)) => {
2002 for value in &**values {
2003 out.extend_from_slice(&value.to_le_bytes());
2004 }
2005 }
2006 (LogicalType::USmallInt, Data::UInt16(values)) => {
2007 for value in &**values {
2008 out.extend_from_slice(&value.to_le_bytes());
2009 }
2010 }
2011 (LogicalType::UInteger, Data::UInt32(values)) => {
2012 for value in &**values {
2013 out.extend_from_slice(&value.to_le_bytes());
2014 }
2015 }
2016 (LogicalType::UBigInt, Data::UInt64(values)) => {
2017 for value in &**values {
2018 out.extend_from_slice(&value.to_le_bytes());
2019 }
2020 }
2021 (LogicalType::Integer | LogicalType::Date, Data::Int32(values)) => {
2022 for value in &**values {
2023 out.extend_from_slice(&value.to_le_bytes());
2024 }
2025 }
2026 (LogicalType::BigInt | LogicalType::Timestamp, Data::Int64(values)) => {
2027 for value in &**values {
2028 out.extend_from_slice(&value.to_le_bytes());
2029 }
2030 }
2031 (LogicalType::Boolean, Data::Bool(values)) => {
2032 for value in &**values {
2033 out.push(u8::from(*value));
2034 }
2035 }
2036 (LogicalType::Varchar, Data::Varlen(values)) => {
2037 let mut bytes = Vec::new();
2038 put_u32(&mut out, 0);
2039 for row in 0..vector.len() {
2040 let value = values.bytes(row).ok_or_else(|| invalid("string view is invalid"))?;
2041 bytes.extend_from_slice(value);
2042 put_u32(
2043 &mut out,
2044 u32::try_from(bytes.len())
2045 .map_err(|_| invalid("string payload exceeds 4GiB"))?,
2046 );
2047 }
2048 out.extend_from_slice(&bytes);
2049 }
2050 _ => return Err(Error::not_implemented(format!("native page for {ty}"))),
2051 }
2052 Ok((out, membership))
2053}
2054
2055fn put_varint(out: &mut Vec<u8>, mut value: u32) {
2056 while value >= 0x80 {
2057 out.push((value as u8 & 0x7f) | 0x80);
2058 value >>= 7;
2059 }
2060 out.push(value as u8);
2061}
2062
2063fn encode_membership(codes: &[u32]) -> Vec<u8> {
2064 let mut unique = codes.to_vec();
2065 unique.sort_unstable();
2066 unique.dedup();
2067 let mut out = Vec::with_capacity(unique.len().saturating_mul(2).saturating_add(5));
2068 put_varint(&mut out, u32::try_from(unique.len()).unwrap_or(u32::MAX));
2069 let mut previous = 0;
2070 for (at, code) in unique.into_iter().enumerate() {
2071 put_varint(&mut out, if at == 0 { code } else { code - previous });
2072 previous = code;
2073 }
2074 out
2075}
2076
2077fn take_varint(bytes: &[u8], at: &mut usize) -> Result<u32> {
2078 let mut value = 0_u32;
2079 for shift in (0..35).step_by(7) {
2080 let byte = *bytes.get(*at).ok_or_else(|| invalid("membership varint is truncated"))?;
2081 *at += 1;
2082 let part = u32::from(byte & 0x7f);
2083 if shift == 28 && part > 0x0f {
2084 return Err(invalid("membership varint overflow"));
2085 }
2086 value = value
2087 .checked_add(
2088 part.checked_shl(shift).ok_or_else(|| invalid("membership varint overflow"))?,
2089 )
2090 .ok_or_else(|| invalid("membership varint overflow"))?;
2091 if byte & 0x80 == 0 {
2092 return Ok(value);
2093 }
2094 }
2095 Err(invalid("membership varint is too long"))
2096}
2097
2098fn decode_membership(bytes: &[u8]) -> Result<Vec<u32>> {
2099 let mut at = 0;
2100 let count = take_varint(bytes, &mut at)? as usize;
2101 let mut codes = Vec::with_capacity(count);
2102 let mut previous = 0_u32;
2103 for index in 0..count {
2104 let delta = take_varint(bytes, &mut at)?;
2105 let code = if index == 0 {
2106 delta
2107 } else {
2108 previous.checked_add(delta).ok_or_else(|| invalid("membership code overflow"))?
2109 };
2110 if index > 0 && code <= previous {
2111 return Err(invalid("membership codes are not increasing"));
2112 }
2113 codes.push(code);
2114 previous = code;
2115 }
2116 if at != bytes.len() {
2117 return Err(invalid("membership page has trailing bytes"));
2118 }
2119 Ok(codes)
2120}
2121
2122fn string_dictionary(vector: &Vector) -> Result<Option<Vec<u8>>> {
2123 let mut by_text = HashMap::new();
2124 let mut values = Vec::new();
2125 let mut codes = Vec::with_capacity(vector.len());
2126 let mut plain_bytes = 0_usize;
2127 for row in 0..vector.len() {
2128 let text = vector.text_at(row).unwrap_or("");
2129 plain_bytes = plain_bytes.saturating_add(text.len());
2130 let code = match by_text.get(text) {
2131 Some(&code) => code,
2132 None => {
2133 let code = u32::try_from(values.len())
2134 .map_err(|_| invalid("too many dictionary values"))?;
2135 by_text.insert(text, code);
2136 values.push(text);
2137 code
2138 }
2139 };
2140 codes.push(code);
2141 }
2142 let dictionary_bytes = values.iter().map(|value| value.len()).sum::<usize>();
2143 let encoded = 8_usize
2144 .saturating_add((values.len() + 1).saturating_mul(4))
2145 .saturating_add(dictionary_bytes)
2146 .saturating_add(codes.len().saturating_mul(4));
2147 let plain = (vector.len() + 1).saturating_mul(4).saturating_add(plain_bytes);
2148 if encoded >= plain {
2149 return Ok(None);
2150 }
2151 let mut out = Vec::with_capacity(encoded);
2152 put_u32(
2153 &mut out,
2154 u32::try_from(values.len()).map_err(|_| invalid("too many dictionary values"))?,
2155 );
2156 put_u32(
2157 &mut out,
2158 u32::try_from(dictionary_bytes).map_err(|_| invalid("dictionary payload exceeds 4GiB"))?,
2159 );
2160 let mut offset = 0_u32;
2161 put_u32(&mut out, offset);
2162 for value in &values {
2163 offset = offset
2164 .checked_add(
2165 u32::try_from(value.len()).map_err(|_| invalid("dictionary value is too long"))?,
2166 )
2167 .ok_or_else(|| invalid("dictionary payload exceeds 4GiB"))?;
2168 put_u32(&mut out, offset);
2169 }
2170 for value in values {
2171 out.extend_from_slice(value.as_bytes());
2172 }
2173 for code in codes {
2174 put_u32(&mut out, code);
2175 }
2176 Ok(Some(out))
2177}
2178
2179struct EncodedDictionary {
2180 index: Vec<u8>,
2181 ranks: Vec<u8>,
2182 payload: Vec<u8>,
2183}
2184
2185fn head(bytes: &[u8]) -> u64 {
2187 let mut word = [0; 8];
2188 let take = bytes.len().min(8);
2189 word[..take].copy_from_slice(&bytes[..take]);
2190 u64::from_be_bytes(word)
2191}
2192
2193fn rankings(dictionaries: &[Option<GlobalDictionary>]) -> Result<Vec<Vec<(u64, u32)>>> {
2201 let present =
2202 dictionaries.iter().enumerate().filter(|(_, held)| held.is_some()).map(|(at, _)| at);
2203 let present = present.collect::<Vec<_>>();
2204 let mut orders = vec![Vec::new(); dictionaries.len()];
2205 let workers = std::thread::available_parallelism()
2206 .map_or(1, usize::from)
2207 .min(MAX_FREQUENCY_WORKERS)
2208 .min(present.len());
2209 if workers <= 1 {
2210 for at in present {
2211 if let Some(dictionary) = &dictionaries[at] {
2212 orders[at] = dictionary.ranked();
2213 }
2214 }
2215 return Ok(orders);
2216 }
2217 let width = present.len().div_ceil(workers);
2218 let pieces = std::thread::scope(|scope| {
2219 present
2220 .chunks(width)
2221 .map(|columns| {
2222 scope.spawn(|| {
2223 columns
2224 .iter()
2225 .filter_map(|&at| dictionaries[at].as_ref().map(|held| (at, held.ranked())))
2226 .collect::<Vec<_>>()
2227 })
2228 })
2229 .collect::<Vec<_>>()
2230 .into_iter()
2231 .map(|handle| {
2232 handle.join().map_err(|_| Error::internal("a dictionary sort worker panicked"))
2233 })
2234 .collect::<Result<Vec<_>>>()
2235 })?;
2236 for piece in pieces {
2237 for (at, order) in piece {
2238 orders[at] = order;
2239 }
2240 }
2241 Ok(orders)
2242}
2243
2244fn encode_global_dictionary(
2245 dictionary: GlobalDictionary,
2246 order: &[(u64, u32)],
2247) -> Result<EncodedDictionary> {
2248 let values = dictionary.offsets.len() - 1;
2249 if order.len() != values {
2250 return Err(invalid("global dictionary order does not cover its values"));
2251 }
2252 let payload_len = dictionary.payload.len();
2253 let blocks = payload_len.div_ceil(TEXT_PAYLOAD_BLOCK);
2254 let ranks = encode_ranks(order);
2255 let rank_blocks = values.div_ceil(TEXT_RANK_BLOCK);
2256 let mut index = Vec::with_capacity(12 + (values + 1) * 4 + (blocks + rank_blocks) * 8);
2257 put_u32(
2258 &mut index,
2259 u32::try_from(values).map_err(|_| invalid("global dictionary has too many values"))?,
2260 );
2261 put_u32(&mut index, TEXT_PAYLOAD_BLOCK as u32);
2262 put_u32(
2263 &mut index,
2264 u32::try_from(blocks).map_err(|_| invalid("global dictionary has too many blocks"))?,
2265 );
2266 for offset in dictionary.offsets {
2267 put_u32(&mut index, offset);
2268 }
2269 for block in dictionary.payload.chunks(TEXT_PAYLOAD_BLOCK) {
2270 put_u64(&mut index, checksum(block));
2271 }
2272 for block in ranks.chunks(TEXT_RANK_BLOCK * RANK_ENTRY) {
2273 put_u64(&mut index, checksum(block));
2274 }
2275 Ok(EncodedDictionary { index, ranks, payload: dictionary.payload })
2276}
2277
2278fn encode_ranks(order: &[(u64, u32)]) -> Vec<u8> {
2285 let mut out = Vec::with_capacity(order.len() * RANK_ENTRY);
2286 for block in order.chunks(TEXT_RANK_BLOCK) {
2287 for &(head, _) in block {
2288 put_u64(&mut out, head);
2289 }
2290 for &(_, code) in block {
2291 put_u32(&mut out, code);
2292 }
2293 }
2294 out
2295}
2296
2297fn open_global_dictionary(
2298 file: Arc<File>,
2299 page: Page,
2300 ty: &LogicalType,
2301 ordered: bool,
2302) -> Result<Vector> {
2303 if ty != &LogicalType::Varchar {
2304 return Err(invalid("global dictionary belongs to a non-string column"));
2305 }
2306 let mut header = [0; 12];
2307 read_at(&file, page.offset, &mut header)?;
2308 let count = u32::from_le_bytes(header[0..4].try_into().expect("four bytes")) as usize;
2309 let block_size = u32::from_le_bytes(header[4..8].try_into().expect("four bytes")) as usize;
2310 let blocks = u32::from_le_bytes(header[8..12].try_into().expect("four bytes")) as usize;
2311 if block_size != TEXT_PAYLOAD_BLOCK {
2312 return Err(invalid("global dictionary block width differs"));
2313 }
2314 let offset_len = (count + 1)
2315 .checked_mul(4)
2316 .ok_or_else(|| invalid("global dictionary offset count overflow"))?;
2317 let ranks = if ordered { count } else { 0 };
2322 let rank_blocks = ranks.div_ceil(TEXT_RANK_BLOCK);
2323 let rank_len =
2324 ranks.checked_mul(RANK_ENTRY).ok_or_else(|| invalid("global dictionary rank overflow"))?;
2325 let hash_len = blocks
2326 .checked_add(rank_blocks)
2327 .and_then(|count| count.checked_mul(8))
2328 .ok_or_else(|| invalid("global dictionary block count overflow"))?;
2329 let index_len = 12usize
2330 .checked_add(offset_len)
2331 .and_then(|len| len.checked_add(hash_len))
2332 .ok_or_else(|| invalid("global dictionary header overflow"))?;
2333 let body_len = index_len
2334 .checked_add(rank_len)
2335 .ok_or_else(|| invalid("global dictionary header overflow"))?;
2336 if body_len > page.length as usize {
2337 return Err(invalid("global dictionary offset index exceeds its page"));
2338 }
2339 let mut index = vec![0; index_len];
2340 index[..12].copy_from_slice(&header);
2341 read_at(&file, page.offset + 12, &mut index[12..])?;
2342 if checksum(&index) != page.hash {
2343 return Err(invalid("global dictionary index checksum differs"));
2344 }
2345 let offsets = index[12..12 + offset_len]
2346 .chunks_exact(4)
2347 .map(|part| u32::from_le_bytes(part.try_into().expect("four bytes")))
2348 .collect::<Vec<_>>();
2349 let mut hashes = index[12 + offset_len..]
2350 .chunks_exact(8)
2351 .map(|part| u64::from_le_bytes(part.try_into().expect("eight bytes")))
2352 .collect::<Vec<_>>();
2353 let rank_hashes = hashes.split_off(blocks);
2354 let payload_len = page.length as usize - body_len;
2355 if blocks != payload_len.div_ceil(TEXT_PAYLOAD_BLOCK) {
2356 return Err(invalid("global dictionary block count differs from its payload"));
2357 }
2358 if offsets.first() != Some(&0)
2359 || offsets.last().copied().map(|last| last as usize) != Some(payload_len)
2360 || offsets.windows(2).any(|pair| pair[0] > pair[1])
2361 {
2362 return Err(invalid("global dictionary offsets do not bound the payload"));
2363 }
2364 let payload_blocks =
2365 (0..payload_len.div_ceil(TEXT_PAYLOAD_BLOCK)).map(|_| OnceLock::new()).collect();
2366 let crossing = (0..count.div_ceil(TEXT_CROSSING_BLOCK)).map(|_| OnceLock::new()).collect();
2367 Vector::external_text(
2368 LogicalType::Varchar,
2369 Arc::new(NativeText {
2370 file,
2371 offsets,
2372 ranks,
2373 rank_at: page.offset + index_len as u64,
2374 rank_hashes,
2375 rank_blocks: (0..rank_blocks).map(|_| OnceLock::new()).collect(),
2376 payload: page.offset + body_len as u64,
2377 payload_len,
2378 hashes,
2379 payload_blocks,
2380 crossing,
2381 }),
2382 )
2383}
2384
2385fn decode(
2386 ty: &LogicalType,
2387 rows: usize,
2388 bytes: &[u8],
2389 global: Option<Arc<Vector>>,
2390) -> Result<Vector> {
2391 let mut cur = Cursor { bytes, at: 0 };
2392 let codec = cur.u8()?;
2393 let flag = cur.u8()?;
2394 let validity = match flag {
2395 0 => Validity::AllValid,
2396 1 => Validity::AllInvalid,
2397 2 => {
2398 let mask = cur.take(rows.div_ceil(8))?;
2399 Validity::from_iter(rows, |row| mask[row / 8] >> (row % 8) & 1 == 1)
2400 }
2401 _ => return Err(invalid("page validity tag differs")),
2402 };
2403 if codec == 1 {
2404 if ty != &LogicalType::Varchar {
2405 return Err(invalid("dictionary codec belongs to a non-string page"));
2406 }
2407 let count = cur.u32()? as usize;
2408 let payload_len = cur.u32()? as usize;
2409 let offset_bytes = cur.take(
2410 (count + 1)
2411 .checked_mul(4)
2412 .ok_or_else(|| invalid("dictionary offset count overflow"))?,
2413 )?;
2414 let offsets = offset_bytes
2415 .chunks_exact(4)
2416 .map(|part| u32::from_le_bytes(part.try_into().expect("four bytes")))
2417 .collect::<Vec<_>>();
2418 let payload = cur.take(payload_len)?.to_vec();
2419 if offsets.first() != Some(&0)
2420 || offsets.last().copied().map(|last| last as usize) != Some(payload.len())
2421 || offsets.windows(2).any(|pair| pair[0] > pair[1])
2422 {
2423 return Err(invalid("dictionary offsets do not bound the payload"));
2424 }
2425 let mut strings = StringColumn::over(Buffer::from_vec(payload));
2426 for pair in offsets.windows(2) {
2427 strings.push_in_place(pair[0] as usize, (pair[1] - pair[0]) as usize)?;
2428 }
2429 let mut codes = Vec::with_capacity(rows);
2430 for _ in 0..rows {
2431 codes.push(cur.u32()?);
2432 }
2433 if codes.iter().any(|code| *code as usize >= count) {
2434 return Err(invalid("dictionary code is out of range"));
2435 }
2436 if cur.at != bytes.len() {
2437 return Err(invalid("dictionary page has trailing bytes"));
2438 }
2439 let dictionary = Vector::flat(LogicalType::Varchar, Data::Varlen(strings))?;
2440 return Ok(Vector::dictionary(codes, dictionary)?.with_validity(validity));
2441 }
2442 if codec == 3 {
2443 let dictionary = global.ok_or_else(|| invalid("global code page has no dictionary"))?;
2444 let mut codes = Vec::with_capacity(rows);
2445 let mut highest = None;
2446 for _ in 0..rows {
2447 let code = cur.u32()?;
2448 highest = Some(highest.map_or(code, |old: u32| old.max(code)));
2449 codes.push(code);
2450 }
2451 if cur.at != bytes.len() {
2452 return Err(invalid("global code page has trailing bytes"));
2453 }
2454 return Ok(Vector::stable_dictionary_validated(codes, dictionary, highest)?
2455 .with_validity(validity));
2456 }
2457 if codec == 2 {
2458 let width = u32::from(cur.u8()?);
2459 let base = i128::from_le_bytes(cur.take(16)?.try_into().expect("sixteen bytes"));
2460 let count = cur.u32()? as usize;
2461 let mut words = Vec::with_capacity(count);
2462 for _ in 0..count {
2463 words.push(cur.u64()?);
2464 }
2465 if cur.at != bytes.len() {
2466 return Err(invalid("packed page has trailing bytes"));
2467 }
2468 return Ok(Vector::packed(ty.clone(), words, width, base, rows)?.with_validity(validity));
2469 }
2470 if codec != 0 {
2471 return Err(invalid("page codec is unknown"));
2472 }
2473 let data = match ty {
2474 LogicalType::TinyInt => {
2475 let values = cur.take(rows)?;
2476 Data::Int8(values.iter().map(|item| *item as i8).collect::<Vec<_>>().into())
2477 }
2478 LogicalType::UTinyInt => Data::UInt8(cur.take(rows)?.to_vec().into()),
2479 LogicalType::SmallInt => {
2480 let values =
2481 cur.take(rows.checked_mul(2).ok_or_else(|| invalid("page size overflow"))?)?;
2482 Data::Int16(
2483 values
2484 .chunks_exact(2)
2485 .map(|item| i16::from_le_bytes(item.try_into().expect("two bytes")))
2486 .collect::<Vec<_>>()
2487 .into(),
2488 )
2489 }
2490 LogicalType::USmallInt => {
2491 let values =
2492 cur.take(rows.checked_mul(2).ok_or_else(|| invalid("page size overflow"))?)?;
2493 Data::UInt16(
2494 values
2495 .chunks_exact(2)
2496 .map(|item| u16::from_le_bytes(item.try_into().expect("two bytes")))
2497 .collect::<Vec<_>>()
2498 .into(),
2499 )
2500 }
2501 LogicalType::UInteger => {
2502 let values =
2503 cur.take(rows.checked_mul(4).ok_or_else(|| invalid("page size overflow"))?)?;
2504 Data::UInt32(
2505 values
2506 .chunks_exact(4)
2507 .map(|item| u32::from_le_bytes(item.try_into().expect("four bytes")))
2508 .collect::<Vec<_>>()
2509 .into(),
2510 )
2511 }
2512 LogicalType::UBigInt => {
2513 let values =
2514 cur.take(rows.checked_mul(8).ok_or_else(|| invalid("page size overflow"))?)?;
2515 Data::UInt64(
2516 values
2517 .chunks_exact(8)
2518 .map(|item| u64::from_le_bytes(item.try_into().expect("eight bytes")))
2519 .collect::<Vec<_>>()
2520 .into(),
2521 )
2522 }
2523 LogicalType::Integer | LogicalType::Date => {
2524 let values =
2525 cur.take(rows.checked_mul(4).ok_or_else(|| invalid("page size overflow"))?)?;
2526 Data::Int32(
2527 values
2528 .chunks_exact(4)
2529 .map(|item| i32::from_le_bytes(item.try_into().expect("four bytes")))
2530 .collect::<Vec<_>>()
2531 .into(),
2532 )
2533 }
2534 LogicalType::BigInt | LogicalType::Timestamp => {
2535 let values =
2536 cur.take(rows.checked_mul(8).ok_or_else(|| invalid("page size overflow"))?)?;
2537 Data::Int64(
2538 values
2539 .chunks_exact(8)
2540 .map(|item| i64::from_le_bytes(item.try_into().expect("eight bytes")))
2541 .collect::<Vec<_>>()
2542 .into(),
2543 )
2544 }
2545 LogicalType::Boolean => {
2546 let values = cur.take(rows)?;
2547 if values.iter().any(|value| *value > 1) {
2548 return Err(invalid("boolean page has another value"));
2549 }
2550 Data::Bool(values.iter().map(|value| *value == 1).collect::<Vec<_>>().into())
2551 }
2552 LogicalType::Varchar => {
2553 let offset_bytes = cur
2554 .take((rows + 1).checked_mul(4).ok_or_else(|| invalid("offset count overflow"))?)?;
2555 let offsets = offset_bytes
2556 .chunks_exact(4)
2557 .map(|part| u32::from_le_bytes(part.try_into().expect("four bytes")))
2558 .collect::<Vec<_>>();
2559 let payload = cur.take(bytes.len() - cur.at)?.to_vec();
2560 if offsets.first() != Some(&0)
2561 || offsets.last().copied().map(|last| last as usize) != Some(payload.len())
2562 || offsets.windows(2).any(|pair| pair[0] > pair[1])
2563 {
2564 return Err(invalid("string offsets do not bound the payload"));
2565 }
2566 let mut values = StringColumn::over(Buffer::from_vec(payload));
2567 for pair in offsets.windows(2) {
2568 values.push_in_place(pair[0] as usize, (pair[1] - pair[0]) as usize)?;
2569 }
2570 Data::Varlen(values)
2571 }
2572 _ => return Err(Error::not_implemented(format!("native page for {ty}"))),
2573 };
2574 if cur.at != bytes.len() {
2575 return Err(invalid("page has trailing bytes"));
2576 }
2577 Ok(Vector::flat(ty.clone(), data)?.with_validity(validity))
2578}
2579
2580#[cfg(test)]
2581mod tests {
2582 use std::fs;
2583 use std::io::{Seek, SeekFrom, Write};
2584 use std::path::PathBuf;
2585 use std::time::{SystemTime, UNIX_EPOCH};
2586
2587 use rudb_common::Value;
2588 use rudb_common::bounds::Op;
2589
2590 use super::*;
2591
2592 #[test]
2593 fn checksum_matches_fixed_vectors() {
2594 assert_eq!(checksum(b""), 0xef46_db37_51d8_e999);
2595 assert_eq!(checksum(b"a"), 0xd24e_c4f1_a98c_6e5b);
2596 assert_eq!(checksum(b"abc"), 0x44bc_2cf5_ad77_0999);
2597 }
2598
2599 fn path(label: &str) -> PathBuf {
2600 let stamp = SystemTime::now().duration_since(UNIX_EPOCH).expect("time advances").as_nanos();
2601 std::env::temp_dir().join(format!("rudb-native-{label}-{}-{stamp}.rdb", std::process::id()))
2602 }
2603
2604 fn sample() -> Chunk {
2605 Chunk::new(vec![
2606 Vector::from_values(
2607 LogicalType::Integer,
2608 &[Value::Integer(4), Value::Integer(9), Value::Integer(-2)],
2609 )
2610 .expect("integers"),
2611 Vector::from_values(
2612 LogicalType::Varchar,
2613 &[
2614 Value::Varchar("alpha".into()),
2615 Value::Null,
2616 Value::Varchar("long text after a slash".into()),
2617 ],
2618 )
2619 .expect("strings"),
2620 ])
2621 .expect("matching rows")
2622 }
2623
2624 #[test]
2625 fn committed_file_reopens_and_reads_only_requested_columns() {
2626 let path = path("reopen");
2627 let mut writer = Writer::create(
2628 &path,
2629 "items",
2630 vec![
2631 Field::required("id", LogicalType::Integer),
2632 Field::new("text", LogicalType::Varchar),
2633 ],
2634 )
2635 .expect("new file");
2636 writer.append(&sample()).expect("first stripe");
2637 writer.append(&sample()).expect("second stripe");
2638 writer.finish().expect("commit");
2639 let reader = Reader::open(&path).expect("reopen from disk");
2640 assert_eq!(reader.table().rows(), 6);
2641 assert_eq!(reader.table().stripes().len(), 2);
2642 let text = reader.read(1, &[1]).expect("only text page");
2643 assert_eq!(text.width(), 1);
2644 assert_eq!(text.value_at(1, 0), Value::Null);
2645 assert_eq!(text.value_at(2, 0), Value::Varchar("long text after a slash".into()));
2646 let sparse = reader.read_sparse(1, &[1]).expect("one page without extent prefetch");
2647 assert_eq!(sparse.width(), 1);
2648 assert_eq!(sparse.value_at(1, 0), Value::Null);
2649 assert_eq!(sparse.value_at(2, 0), Value::Varchar("long text after a slash".into()));
2650 assert!(!reader.skips_codes(0, 1, &[0]).expect("alpha is in the stripe"));
2651 assert!(!reader.skips_codes(0, 1, &[2]).expect("long text is in the stripe"));
2652 assert!(reader.skips_codes(0, 1, &[3]).expect("unknown code is absent"));
2653 let count = reader.read(0, &[]).expect("no page is needed for count");
2654 assert_eq!(count.len(), 3);
2655 assert!(reader.skips(0, &[Probe { column: 0, op: Op::Greater, value: Bound::Int(100) }]));
2656 assert!(!reader.skips(0, &[Probe { column: 0, op: Op::Greater, value: Bound::Int(0) }]));
2657 let integers = reader.top_frequencies(0, 1).expect("valid integer synopsis").expect("kept");
2658 assert_eq!(
2659 integers,
2660 vec![(Value::Integer(-2), 2), (Value::Integer(4), 2), (Value::Integer(9), 2),]
2661 );
2662 let strings = reader.top_frequencies(1, 1).expect("valid string synopsis").expect("kept");
2663 assert_eq!(strings.len(), 3);
2664 assert!(strings.contains(&(Value::Null, 2)));
2665 assert!(strings.contains(&(Value::Varchar("alpha".into()), 2)));
2666 assert!(strings.contains(&(Value::Varchar("long text after a slash".into()), 2)));
2667 fs::remove_file(path).expect("remove scratch file");
2668 }
2669
2670 #[test]
2677 fn every_integer_width_round_trips_through_a_page() {
2678 let path = path("integer-widths");
2679 let columns = [
2680 (LogicalType::TinyInt, vec![Value::TinyInt(i8::MIN), Value::TinyInt(i8::MAX)]),
2681 (LogicalType::UTinyInt, vec![Value::UTinyInt(0), Value::UTinyInt(u8::MAX)]),
2682 (LogicalType::SmallInt, vec![Value::SmallInt(i16::MIN), Value::SmallInt(i16::MAX)]),
2683 (LogicalType::USmallInt, vec![Value::USmallInt(0), Value::USmallInt(u16::MAX)]),
2684 (LogicalType::Integer, vec![Value::Integer(i32::MIN), Value::Integer(i32::MAX)]),
2685 (LogicalType::UInteger, vec![Value::UInteger(0), Value::UInteger(u32::MAX)]),
2686 (LogicalType::BigInt, vec![Value::BigInt(i64::MIN), Value::BigInt(i64::MAX)]),
2687 (LogicalType::UBigInt, vec![Value::UBigInt(0), Value::UBigInt(u64::MAX)]),
2688 ];
2689 let fields = columns
2690 .iter()
2691 .enumerate()
2692 .map(|(at, (ty, _))| Field::required(format!("c{at}"), ty.clone()))
2693 .collect::<Vec<_>>();
2694 let vectors = columns
2695 .iter()
2696 .map(|(ty, values)| Vector::from_values(ty.clone(), values).expect("a vector"))
2697 .collect::<Vec<_>>();
2698 let mut writer = Writer::create(&path, "widths", fields).expect("new file");
2699 writer.append(&Chunk::new(vectors).expect("matching rows")).expect("one stripe");
2700 writer.finish().expect("commit");
2701
2702 let reader = Reader::open(&path).expect("reopen from disk");
2703 let wanted = (0..columns.len()).collect::<Vec<_>>();
2704 let read = reader.read(0, &wanted).expect("every column");
2705 assert_eq!(read.len(), 2);
2706 for (at, (ty, values)) in columns.iter().enumerate() {
2708 assert_eq!(read.value_at(0, at), values[0], "the low end of {ty}");
2709 assert_eq!(read.value_at(1, at), values[1], "the high end of {ty}");
2710 }
2711 fs::remove_file(path).expect("remove scratch file");
2712 }
2713
2714 #[test]
2715 fn numeric_frequency_candidates_keep_bounded_row_ordinals() {
2716 let path = path("frequency-ordinals");
2717 let mut writer =
2718 Writer::create(&path, "items", vec![Field::required("id", LogicalType::BigInt)])
2719 .expect("new file");
2720 let mut values = Vec::new();
2721 for leader in 0..10_i64 {
2722 values.extend(std::iter::repeat_n(leader, 100));
2723 }
2724 values.extend(1_000_i64..41_000);
2725 for part in values.chunks(1_024) {
2726 let vector = Vector::flat(LogicalType::BigInt, Data::Int64(part.to_vec().into()))
2727 .expect("big integers");
2728 writer.append(&Chunk::new(vec![vector]).expect("one column")).expect("one stripe");
2729 }
2730 writer.finish().expect("commit");
2731
2732 let reader = Reader::open(&path).expect("reopen from disk");
2733 let occurrences =
2734 reader.frequency_occurrences(0).expect("valid metadata").expect("bounded ordinals");
2735 assert!(occurrences.omitted_max < 100);
2736 assert!(occurrences.ordinals.len() <= FREQUENCY_ORDINALS);
2737 assert!(occurrences.ordinals.windows(2).all(|pair| pair[0] < pair[1]));
2738 assert_eq!(&occurrences.ordinals[..1_000], &(0_u64..1_000).collect::<Vec<_>>());
2739 fs::remove_file(path).expect("remove scratch file");
2740 }
2741
2742 #[test]
2743 fn an_unfinished_or_damaged_file_does_not_answer_with_partial_rows() {
2744 let unfinished = path("unfinished");
2745 let mut writer =
2746 Writer::create(&unfinished, "items", vec![Field::new("id", LogicalType::Integer)])
2747 .expect("new file");
2748 let chunk = Chunk::new(vec![
2749 Vector::flat(LogicalType::Integer, Data::Int32(vec![1, 2, 3].into()))
2750 .expect("integers"),
2751 ])
2752 .expect("chunk");
2753 writer.append(&chunk).expect("page written");
2754 drop(writer);
2755 assert!(Reader::open(&unfinished).is_err(), "no directory was committed");
2756 fs::remove_file(unfinished).expect("remove scratch file");
2757
2758 let damaged = path("damaged");
2759 let mut writer =
2760 Writer::create(&damaged, "items", vec![Field::new("id", LogicalType::Integer)])
2761 .expect("new file");
2762 writer.append(&chunk).expect("page written");
2763 writer.finish().expect("commit");
2764 let reader = Reader::open(&damaged).expect("valid directory");
2765 let mut file =
2766 OpenOptions::new().write(true).open(&damaged).expect("open for a damaged page");
2767 file.seek(SeekFrom::Start(HEADER + 1)).expect("inside first page");
2768 file.write_all(&[255]).expect("damage one byte");
2769 assert!(reader.read(0, &[0]).is_err(), "page checksum rejects corruption");
2770 fs::remove_file(damaged).expect("remove scratch file");
2771 }
2772
2773 #[test]
2774 fn damaged_lazy_dictionary_payload_is_an_error() {
2775 let path = path("damaged-dictionary");
2776 let mut writer = Writer::create(
2777 &path,
2778 "items",
2779 vec![
2780 Field::required("id", LogicalType::Integer),
2781 Field::new("text", LogicalType::Varchar),
2782 ],
2783 )
2784 .expect("new file");
2785 writer.append(&sample()).expect("stripe written");
2786 writer.finish().expect("commit");
2787
2788 let reader = Reader::open(&path).expect("valid directory");
2789 let dictionary = reader.table.dictionaries[1].expect("string dictionary page");
2790 let mut header = [0; 12];
2793 read_at(&reader.file, dictionary.offset, &mut header).expect("dictionary header");
2794 let count = u64::from(u32::from_le_bytes(header[0..4].try_into().expect("four bytes")));
2795 let blocks = u64::from(u32::from_le_bytes(header[8..12].try_into().expect("four bytes")));
2796 let rank_blocks = count.div_ceil(TEXT_RANK_BLOCK as u64);
2797 let index_len =
2798 12 + (count + 1) * 4 + (blocks + rank_blocks) * 8 + count * RANK_ENTRY as u64;
2799 let mut file = OpenOptions::new().write(true).open(&path).expect("open dictionary page");
2800 file.seek(SeekFrom::Start(dictionary.offset + index_len))
2801 .expect("inside dictionary payload");
2802 file.write_all(&[255]).expect("damage dictionary payload");
2803
2804 let chunk = reader.read(0, &[1]).expect("code page and dictionary index remain valid");
2805 let error =
2806 chunk.validate_external().expect_err("payload corruption must reach the caller");
2807 assert!(error.message().contains("payload checksum differs"), "{error}");
2808 fs::remove_file(path).expect("remove scratch file");
2809 }
2810
2811 #[test]
2816 fn a_damaged_sorted_order_is_an_error() {
2817 let path = path("damaged-order");
2818 let mut writer = Writer::create(
2819 &path,
2820 "items",
2821 vec![
2822 Field::required("id", LogicalType::Integer),
2823 Field::new("text", LogicalType::Varchar),
2824 ],
2825 )
2826 .expect("new file");
2827 writer.append(&sample()).expect("stripe written");
2828 writer.finish().expect("commit");
2829
2830 let reader = Reader::open(&path).expect("valid directory");
2831 let page = reader.table.dictionaries[1].expect("string dictionary page");
2832 let mut header = [0; 12];
2833 read_at(&reader.file, page.offset, &mut header).expect("dictionary header");
2834 let count = u64::from(u32::from_le_bytes(header[0..4].try_into().expect("four bytes")));
2835 let blocks = u64::from(u32::from_le_bytes(header[8..12].try_into().expect("four bytes")));
2836 let rank_blocks = count.div_ceil(TEXT_RANK_BLOCK as u64);
2837 let index_len = 12 + (count + 1) * 4 + (blocks + rank_blocks) * 8;
2838 let mut file = OpenOptions::new().write(true).open(&path).expect("open dictionary page");
2839 file.seek(SeekFrom::Start(page.offset + index_len)).expect("the first head");
2840 file.write_all(&[255]).expect("damage the order");
2841
2842 let dictionary = reader.dictionary(1).expect("read").expect("a string column has one");
2843 let error = dictionary.compare_rank(0, b"anything").expect_err("a damaged order is caught");
2844 assert!(error.message().contains("rank checksum differs"), "{error}");
2845 fs::remove_file(path).expect("remove scratch file");
2846 }
2847
2848 #[test]
2852 fn a_global_dictionary_carries_the_sorted_order_of_its_values() {
2853 let spellings = ["overlong1z", "b", "", "overlong1a", "overlong", "ab", "a", "overlong1"];
2856 let path = path("dictionary-order");
2857 let mut writer =
2858 Writer::create(&path, "items", vec![Field::new("text", LogicalType::Varchar)])
2859 .expect("new file");
2860 writer
2861 .append(
2862 &Chunk::new(vec![
2863 Vector::from_values(
2864 LogicalType::Varchar,
2865 &spellings.map(|text| Value::Varchar(text.into())),
2866 )
2867 .expect("strings"),
2868 ])
2869 .expect("one column"),
2870 )
2871 .expect("stripe written");
2872 writer.finish().expect("commit");
2873
2874 let reader = Reader::open(&path).expect("valid directory");
2875 let dictionary = reader.dictionary(0).expect("read").expect("a string column has one");
2876 let count = dictionary.ranks().expect("a v9 file stores one");
2877 assert_eq!(count, spellings.len(), "every distinct value has a rank");
2878 let order = (0..count)
2879 .map(|rank| dictionary.code_at_rank(rank).expect("a code"))
2880 .collect::<Vec<_>>();
2881 let mut seen = order.clone();
2882 seen.sort_unstable();
2883 assert_eq!(seen, (0..spellings.len() as u32).collect::<Vec<_>>(), "a permutation of codes");
2884
2885 let ranked = order
2886 .iter()
2887 .map(|&code| {
2888 dictionary.try_bytes_at(code as usize).expect("read").expect("a value").to_vec()
2889 })
2890 .collect::<Vec<_>>();
2891 let mut expected = spellings.map(|text| text.as_bytes().to_vec()).to_vec();
2892 expected.sort();
2893 assert_eq!(ranked, expected, "rank order is value order");
2894
2895 for (rank, value) in expected.iter().enumerate() {
2898 assert_eq!(
2899 dictionary.compare_rank(rank, value).expect("compare"),
2900 Ordering::Equal,
2901 "rank {rank} is its own value"
2902 );
2903 if rank > 0 {
2904 assert_eq!(
2905 dictionary.compare_rank(rank - 1, value).expect("compare"),
2906 Ordering::Less,
2907 "rank {rank} follows the one before it"
2908 );
2909 }
2910 }
2911 fs::remove_file(path).expect("remove scratch file");
2912 }
2913
2914 #[test]
2915 fn damaged_membership_cannot_skip_a_string_page() {
2916 let path = path("damaged-membership");
2917 let mut writer = Writer::create(
2918 &path,
2919 "items",
2920 vec![
2921 Field::required("id", LogicalType::Integer),
2922 Field::new("text", LogicalType::Varchar),
2923 ],
2924 )
2925 .expect("new file");
2926 writer.append(&sample()).expect("stripe written");
2927 writer.finish().expect("commit");
2928
2929 let reader = Reader::open(&path).expect("valid directory");
2930 let membership = reader.table.stripes[0].memberships[1].expect("string membership");
2931 let mut file = OpenOptions::new().write(true).open(&path).expect("open membership page");
2932 file.seek(SeekFrom::Start(membership.offset)).expect("membership start");
2933 file.write_all(&[255]).expect("damage membership");
2934 let error = reader.skips_codes(0, 1, &[3]).expect_err("corruption must not skip rows");
2935 assert!(error.message().contains("membership page checksum differs"), "{error}");
2936 fs::remove_file(path).expect("remove scratch file");
2937 }
2938
2939 #[test]
2940 fn membership_delta_stream_is_sorted_exact_and_bounded() {
2941 let encoded = encode_membership(&[900, 4, 4, 72, 9, u32::MAX]);
2942 assert_eq!(
2943 decode_membership(&encoded).expect("valid membership"),
2944 [4, 9, 72, 900, u32::MAX]
2945 );
2946 assert!(decode_membership(&[1, 0x80]).is_err(), "a truncated varint is invalid");
2947 assert!(
2948 decode_membership(&[1, 0xff, 0xff, 0xff, 0xff, 0x10]).is_err(),
2949 "a value past u32 is invalid"
2950 );
2951 }
2952
2953 #[test]
2954 fn a_global_dictionary_may_be_larger_than_one_column_page() {
2955 let dictionary = Page {
2956 offset: HEADER,
2957 length: u32::try_from(MAX_PAGE + 1).expect("the page bound fits on disk"),
2958 hash: 0,
2959 };
2960 let table = Table {
2961 name: "items".to_owned(),
2962 fields: vec![Field::new("text", LogicalType::Varchar)],
2963 stripes: Vec::new(),
2964 rows: 0,
2965 dictionaries: vec![Some(dictionary)],
2966 frequencies: vec![None],
2967 version: FORMAT,
2968 };
2969 let directory = encode_directory(&table).expect("directory");
2970 let file_size = dictionary.offset + u64::from(dictionary.length) + 1;
2971
2972 let decoded =
2973 decode_directory(&directory, file_size, FORMAT).expect("large lazy dictionary");
2974 assert_eq!(decoded.dictionaries[0].expect("dictionary").length, dictionary.length);
2975 for old in [8, 7] {
2976 let legacy = encode_directory_version(&table, old).expect("legacy directory");
2977 let decoded =
2978 decode_directory(&legacy, file_size, old).expect("an older directory still reads");
2979 assert_eq!(
2980 decoded.dictionaries[0].expect("legacy dictionary").length,
2981 dictionary.length
2982 );
2983 }
2984 }
2985}