1#[cfg(test)]
22mod dictionary_config_tests;
23
24use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
25use parking_lot::RwLock;
26use rustc_hash::FxHashMap;
27use std::io::{self, Read, Write};
28use std::sync::Arc;
29
30#[cfg(feature = "fst-index")]
31use super::sstable_index::FstBlockIndex;
32use super::sstable_index::{BlockAddr, BlockIndex, MmapBlockIndex};
33use super::vint::{read_vint, write_vint};
34use crate::compression::{CompressionDict, CompressionLevel};
35use crate::directories::{FileHandle, OwnedBytes};
36
37pub const SSTABLE_MAGIC: u32 = 0x53544235; pub const RESTART_INTERVAL: usize = 16;
45
46pub const BLOCK_SIZE: usize = 16 * 1024;
48
49#[derive(Debug, Clone, Copy, PartialEq, Eq)]
52pub struct SSTableBlockSize(usize);
53
54impl SSTableBlockSize {
55 pub fn bytes(self) -> usize {
56 self.0
57 }
58}
59
60impl std::str::FromStr for SSTableBlockSize {
61 type Err = io::Error;
62
63 fn from_str(value: &str) -> io::Result<Self> {
64 let bytes = value
65 .parse::<usize>()
66 .map_err(|error| io::Error::new(io::ErrorKind::InvalidInput, error))?;
67 Self::try_from(bytes)
68 }
69}
70
71impl std::fmt::Display for SSTableBlockSize {
72 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
73 self.bytes().fmt(formatter)
74 }
75}
76
77impl Default for SSTableBlockSize {
78 fn default() -> Self {
79 Self(BLOCK_SIZE)
80 }
81}
82
83impl TryFrom<usize> for SSTableBlockSize {
84 type Error = io::Error;
85
86 fn try_from(bytes: usize) -> io::Result<Self> {
87 if !(512..=1024 * 1024).contains(&bytes) {
88 return Err(io::Error::new(
89 io::ErrorKind::InvalidInput,
90 "SSTable block target must be in 512..=1048576 bytes",
91 ));
92 }
93 Ok(Self(bytes))
94 }
95}
96
97pub const DEFAULT_DICT_SIZE: usize = 64 * 1024;
99
100pub const BLOOM_BITS_PER_KEY: usize = 10;
102
103pub const BLOOM_HASH_COUNT: usize = 7;
105const BLOOM_FILTER_HEADER_SIZE: usize = 16;
106
107const MAX_SSTABLE_BLOCK_BYTES: usize = 64 * 1024 * 1024;
108const MAX_SSTABLE_DICTIONARY_BYTES: u64 = 16 * 1024 * 1024;
109
110pub type PrefixScanResult<V> = (Vec<(Vec<u8>, V)>, bool);
112
113#[derive(Debug, Clone)]
119pub struct BloomFilter {
120 bits: BloomBits,
121 num_bits: usize,
122 num_hashes: usize,
123}
124
125#[derive(Debug, Clone)]
127enum BloomBits {
128 Vec(Vec<u64>),
130 Bytes(OwnedBytes),
132}
133
134impl BloomBits {
135 #[inline]
136 fn len(&self) -> usize {
137 match self {
138 BloomBits::Vec(v) => v.len(),
139 BloomBits::Bytes(b) => b.len() / 8,
140 }
141 }
142
143 #[inline]
144 fn get(&self, word_idx: usize) -> u64 {
145 match self {
146 BloomBits::Vec(v) => v[word_idx],
147 BloomBits::Bytes(b) => {
148 let off = word_idx * 8;
149 u64::from_le_bytes([
150 b[off],
151 b[off + 1],
152 b[off + 2],
153 b[off + 3],
154 b[off + 4],
155 b[off + 5],
156 b[off + 6],
157 b[off + 7],
158 ])
159 }
160 }
161 }
162
163 #[inline]
164 fn set_bit(&mut self, word_idx: usize, bit_idx: usize) {
165 match self {
166 BloomBits::Vec(v) => v[word_idx] |= 1u64 << bit_idx,
167 BloomBits::Bytes(_) => panic!("cannot mutate read-only bloom filter"),
168 }
169 }
170
171 fn size_bytes(&self) -> usize {
172 match self {
173 BloomBits::Vec(v) => v.len() * 8,
174 BloomBits::Bytes(b) => b.len(),
175 }
176 }
177
178 fn write_to(&self, writer: &mut (impl Write + ?Sized)) -> io::Result<()> {
179 match self {
180 BloomBits::Vec(words) => {
181 #[cfg(target_endian = "little")]
182 {
183 let bytes = unsafe {
186 std::slice::from_raw_parts(
187 words.as_ptr().cast::<u8>(),
188 words.len().saturating_mul(8),
189 )
190 };
191 writer.write_all(bytes)
192 }
193 #[cfg(target_endian = "big")]
194 {
195 for &word in words {
196 writer.write_u64::<LittleEndian>(word)?;
197 }
198 Ok(())
199 }
200 }
201 BloomBits::Bytes(bytes) => writer.write_all(bytes.as_slice()),
202 }
203 }
204}
205
206impl BloomFilter {
207 pub(crate) const SERIALIZED_HEADER_SIZE: usize = BLOOM_FILTER_HEADER_SIZE;
208
209 pub fn new(expected_keys: usize, bits_per_key: usize) -> Self {
211 let num_bits = expected_keys.saturating_mul(bits_per_key).max(64);
212 let num_words = num_bits.div_ceil(64);
213 Self {
214 bits: BloomBits::Vec(vec![0u64; num_words]),
215 num_bits,
216 num_hashes: BLOOM_HASH_COUNT,
217 }
218 }
219
220 pub fn from_bytes_mutable(data: &[u8]) -> io::Result<Self> {
224 if data.len() < BLOOM_FILTER_HEADER_SIZE {
225 return Err(io::Error::new(
226 io::ErrorKind::InvalidData,
227 "Bloom filter data too short",
228 ));
229 }
230 let num_bits = usize::try_from(u64::from_le_bytes(data[0..8].try_into().unwrap()))
231 .map_err(|_| {
232 io::Error::new(
233 io::ErrorKind::InvalidData,
234 "Bloom filter bit count exceeds addressable memory",
235 )
236 })?;
237 let num_hashes = u32::from_le_bytes(data[8..12].try_into().unwrap()) as usize;
238 let num_words = u32::from_le_bytes(data[12..16].try_into().unwrap()) as usize;
239
240 validate_bloom_header(data.len(), num_bits, num_hashes, num_words)?;
241 let expected_len = BLOOM_FILTER_HEADER_SIZE + num_words * 8;
242 if data.len() != expected_len {
243 return Err(io::Error::new(
244 io::ErrorKind::InvalidData,
245 "Bloom filter data truncated",
246 ));
247 }
248
249 let mut vec = vec![0u64; num_words];
250 for (i, v) in vec.iter_mut().enumerate() {
251 let off = BLOOM_FILTER_HEADER_SIZE + i * 8;
252 *v = u64::from_le_bytes(data[off..off + 8].try_into().unwrap());
253 }
254
255 Ok(Self {
256 bits: BloomBits::Vec(vec),
257 num_bits,
258 num_hashes,
259 })
260 }
261
262 pub fn from_owned_bytes(data: OwnedBytes) -> io::Result<Self> {
264 if data.len() < BLOOM_FILTER_HEADER_SIZE {
265 return Err(io::Error::new(
266 io::ErrorKind::InvalidData,
267 "Bloom filter data too short",
268 ));
269 }
270 let d = data.as_slice();
271 let num_bits =
272 usize::try_from(u64::from_le_bytes(d[0..8].try_into().unwrap())).map_err(|_| {
273 io::Error::new(
274 io::ErrorKind::InvalidData,
275 "Bloom filter bit count exceeds addressable memory",
276 )
277 })?;
278 let num_hashes = u32::from_le_bytes(d[8..12].try_into().unwrap()) as usize;
279 let num_words = u32::from_le_bytes(d[12..16].try_into().unwrap()) as usize;
280
281 validate_bloom_header(d.len(), num_bits, num_hashes, num_words)?;
282 let expected_len = BLOOM_FILTER_HEADER_SIZE + num_words * 8;
283 if d.len() != expected_len {
284 return Err(io::Error::new(
285 io::ErrorKind::InvalidData,
286 "Bloom filter data truncated",
287 ));
288 }
289
290 let bits_bytes =
292 data.slice(BLOOM_FILTER_HEADER_SIZE..BLOOM_FILTER_HEADER_SIZE + num_words * 8);
293
294 Ok(Self {
295 bits: BloomBits::Bytes(bits_bytes),
296 num_bits,
297 num_hashes,
298 })
299 }
300
301 pub fn serialized_len(&self) -> usize {
303 BLOOM_FILTER_HEADER_SIZE + self.bits.len() * 8
304 }
305
306 pub fn write_to(&self, writer: &mut (impl Write + ?Sized)) -> io::Result<()> {
308 let num_words = self.bits.len();
309 write_bloom_header(writer, self.num_bits, self.num_hashes, num_words)?;
310 self.bits.write_to(writer)
311 }
312
313 pub fn to_bytes(&self) -> Vec<u8> {
315 let mut data = Vec::with_capacity(self.serialized_len());
316 self.write_to(&mut data)
317 .expect("writing a bloom filter to Vec cannot fail");
318 data
319 }
320
321 pub fn insert(&mut self, key: &[u8]) {
323 let (h1, h2) = bloom_hash_pair(key);
324 self.insert_hashed(h1, h2);
325 }
326
327 pub fn may_contain(&self, key: &[u8]) -> bool {
330 let (h1, h2) = bloom_hash_pair(key);
331 for i in 0..self.num_hashes {
332 let bit_pos = self.get_bit_pos(h1, h2, i);
333 let word_idx = bit_pos / 64;
334 let bit_idx = bit_pos % 64;
335 if word_idx >= self.bits.len() || (self.bits.get(word_idx) & (1u64 << bit_idx)) == 0 {
336 return false;
337 }
338 }
339 true
340 }
341
342 pub fn size_bytes(&self) -> usize {
344 BLOOM_FILTER_HEADER_SIZE + self.bits.size_bytes()
345 }
346
347 pub fn insert_hashed(&mut self, h1: u64, h2: u64) {
349 for i in 0..self.num_hashes {
350 let bit_pos = self.get_bit_pos(h1, h2, i);
351 let word_idx = bit_pos / 64;
352 let bit_idx = bit_pos % 64;
353 if word_idx < self.bits.len() {
354 self.bits.set_bit(word_idx, bit_idx);
355 }
356 }
357 }
358
359 #[inline]
361 fn get_bit_pos(&self, h1: u64, h2: u64, i: usize) -> usize {
362 (h1.wrapping_add((i as u64).wrapping_mul(h2)) % (self.num_bits as u64)) as usize
363 }
364}
365
366fn write_bloom_header(
367 writer: &mut (impl Write + ?Sized),
368 num_bits: usize,
369 num_hashes: usize,
370 num_words: usize,
371) -> io::Result<()> {
372 writer.write_u64::<LittleEndian>(u64::try_from(num_bits).map_err(|_| {
373 io::Error::new(
374 io::ErrorKind::InvalidInput,
375 "Bloom filter bit count exceeds u64",
376 )
377 })?)?;
378 writer.write_u32::<LittleEndian>(u32::try_from(num_hashes).map_err(|_| {
379 io::Error::new(
380 io::ErrorKind::InvalidInput,
381 "Bloom filter hash count exceeds u32",
382 )
383 })?)?;
384 writer.write_u32::<LittleEndian>(u32::try_from(num_words).map_err(|_| {
385 io::Error::new(
386 io::ErrorKind::InvalidInput,
387 "Bloom filter word count exceeds u32",
388 )
389 })?)?;
390 Ok(())
391}
392
393fn validate_bloom_header(
394 data_len: usize,
395 num_bits: usize,
396 num_hashes: usize,
397 num_words: usize,
398) -> io::Result<()> {
399 if num_bits == 0 || num_hashes == 0 || num_hashes > 32 || num_words == 0 {
400 return Err(io::Error::new(
401 io::ErrorKind::InvalidData,
402 "invalid bloom filter parameters",
403 ));
404 }
405 let word_bytes = num_words
406 .checked_mul(8)
407 .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "bloom filter size overflow"))?;
408 let expected_len = BLOOM_FILTER_HEADER_SIZE
409 .checked_add(word_bytes)
410 .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "bloom filter size overflow"))?;
411 let capacity_bits = num_words
412 .checked_mul(64)
413 .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "bloom filter size overflow"))?;
414 if expected_len > data_len
415 || num_bits > capacity_bits
416 || num_bits <= capacity_bits.saturating_sub(64)
417 {
418 return Err(io::Error::new(
419 io::ErrorKind::InvalidData,
420 "inconsistent bloom filter dimensions",
421 ));
422 }
423 Ok(())
424}
425
426#[inline]
429fn bloom_hash_pair(key: &[u8]) -> (u64, u64) {
430 let mut h1: u64 = 0xcbf29ce484222325;
431 let mut h2: u64 = 0x84222325cbf29ce4;
432 for &byte in key {
433 h1 ^= byte as u64;
434 h1 = h1.wrapping_mul(0x100000001b3);
435 h2 = h2.wrapping_mul(0x100000001b3);
436 h2 ^= byte as u64;
437 }
438 (h1, h2)
439}
440
441pub trait SSTableValue: Clone + Send + Sync {
448 fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()>;
450
451 fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self>;
453}
454
455impl SSTableValue for u64 {
457 fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
458 write_vint(writer, *self)
459 }
460
461 fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
462 read_vint(reader)
463 }
464}
465
466impl SSTableValue for Vec<u8> {
468 fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
469 write_vint(writer, self.len() as u64)?;
470 writer.write_all(self)
471 }
472
473 fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
474 let len = usize::try_from(read_vint(reader)?).map_err(|_| {
475 io::Error::new(io::ErrorKind::InvalidData, "SSTable value length too large")
476 })?;
477 if len > MAX_SSTABLE_BLOCK_BYTES {
478 return Err(io::Error::new(
479 io::ErrorKind::InvalidData,
480 "SSTable value exceeds block safety limit",
481 ));
482 }
483 let mut data = vec![0u8; len];
484 reader.read_exact(&mut data)?;
485 Ok(data)
486 }
487}
488
489#[derive(Debug, Clone, Copy, PartialEq, Eq)]
492pub struct SparseDimInfo {
493 pub offset: u64,
495 pub length: u32,
497}
498
499impl SparseDimInfo {
500 pub fn new(offset: u64, length: u32) -> Self {
501 Self { offset, length }
502 }
503}
504
505impl SSTableValue for SparseDimInfo {
506 fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
507 write_vint(writer, self.offset)?;
508 write_vint(writer, self.length as u64)
509 }
510
511 fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
512 let offset = read_vint(reader)?;
513 let length = u32::try_from(read_vint(reader)?).map_err(|_| {
514 io::Error::new(
515 io::ErrorKind::InvalidData,
516 "sparse posting length exceeds u32",
517 )
518 })?;
519 Ok(Self { offset, length })
520 }
521}
522
523pub const MAX_INLINE_POSTINGS: usize = 3;
525
526#[derive(Clone, Debug)]
527pub(crate) struct DecodedInlinePostings {
528 docs: [u32; MAX_INLINE_POSTINGS],
529 frequencies: [u32; MAX_INLINE_POSTINGS],
530 len: usize,
531}
532
533impl DecodedInlinePostings {
534 pub(crate) fn docs(&self) -> &[u32] {
535 &self.docs[..self.len]
536 }
537 pub(crate) fn frequencies(&self) -> &[u32] {
538 &self.frequencies[..self.len]
539 }
540}
541
542#[derive(Debug, Clone, PartialEq, Eq)]
550pub enum TermInfo {
551 Inline {
554 doc_freq: u8,
556 data: [u8; 16],
559 data_len: u8,
561 },
562 External {
564 posting_offset: u64,
565 posting_len: u64,
566 doc_freq: u32,
567 position_offset: u64,
569 position_len: u64,
571 },
572}
573
574impl TermInfo {
575 pub fn external(posting_offset: u64, posting_len: u64, doc_freq: u32) -> Self {
577 TermInfo::External {
578 posting_offset,
579 posting_len,
580 doc_freq,
581 position_offset: 0,
582 position_len: 0,
583 }
584 }
585
586 pub fn external_with_positions(
588 posting_offset: u64,
589 posting_len: u64,
590 doc_freq: u32,
591 position_offset: u64,
592 position_len: u64,
593 ) -> Self {
594 TermInfo::External {
595 posting_offset,
596 posting_len,
597 doc_freq,
598 position_offset,
599 position_len,
600 }
601 }
602
603 pub fn try_inline(doc_ids: &[u32], term_freqs: &[u32]) -> Option<Self> {
606 if doc_ids.len() > MAX_INLINE_POSTINGS
607 || doc_ids.is_empty()
608 || doc_ids.len() != term_freqs.len()
609 {
610 return None;
611 }
612
613 let mut data = [0u8; 16];
614 let mut cursor = std::io::Cursor::new(&mut data[..]);
615 let mut prev_doc_id = 0u32;
616
617 for (i, &doc_id) in doc_ids.iter().enumerate() {
618 let delta = doc_id.checked_sub(prev_doc_id)?;
619 if write_vint(&mut cursor, delta as u64).is_err() {
620 return None;
621 }
622 if write_vint(&mut cursor, term_freqs[i] as u64).is_err() {
623 return None;
624 }
625 prev_doc_id = doc_id;
626 }
627
628 let data_len = cursor.position() as u8;
629 if data_len > 16 {
630 return None;
631 }
632
633 Some(TermInfo::Inline {
634 doc_freq: doc_ids.len() as u8,
635 data,
636 data_len,
637 })
638 }
639
640 pub fn try_inline_iter(count: usize, iter: impl Iterator<Item = (u32, u32)>) -> Option<Self> {
644 if count > MAX_INLINE_POSTINGS || count == 0 {
645 return None;
646 }
647
648 let mut data = [0u8; 16];
649 let mut cursor = std::io::Cursor::new(&mut data[..]);
650 let mut prev_doc_id = 0u32;
651
652 let mut actual_count = 0usize;
653 for (doc_id, tf) in iter {
654 if actual_count >= count {
655 return None;
656 }
657 let delta = doc_id.checked_sub(prev_doc_id)?;
658 if write_vint(&mut cursor, delta as u64).is_err() {
659 return None;
660 }
661 if write_vint(&mut cursor, tf as u64).is_err() {
662 return None;
663 }
664 prev_doc_id = doc_id;
665 actual_count += 1;
666 }
667
668 if actual_count != count {
669 return None;
670 }
671
672 let data_len = cursor.position() as u8;
673
674 Some(TermInfo::Inline {
675 doc_freq: count as u8,
676 data,
677 data_len,
678 })
679 }
680
681 pub fn doc_freq(&self) -> u32 {
683 match self {
684 TermInfo::Inline { doc_freq, .. } => *doc_freq as u32,
685 TermInfo::External { doc_freq, .. } => *doc_freq,
686 }
687 }
688
689 pub fn is_inline(&self) -> bool {
691 matches!(self, TermInfo::Inline { .. })
692 }
693
694 pub fn external_info(&self) -> Option<(u64, u64)> {
696 match self {
697 TermInfo::External {
698 posting_offset,
699 posting_len,
700 ..
701 } => Some((*posting_offset, *posting_len)),
702 TermInfo::Inline { .. } => None,
703 }
704 }
705
706 pub fn position_info(&self) -> Option<(u64, u64)> {
708 match self {
709 TermInfo::External {
710 position_offset,
711 position_len,
712 ..
713 } if *position_len > 0 => Some((*position_offset, *position_len)),
714 _ => None,
715 }
716 }
717
718 pub fn decode_inline(&self) -> Option<(Vec<u32>, Vec<u32>)> {
721 self.decode_inline_fixed()
722 .map(|decoded| (decoded.docs().to_vec(), decoded.frequencies().to_vec()))
723 }
724
725 pub(crate) fn decode_inline_fixed(&self) -> Option<DecodedInlinePostings> {
726 match self {
727 TermInfo::Inline {
728 doc_freq,
729 data,
730 data_len,
731 } => {
732 if *doc_freq == 0
733 || *doc_freq as usize > MAX_INLINE_POSTINGS
734 || *data_len as usize > data.len()
735 {
736 return None;
737 }
738 let mut decoded = DecodedInlinePostings {
739 docs: [0; MAX_INLINE_POSTINGS],
740 frequencies: [0; MAX_INLINE_POSTINGS],
741 len: usize::from(*doc_freq),
742 };
743 let mut reader = &data[..*data_len as usize];
744 let mut prev_doc_id = 0u32;
745
746 for i in 0..decoded.len {
747 let delta = u32::try_from(read_vint(&mut reader).ok()?).ok()?;
748 let tf = u32::try_from(read_vint(&mut reader).ok()?).ok()?;
749 let doc_id = prev_doc_id.checked_add(delta)?;
750 decoded.docs[i] = doc_id;
751 decoded.frequencies[i] = tf;
752 prev_doc_id = doc_id;
753 }
754
755 Some(decoded)
756 }
757 TermInfo::External { .. } => None,
758 }
759 }
760}
761
762impl SSTableValue for TermInfo {
763 fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
764 match self {
765 TermInfo::Inline {
766 doc_freq,
767 data,
768 data_len,
769 } => {
770 if *doc_freq == 0
771 || *doc_freq as usize > MAX_INLINE_POSTINGS
772 || *data_len as usize > data.len()
773 {
774 return Err(io::Error::new(
775 io::ErrorKind::InvalidInput,
776 "invalid inline TermInfo",
777 ));
778 }
779 writer.write_u8(0xFF)?;
781 writer.write_u8(*doc_freq)?;
782 writer.write_u8(*data_len)?;
783 writer.write_all(&data[..*data_len as usize])?;
784 }
785 TermInfo::External {
786 posting_offset,
787 posting_len,
788 doc_freq,
789 position_offset,
790 position_len,
791 } => {
792 if *position_len > 0 {
795 writer.write_u8(0x01)?;
796 write_vint(writer, *doc_freq as u64)?;
797 write_vint(writer, *posting_offset)?;
798 write_vint(writer, *posting_len)?;
799 write_vint(writer, *position_offset)?;
800 write_vint(writer, *position_len)?;
801 } else {
802 writer.write_u8(0x00)?;
803 write_vint(writer, *doc_freq as u64)?;
804 write_vint(writer, *posting_offset)?;
805 write_vint(writer, *posting_len)?;
806 }
807 }
808 }
809 Ok(())
810 }
811
812 fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
813 let tag = reader.read_u8()?;
814
815 if tag == 0xFF {
816 let doc_freq = reader.read_u8()?;
818 let data_len = reader.read_u8()?;
819 if doc_freq == 0 || doc_freq as usize > MAX_INLINE_POSTINGS || data_len as usize > 16 {
820 return Err(io::Error::new(
821 io::ErrorKind::InvalidData,
822 "invalid inline TermInfo lengths",
823 ));
824 }
825 let mut data = [0u8; 16];
826 reader.read_exact(&mut data[..data_len as usize])?;
827 Ok(TermInfo::Inline {
828 doc_freq,
829 data,
830 data_len,
831 })
832 } else if tag == 0x00 {
833 let doc_freq = read_vint(reader)? as u32;
835 let posting_offset = read_vint(reader)?;
836 let posting_len = read_vint(reader)?;
837 Ok(TermInfo::External {
838 posting_offset,
839 posting_len,
840 doc_freq,
841 position_offset: 0,
842 position_len: 0,
843 })
844 } else if tag == 0x01 {
845 let doc_freq = read_vint(reader)? as u32;
847 let posting_offset = read_vint(reader)?;
848 let posting_len = read_vint(reader)?;
849 let position_offset = read_vint(reader)?;
850 let position_len = read_vint(reader)?;
851 Ok(TermInfo::External {
852 posting_offset,
853 posting_len,
854 doc_freq,
855 position_offset,
856 position_len,
857 })
858 } else {
859 Err(io::Error::new(
860 io::ErrorKind::InvalidData,
861 format!("Invalid TermInfo tag: {}", tag),
862 ))
863 }
864 }
865}
866
867pub fn common_prefix_len(a: &[u8], b: &[u8]) -> usize {
869 a.iter().zip(b.iter()).take_while(|(x, y)| x == y).count()
870}
871
872fn decode_block_entry<V: SSTableValue>(
878 reader: &mut &[u8],
879 current_key: &mut Vec<u8>,
880) -> io::Result<V> {
881 let common_prefix_len = read_vint(reader)? as usize;
882 let suffix_len = read_vint(reader)? as usize;
883
884 if suffix_len > reader.len() {
885 return Err(io::Error::new(
886 io::ErrorKind::UnexpectedEof,
887 "SSTable block suffix truncated",
888 ));
889 }
890
891 current_key.truncate(common_prefix_len);
892 current_key.extend_from_slice(&reader[..suffix_len]);
893 *reader = &reader[suffix_len..];
894
895 V::deserialize(reader)
896}
897
898#[derive(Debug, Clone)]
900pub struct SSTableStats {
901 pub num_blocks: usize,
903 pub num_sparse_entries: usize,
905 pub num_entries: u64,
907 pub has_bloom_filter: bool,
909 pub has_dictionary: bool,
911 pub bloom_filter_size: usize,
913 pub dictionary_size: usize,
915 pub cache_insert_bypasses: u64,
920}
921
922#[derive(Debug, Clone)]
924pub struct SSTableWriterConfig {
925 pub block_size: SSTableBlockSize,
927 pub compression_level: CompressionLevel,
929 pub use_dictionary: bool,
931 pub dict_size: usize,
933 pub use_bloom_filter: bool,
935 pub bloom_bits_per_key: usize,
937}
938
939impl Default for SSTableWriterConfig {
940 fn default() -> Self {
941 Self::from_optimization(crate::structures::IndexOptimization::default())
942 }
943}
944
945impl SSTableWriterConfig {
946 pub fn from_optimization(optimization: crate::structures::IndexOptimization) -> Self {
948 use crate::structures::IndexOptimization;
949 match optimization {
950 IndexOptimization::Adaptive => Self {
951 block_size: SSTableBlockSize::default(),
952 compression_level: CompressionLevel::BETTER, use_dictionary: false,
954 dict_size: DEFAULT_DICT_SIZE,
955 use_bloom_filter: true, bloom_bits_per_key: BLOOM_BITS_PER_KEY,
957 },
958 IndexOptimization::SizeOptimized => Self {
959 block_size: SSTableBlockSize::default(),
960 compression_level: CompressionLevel::MAX, use_dictionary: true,
962 dict_size: DEFAULT_DICT_SIZE,
963 use_bloom_filter: true,
964 bloom_bits_per_key: BLOOM_BITS_PER_KEY,
965 },
966 IndexOptimization::PerformanceOptimized => Self {
967 block_size: SSTableBlockSize::default(),
968 compression_level: CompressionLevel::FAST, use_dictionary: false,
970 dict_size: DEFAULT_DICT_SIZE,
971 use_bloom_filter: true, bloom_bits_per_key: BLOOM_BITS_PER_KEY,
973 },
974 }
975 }
976
977 pub fn fast() -> Self {
979 Self::from_optimization(crate::structures::IndexOptimization::PerformanceOptimized)
980 }
981
982 pub fn max_compression() -> Self {
984 Self::from_optimization(crate::structures::IndexOptimization::SizeOptimized)
985 }
986}
987
988pub struct SSTableWriter<W: Write, V: SSTableValue> {
994 writer: W,
995 block_buffer: Vec<u8>,
996 prev_key: Vec<u8>,
997 index: Vec<BlockIndexEntry>,
998 current_offset: u64,
999 num_entries: u64,
1000 block_first_key: Option<Vec<u8>>,
1001 config: SSTableWriterConfig,
1002 dictionary: Option<CompressionDict>,
1004 bloom_hashes: Vec<(u64, u64)>,
1007 block_restarts: Vec<u32>,
1010 block_entry_count: usize,
1012 failed: bool,
1014 _phantom: std::marker::PhantomData<V>,
1015}
1016
1017struct BlockEntryWriter<'a> {
1020 buffer: &'a mut Vec<u8>,
1021 limit: usize,
1022}
1023
1024impl Write for BlockEntryWriter<'_> {
1025 fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
1026 let needed = self
1027 .buffer
1028 .len()
1029 .checked_add(bytes.len())
1030 .filter(|&n| n <= self.limit)
1031 .ok_or_else(|| {
1032 io::Error::new(
1033 io::ErrorKind::InvalidInput,
1034 "SSTable entry exceeds the reader block safety limit",
1035 )
1036 })?;
1037 if needed > self.buffer.capacity() {
1038 let capacity = needed
1040 .max(self.buffer.capacity().saturating_mul(2))
1041 .min(self.limit);
1042 self.buffer
1043 .try_reserve_exact(capacity - self.buffer.len())
1044 .map_err(io::Error::other)?;
1045 }
1046 self.buffer.extend_from_slice(bytes);
1047 Ok(bytes.len())
1048 }
1049
1050 fn flush(&mut self) -> io::Result<()> {
1051 Ok(())
1052 }
1053}
1054
1055impl<W: Write, V: SSTableValue> SSTableWriter<W, V> {
1056 pub fn new(writer: W) -> Self {
1058 Self::with_config(writer, SSTableWriterConfig::default())
1059 }
1060
1061 pub fn with_config(writer: W, config: SSTableWriterConfig) -> Self {
1063 Self {
1064 writer,
1065 block_buffer: Vec::with_capacity(config.block_size.bytes()),
1066 prev_key: Vec::new(),
1067 index: Vec::new(),
1068 current_offset: 0,
1069 num_entries: 0,
1070 block_first_key: None,
1071 config,
1072 dictionary: None,
1073 bloom_hashes: Vec::new(),
1074 block_restarts: Vec::new(),
1075 block_entry_count: 0,
1076 failed: false,
1077 _phantom: std::marker::PhantomData,
1078 }
1079 }
1080
1081 pub fn with_dictionary(
1083 writer: W,
1084 config: SSTableWriterConfig,
1085 dictionary: CompressionDict,
1086 ) -> Self {
1087 Self {
1088 writer,
1089 block_buffer: Vec::with_capacity(config.block_size.bytes()),
1090 prev_key: Vec::new(),
1091 index: Vec::new(),
1092 current_offset: 0,
1093 num_entries: 0,
1094 block_first_key: None,
1095 config,
1096 dictionary: Some(dictionary),
1097 bloom_hashes: Vec::new(),
1098 block_restarts: Vec::new(),
1099 block_entry_count: 0,
1100 failed: false,
1101 _phantom: std::marker::PhantomData,
1102 }
1103 }
1104
1105 pub fn insert(&mut self, key: &[u8], value: &V) -> io::Result<()> {
1106 if self.failed {
1107 return Err(io::Error::other(
1108 "SSTable writer failed during an earlier insert",
1109 ));
1110 }
1111 self.failed = true;
1113 self.insert_entry(key, value)?;
1114 self.failed = false;
1115 Ok(())
1116 }
1117
1118 fn insert_entry(&mut self, key: &[u8], value: &V) -> io::Result<()> {
1119 if key.len() > MAX_SSTABLE_BLOCK_BYTES || self.block_buffer.len() > MAX_SSTABLE_BLOCK_BYTES
1120 {
1121 return Err(io::Error::new(
1122 io::ErrorKind::InvalidInput,
1123 "SSTable entry exceeds the reader block safety limit",
1124 ));
1125 }
1126 if self.block_first_key.is_none() {
1127 self.block_first_key = Some(key.to_vec());
1128 }
1129
1130 if self.config.use_bloom_filter {
1132 self.bloom_hashes.push(bloom_hash_pair(key));
1133 }
1134
1135 let restart = self.block_entry_count.is_multiple_of(RESTART_INTERVAL);
1138 let prefix_len = if restart {
1139 self.block_restarts.push(self.block_buffer.len() as u32);
1140 0
1141 } else {
1142 common_prefix_len(&self.prev_key, key)
1143 };
1144 let suffix = &key[prefix_len..];
1145
1146 let trailer_bytes = self.block_restarts.len() * 4 + 4;
1147 let mut entry = BlockEntryWriter {
1148 buffer: &mut self.block_buffer,
1149 limit: MAX_SSTABLE_BLOCK_BYTES - trailer_bytes,
1150 };
1151 write_vint(&mut entry, prefix_len as u64)?;
1152 write_vint(&mut entry, suffix.len() as u64)?;
1153 entry.write_all(suffix)?;
1154 value.serialize(&mut entry)?;
1155
1156 self.prev_key.clear();
1157 self.prev_key.extend_from_slice(key);
1158 self.num_entries += 1;
1159 self.block_entry_count += 1;
1160
1161 if self.block_buffer.len() >= self.config.block_size.bytes() {
1162 self.flush_block()?;
1163 }
1164
1165 Ok(())
1166 }
1167
1168 fn flush_block(&mut self) -> io::Result<()> {
1170 if self.block_buffer.is_empty() {
1171 return Ok(());
1172 }
1173
1174 let trailer_bytes = self
1175 .block_restarts
1176 .len()
1177 .checked_mul(4)
1178 .and_then(|bytes| bytes.checked_add(4));
1179 if trailer_bytes
1180 .and_then(|bytes| self.block_buffer.len().checked_add(bytes))
1181 .is_none_or(|bytes| bytes > MAX_SSTABLE_BLOCK_BYTES)
1182 {
1183 return Err(io::Error::new(
1184 io::ErrorKind::InvalidInput,
1185 "SSTable block including restart trailer exceeds the reader safety limit",
1186 ));
1187 }
1188 for offset in &self.block_restarts {
1190 self.block_buffer.extend_from_slice(&offset.to_le_bytes());
1191 }
1192 self.block_buffer
1193 .extend_from_slice(&(self.block_restarts.len() as u32).to_le_bytes());
1194 self.block_restarts.clear();
1195 self.block_entry_count = 0;
1196
1197 let compressed = if let Some(ref dict) = self.dictionary {
1199 crate::compression::compress_with_dict(
1200 &self.block_buffer,
1201 self.config.compression_level,
1202 dict,
1203 )?
1204 } else {
1205 crate::compression::compress(&self.block_buffer, self.config.compression_level)?
1206 };
1207
1208 if let Some(first_key) = self.block_first_key.take() {
1209 self.index.push(BlockIndexEntry {
1210 first_key,
1211 offset: self.current_offset,
1212 length: compressed.len() as u32,
1213 });
1214 }
1215
1216 self.writer.write_all(&compressed)?;
1217 self.current_offset += compressed.len() as u64;
1218 self.block_buffer.clear();
1219 self.prev_key.clear();
1220
1221 Ok(())
1222 }
1223
1224 pub fn finish(mut self) -> io::Result<W> {
1225 if self.failed {
1226 return Err(io::Error::other("cannot finish a failed SSTable writer"));
1227 }
1228 self.flush_block()?;
1230
1231 let bloom_filter = if self.config.use_bloom_filter && !self.bloom_hashes.is_empty() {
1233 let mut bloom =
1234 BloomFilter::new(self.bloom_hashes.len(), self.config.bloom_bits_per_key);
1235 for (h1, h2) in &self.bloom_hashes {
1236 bloom.insert_hashed(*h1, *h2);
1237 }
1238 Some(bloom)
1239 } else {
1240 None
1241 };
1242
1243 let data_end_offset = self.current_offset;
1244
1245 let entries: Vec<(Vec<u8>, BlockAddr)> = self
1248 .index
1249 .iter()
1250 .map(|e| {
1251 (
1252 e.first_key.clone(),
1253 BlockAddr {
1254 offset: e.offset,
1255 length: e.length,
1256 },
1257 )
1258 })
1259 .collect();
1260
1261 #[cfg(feature = "native")]
1263 let index_bytes = FstBlockIndex::build(&entries)?;
1264 #[cfg(not(feature = "native"))]
1265 let index_bytes = MmapBlockIndex::build(&entries)?;
1266
1267 self.writer
1269 .write_u32::<LittleEndian>(index_bytes.len() as u32)?;
1270 self.writer.write_all(&index_bytes)?;
1271 self.current_offset += 4 + index_bytes.len() as u64;
1272
1273 let bloom_offset = if let Some(ref bloom) = bloom_filter {
1275 let offset = self.current_offset;
1276 bloom.write_to(&mut self.writer)?;
1277 self.current_offset += bloom.serialized_len() as u64;
1278 offset
1279 } else {
1280 0
1281 };
1282
1283 let dict_offset = if let Some(ref dict) = self.dictionary {
1285 let dict_bytes = dict.as_bytes();
1286 let offset = self.current_offset;
1287 self.writer
1288 .write_u32::<LittleEndian>(dict_bytes.len() as u32)?;
1289 self.writer.write_all(dict_bytes)?;
1290 self.current_offset += 4 + dict_bytes.len() as u64;
1291 offset
1292 } else {
1293 0
1294 };
1295
1296 self.writer.write_u64::<LittleEndian>(data_end_offset)?;
1298 self.writer.write_u64::<LittleEndian>(self.num_entries)?;
1299 self.writer.write_u64::<LittleEndian>(bloom_offset)?; self.writer.write_u64::<LittleEndian>(dict_offset)?; self.writer
1302 .write_u8(self.config.compression_level.0 as u8)?;
1303 self.writer.write_u32::<LittleEndian>(SSTABLE_MAGIC)?;
1304
1305 Ok(self.writer)
1306 }
1307}
1308
1309#[derive(Debug, Clone)]
1311struct BlockIndexEntry {
1312 first_key: Vec<u8>,
1313 offset: u64,
1314 length: u32,
1315}
1316
1317pub struct AsyncSSTableReader<V: SSTableValue> {
1324 data_slice: FileHandle,
1326 block_index: BlockIndex,
1328 num_entries: u64,
1329 cache: RwLock<BlockCache>,
1331 bloom_filter: Option<BloomFilter>,
1333 dictionary: Option<CompressionDict>,
1335 _phantom: std::marker::PhantomData<V>,
1336}
1337
1338struct BlockParts<'b> {
1340 entries: &'b [u8],
1341 restart_table: &'b [u8],
1343}
1344
1345impl<'b> BlockParts<'b> {
1346 fn split(block: &'b [u8]) -> io::Result<Self> {
1347 if block.len() < 4 {
1348 return Err(io::Error::new(
1349 io::ErrorKind::InvalidData,
1350 "SSTable block shorter than its restart trailer",
1351 ));
1352 }
1353 let count_at = block.len() - 4;
1354 let count = u32::from_le_bytes(block[count_at..].try_into().unwrap()) as usize;
1355 let table_len = count.checked_mul(4).ok_or_else(|| {
1356 io::Error::new(io::ErrorKind::InvalidData, "SSTable restart table overflow")
1357 })?;
1358 let table_at = count_at.checked_sub(table_len).ok_or_else(|| {
1359 io::Error::new(
1360 io::ErrorKind::InvalidData,
1361 "SSTable restart table exceeds its block",
1362 )
1363 })?;
1364 Ok(Self {
1365 entries: &block[..table_at],
1366 restart_table: &block[table_at..count_at],
1367 })
1368 }
1369
1370 #[inline]
1371 fn num_restarts(&self) -> usize {
1372 self.restart_table.len() / 4
1373 }
1374
1375 #[inline]
1376 fn restart_offset(&self, i: usize) -> io::Result<usize> {
1377 let at = i * 4;
1378 let offset =
1379 u32::from_le_bytes(self.restart_table[at..at + 4].try_into().unwrap()) as usize;
1380 if offset >= self.entries.len() {
1381 return Err(io::Error::new(
1382 io::ErrorKind::InvalidData,
1383 "SSTable restart offset outside the block",
1384 ));
1385 }
1386 Ok(offset)
1387 }
1388
1389 fn restart_key(&self, i: usize) -> io::Result<&'b [u8]> {
1391 let offset = self.restart_offset(i)?;
1392 let mut reader = &self.entries[offset..];
1393 let prefix_len = read_vint(&mut reader)?;
1394 if prefix_len != 0 {
1395 return Err(io::Error::new(
1396 io::ErrorKind::InvalidData,
1397 "SSTable restart entry has a non-zero key prefix",
1398 ));
1399 }
1400 let suffix_len = read_vint(&mut reader)? as usize;
1401 if suffix_len > reader.len() {
1402 return Err(io::Error::new(
1403 io::ErrorKind::UnexpectedEof,
1404 "SSTable restart key truncated",
1405 ));
1406 }
1407 Ok(&reader[..suffix_len])
1408 }
1409}
1410
1411struct BlockCache {
1417 blocks: FxHashMap<u64, Arc<[u8]>>,
1418 lru_order: std::collections::VecDeque<u64>,
1419 max_blocks: usize,
1420 max_bytes: Option<usize>,
1421 retained_bytes: usize,
1422 insert_bypasses: u64,
1425}
1426
1427pub const MAX_CACHE_BLOCKS: usize = 65_536;
1433
1434impl BlockCache {
1435 fn new(max_blocks: usize, max_bytes: Option<usize>) -> Self {
1436 let max_blocks = if max_bytes == Some(0) { 0 } else { max_blocks };
1437 if max_blocks > MAX_CACHE_BLOCKS {
1438 log::warn!(
1439 "SSTable block cache cap {max_blocks} exceeds the {MAX_CACHE_BLOCKS} block \
1440 maximum; clamping"
1441 );
1442 }
1443 let max_blocks = max_blocks.min(MAX_CACHE_BLOCKS);
1444 Self {
1445 blocks: FxHashMap::default(),
1446 lru_order: std::collections::VecDeque::with_capacity(
1447 max_blocks.min(max_bytes.unwrap_or(usize::MAX)),
1448 ),
1449 max_blocks,
1450 max_bytes,
1451 retained_bytes: 0,
1452 insert_bypasses: 0,
1453 }
1454 }
1455
1456 fn peek(&self, offset: u64) -> Option<Arc<[u8]>> {
1458 self.blocks.get(&offset).map(Arc::clone)
1459 }
1460
1461 fn insert(&mut self, offset: u64, block: Arc<[u8]>) {
1462 if self.max_blocks == 0 {
1463 self.insert_bypasses += 1;
1464 return;
1465 }
1466 if self.blocks.contains_key(&offset) {
1467 self.promote(offset);
1468 return;
1469 }
1470 if self.max_bytes.is_some_and(|budget| block.len() > budget) {
1471 self.insert_bypasses += 1;
1472 return;
1473 }
1474 while self.blocks.len() >= self.max_blocks
1475 || self
1476 .max_bytes
1477 .is_some_and(|budget| self.retained_bytes > budget - block.len())
1478 {
1479 let evict_offset = self
1480 .lru_order
1481 .pop_front()
1482 .expect("cache order missing block");
1483 let removed = self
1484 .blocks
1485 .remove(&evict_offset)
1486 .expect("cache block missing from order");
1487 self.retained_bytes -= removed.len();
1488 }
1489 self.retained_bytes += block.len();
1490 self.blocks.insert(offset, block);
1491 self.lru_order.push_back(offset);
1492 }
1493
1494 fn promote(&mut self, offset: u64) {
1496 if let Some(pos) = self.lru_order.iter().position(|&k| k == offset) {
1497 self.lru_order.remove(pos);
1498 self.lru_order.push_back(offset);
1499 }
1500 }
1501}
1502
1503impl<V: SSTableValue> AsyncSSTableReader<V> {
1504 pub const PREFETCH_LEADING_MAX_BYTES: u64 = 4 * 1024 * 1024;
1508
1509 pub const PREFETCH_LEADING_FANOUT: usize = 4;
1513
1514 pub async fn open(file_handle: FileHandle, cache_blocks: usize) -> io::Result<Self> {
1519 Self::open_with_cache_budget(file_handle, cache_blocks, None).await
1520 }
1521
1522 pub async fn open_with_cache_budget(
1526 file_handle: FileHandle,
1527 cache_blocks: usize,
1528 cache_budget_bytes: Option<usize>,
1529 ) -> io::Result<Self> {
1530 let file_len = file_handle.len();
1531 if file_len < 37 {
1532 return Err(io::Error::new(
1533 io::ErrorKind::InvalidData,
1534 "SSTable too small",
1535 ));
1536 }
1537
1538 let footer_bytes = file_handle
1541 .read_bytes_range(file_len - 37..file_len)
1542 .await?;
1543
1544 let mut reader = footer_bytes.as_slice();
1545 let data_end_offset = reader.read_u64::<LittleEndian>()?;
1546 let num_entries = reader.read_u64::<LittleEndian>()?;
1547 let bloom_offset = reader.read_u64::<LittleEndian>()?;
1548 let dict_offset = reader.read_u64::<LittleEndian>()?;
1549 let _compression_level = reader.read_u8()?;
1552 let magic = reader.read_u32::<LittleEndian>()?;
1553
1554 if magic != SSTABLE_MAGIC {
1555 return Err(io::Error::new(
1556 io::ErrorKind::InvalidData,
1557 format!(
1558 "Invalid SSTable magic: 0x{magic:08X} (required STB5); the term dictionary \
1559 was written by an incompatible Summa"
1560 ),
1561 ));
1562 }
1563
1564 let footer_start = file_len - 37;
1565 if data_end_offset > footer_start {
1566 return Err(io::Error::new(
1567 io::ErrorKind::InvalidData,
1568 "SSTable data section extends past its footer",
1569 ));
1570 }
1571 if bloom_offset != 0 && (bloom_offset < data_end_offset || bloom_offset >= footer_start) {
1572 return Err(io::Error::new(
1573 io::ErrorKind::InvalidData,
1574 "SSTable bloom filter offset is out of bounds",
1575 ));
1576 }
1577 if dict_offset != 0
1578 && (dict_offset < data_end_offset
1579 || dict_offset >= footer_start
1580 || (bloom_offset != 0 && dict_offset <= bloom_offset))
1581 {
1582 return Err(io::Error::new(
1583 io::ErrorKind::InvalidData,
1584 "SSTable dictionary offset is out of bounds",
1585 ));
1586 }
1587
1588 let index_start = data_end_offset;
1590 let index_end = if bloom_offset != 0 {
1591 bloom_offset
1592 } else if dict_offset != 0 {
1593 dict_offset
1594 } else {
1595 footer_start
1596 };
1597 let index_bytes = file_handle.read_bytes_range(index_start..index_end).await?;
1598
1599 let mut idx_reader = index_bytes.as_slice();
1601 let index_len = idx_reader.read_u32::<LittleEndian>()? as usize;
1602
1603 if index_len != idx_reader.len() {
1604 return Err(io::Error::new(
1605 io::ErrorKind::InvalidData,
1606 "Index data truncated",
1607 ));
1608 }
1609
1610 let index_data = index_bytes.slice(4..4 + index_len);
1611
1612 #[cfg(feature = "fst-index")]
1614 let block_index = match FstBlockIndex::load(index_data.clone()) {
1615 Ok(fst_idx) => BlockIndex::Fst(fst_idx),
1616 Err(_) => BlockIndex::Mmap(MmapBlockIndex::load(index_data)?),
1617 };
1618 #[cfg(not(feature = "fst-index"))]
1619 let block_index = BlockIndex::Mmap(MmapBlockIndex::load(index_data)?);
1620
1621 let mut expected_offset = 0u64;
1622 for addr in block_index.all_addrs() {
1623 let end = addr.offset.checked_add(addr.length as u64).ok_or_else(|| {
1624 io::Error::new(io::ErrorKind::InvalidData, "SSTable block range overflow")
1625 })?;
1626 if addr.length == 0 || addr.offset != expected_offset || end > data_end_offset {
1627 return Err(io::Error::new(
1628 io::ErrorKind::InvalidData,
1629 "SSTable block addresses are inconsistent",
1630 ));
1631 }
1632 expected_offset = end;
1633 }
1634 if expected_offset != data_end_offset {
1635 return Err(io::Error::new(
1636 io::ErrorKind::InvalidData,
1637 "SSTable block addresses do not cover the data section",
1638 ));
1639 }
1640
1641 let bloom_filter = if bloom_offset > 0 {
1643 let bloom_start = bloom_offset;
1644 let bloom_end = if dict_offset != 0 {
1645 dict_offset
1646 } else {
1647 footer_start
1648 };
1649 let header_size = BloomFilter::SERIALIZED_HEADER_SIZE as u64;
1651 let bloom_header_end = bloom_start.checked_add(header_size).ok_or_else(|| {
1652 io::Error::new(io::ErrorKind::InvalidData, "bloom filter range overflow")
1653 })?;
1654 if bloom_header_end > bloom_end {
1655 return Err(io::Error::new(
1656 io::ErrorKind::UnexpectedEof,
1657 "bloom filter header is truncated",
1658 ));
1659 }
1660 let bloom_header = file_handle
1661 .read_bytes_range(bloom_start..bloom_header_end)
1662 .await?;
1663 let num_words = u32::from_le_bytes([
1664 bloom_header[12],
1665 bloom_header[13],
1666 bloom_header[14],
1667 bloom_header[15],
1668 ]) as u64;
1669 let bloom_size = num_words
1670 .checked_mul(8)
1671 .and_then(|bytes| bytes.checked_add(header_size))
1672 .ok_or_else(|| {
1673 io::Error::new(io::ErrorKind::InvalidData, "bloom filter size overflow")
1674 })?;
1675 let actual_bloom_size = bloom_end - bloom_start;
1676 if bloom_size != actual_bloom_size {
1677 return Err(io::Error::new(
1678 io::ErrorKind::InvalidData,
1679 "bloom filter length is inconsistent",
1680 ));
1681 }
1682 let bloom_data = file_handle.read_bytes_range(bloom_start..bloom_end).await?;
1683 Some(BloomFilter::from_owned_bytes(bloom_data)?)
1684 } else {
1685 None
1686 };
1687
1688 let dictionary = if dict_offset > 0 {
1690 let dict_start = dict_offset;
1691 let dict_header_end = dict_start.checked_add(4).ok_or_else(|| {
1693 io::Error::new(io::ErrorKind::InvalidData, "dictionary range overflow")
1694 })?;
1695 if dict_header_end > footer_start {
1696 return Err(io::Error::new(
1697 io::ErrorKind::UnexpectedEof,
1698 "dictionary header is truncated",
1699 ));
1700 }
1701 let dict_len_bytes = file_handle
1702 .read_bytes_range(dict_start..dict_header_end)
1703 .await?;
1704 let dict_len = u32::from_le_bytes([
1705 dict_len_bytes[0],
1706 dict_len_bytes[1],
1707 dict_len_bytes[2],
1708 dict_len_bytes[3],
1709 ]) as u64;
1710 if dict_len > MAX_SSTABLE_DICTIONARY_BYTES {
1711 return Err(io::Error::new(
1712 io::ErrorKind::InvalidData,
1713 "SSTable dictionary exceeds safety limit",
1714 ));
1715 }
1716 let dict_end = dict_header_end.checked_add(dict_len).ok_or_else(|| {
1717 io::Error::new(io::ErrorKind::InvalidData, "dictionary range overflow")
1718 })?;
1719 if dict_end != footer_start {
1720 return Err(io::Error::new(
1721 io::ErrorKind::InvalidData,
1722 "SSTable dictionary length is inconsistent",
1723 ));
1724 }
1725 let dict_data = file_handle
1726 .read_bytes_range(dict_header_end..dict_end)
1727 .await?;
1728 Some(CompressionDict::from_owned_bytes(dict_data))
1729 } else {
1730 None
1731 };
1732
1733 let data_slice = file_handle.slice(0..data_end_offset);
1735
1736 Ok(Self {
1737 data_slice,
1738 block_index,
1739 num_entries,
1740 cache: RwLock::new(BlockCache::new(cache_blocks, cache_budget_bytes)),
1741 bloom_filter,
1742 dictionary,
1743 _phantom: std::marker::PhantomData,
1744 })
1745 }
1746
1747 pub fn num_entries(&self) -> u64 {
1749 self.num_entries
1750 }
1751
1752 pub fn stats(&self) -> SSTableStats {
1754 SSTableStats {
1755 num_blocks: self.block_index.len(),
1756 num_sparse_entries: 0, num_entries: self.num_entries,
1758 has_bloom_filter: self.bloom_filter.is_some(),
1759 has_dictionary: self.dictionary.is_some(),
1760 bloom_filter_size: self
1761 .bloom_filter
1762 .as_ref()
1763 .map(|b| b.size_bytes())
1764 .unwrap_or(0),
1765 dictionary_size: self.dictionary.as_ref().map(|d| d.len()).unwrap_or(0),
1766 cache_insert_bypasses: self.cache.read().insert_bypasses,
1767 }
1768 }
1769
1770 pub fn cached_blocks(&self) -> usize {
1772 self.cache.read().blocks.len()
1773 }
1774
1775 pub fn cached_bytes(&self) -> usize {
1781 self.cache.read().retained_bytes
1782 }
1783
1784 pub async fn get(&self, key: &[u8]) -> io::Result<Option<V>> {
1789 log::debug!(
1790 "SSTable::get called, key_len={}, total_blocks={}",
1791 key.len(),
1792 self.block_index.len()
1793 );
1794
1795 if let Some(ref bloom) = self.bloom_filter
1797 && !bloom.may_contain(key)
1798 {
1799 log::debug!("SSTable::get bloom filter negative");
1800 return Ok(None);
1801 }
1802
1803 let block_idx = match self.block_index.locate(key) {
1805 Some(idx) => idx,
1806 None => {
1807 log::debug!("SSTable::get key not found (before first block)");
1808 return Ok(None);
1809 }
1810 };
1811
1812 log::debug!("SSTable::get loading block_idx={}", block_idx);
1813
1814 let block_data = self.load_block(block_idx).await?;
1816 self.search_block(&block_data, key)
1817 }
1818
1819 pub async fn get_batch(&self, keys: &[&[u8]]) -> io::Result<Vec<Option<V>>> {
1825 if keys.is_empty() {
1826 return Ok(Vec::new());
1827 }
1828
1829 let mut key_to_block: Vec<(usize, usize)> = Vec::with_capacity(keys.len());
1831 for (key_idx, key) in keys.iter().enumerate() {
1832 if let Some(ref bloom) = self.bloom_filter
1834 && !bloom.may_contain(key)
1835 {
1836 key_to_block.push((key_idx, usize::MAX)); continue;
1838 }
1839
1840 match self.block_index.locate(key) {
1841 Some(block_idx) => key_to_block.push((key_idx, block_idx)),
1842 None => key_to_block.push((key_idx, usize::MAX)), }
1844 }
1845
1846 let mut blocks_to_load: Vec<usize> = key_to_block
1848 .iter()
1849 .filter(|(_, b)| *b != usize::MAX)
1850 .map(|(_, b)| *b)
1851 .collect();
1852 blocks_to_load.sort_unstable();
1853 blocks_to_load.dedup();
1854
1855 for &block_idx in &blocks_to_load {
1857 let _ = self.load_block(block_idx).await?;
1858 }
1859
1860 let mut results = vec![None; keys.len()];
1862 for (key_idx, block_idx) in key_to_block {
1863 if block_idx == usize::MAX {
1864 continue;
1865 }
1866 let block_data = self.load_block(block_idx).await?; results[key_idx] = self.search_block(&block_data, keys[key_idx])?;
1868 }
1869
1870 Ok(results)
1871 }
1872
1873 pub async fn preload_all_blocks(&self) -> io::Result<()> {
1878 for block_idx in 0..self.block_index.len() {
1879 self.load_block(block_idx).await?;
1880 }
1881 Ok(())
1882 }
1883
1884 pub async fn prefetch_leading_blocks(&self) -> io::Result<()> {
1893 let num_blocks = self.block_index.len();
1894 let max_blocks = {
1895 let cache = self.cache.read();
1896 if cache.blocks.len() >= cache.max_blocks
1897 || cache
1898 .max_bytes
1899 .is_some_and(|budget| cache.retained_bytes >= budget)
1900 {
1901 log::debug!("SSTable bulk prefetch skipped: cache retention capacity exhausted");
1902 return Ok(());
1903 }
1904 cache.max_blocks
1905 };
1906 let mut start = self.data_slice.len();
1907 let mut end = 0;
1908 let mut planned = 0;
1909 for i in 0..num_blocks.min(max_blocks) {
1910 let addr = self.block_index.get_addr(i).ok_or_else(|| {
1911 io::Error::new(io::ErrorKind::InvalidData, "SSTable prefetch block missing")
1912 })?;
1913 let block_end = addr
1914 .offset
1915 .checked_add(addr.length as u64)
1916 .filter(|end| *end <= self.data_slice.len())
1917 .ok_or_else(|| {
1918 io::Error::new(
1919 io::ErrorKind::InvalidData,
1920 "SSTable prefetch block out of bounds",
1921 )
1922 })?;
1923 let next_start = start.min(addr.offset);
1924 let next_end = end.max(block_end);
1925 if next_end - next_start > Self::PREFETCH_LEADING_MAX_BYTES {
1926 break;
1927 }
1928 start = next_start;
1929 end = next_end;
1930 planned += 1;
1931 }
1932 if planned == 0 {
1933 log::debug!("SSTable bulk prefetch skipped: no block fits the bounded input range");
1934 return Ok(());
1935 }
1936 let all_data = self.data_slice.read_bytes_range(start..end).await?;
1937 let mut inserted = 0;
1938 for i in 0..planned {
1939 let addr = self.block_index.get_addr(i).unwrap();
1940 if self.cache.read().blocks.contains_key(&addr.offset) {
1941 continue;
1942 }
1943 let begin = (addr.offset - start) as usize;
1944 let limit = begin + addr.length as usize;
1945 let compressed = all_data.get(begin..limit).ok_or_else(|| {
1946 io::Error::new(
1947 io::ErrorKind::UnexpectedEof,
1948 "SSTable prefetch range is truncated",
1949 )
1950 })?;
1951 let decompressed = if let Some(ref dict) = self.dictionary {
1952 crate::compression::decompress_with_dict_limited(
1953 compressed,
1954 dict,
1955 MAX_SSTABLE_BLOCK_BYTES,
1956 )?
1957 } else {
1958 crate::compression::decompress_limited(compressed, MAX_SSTABLE_BLOCK_BYTES)?
1959 };
1960 let mut cache = self.cache.write();
1961 if cache.blocks.contains_key(&addr.offset) {
1962 continue;
1963 }
1964 if cache.blocks.len() >= cache.max_blocks
1965 || cache.max_bytes.is_some_and(|budget| {
1966 decompressed.len() > budget.saturating_sub(cache.retained_bytes)
1967 })
1968 {
1969 break;
1970 }
1971 cache.insert(addr.offset, Arc::from(decompressed));
1972 inserted += 1;
1973 }
1974 log::debug!(
1975 "SSTable bulk prefetch planned {planned}/{num_blocks} blocks, retained {inserted} new blocks within cache caps"
1976 );
1977 Ok(())
1978 }
1979
1980 async fn load_block(&self, block_idx: usize) -> io::Result<Arc<[u8]>> {
1983 let addr = self.block_index.get_addr(block_idx).ok_or_else(|| {
1984 io::Error::new(io::ErrorKind::InvalidInput, "Block index out of range")
1985 })?;
1986
1987 {
1989 if let Some(block) = self.cache.read().peek(addr.offset) {
1990 return Ok(block);
1991 }
1992 }
1993
1994 log::debug!(
1995 "SSTable::load_block idx={} CACHE MISS, reading bytes [{}-{}]",
1996 block_idx,
1997 addr.offset,
1998 addr.offset + addr.length as u64
1999 );
2000
2001 let range = addr.byte_range();
2003 let compressed = self.data_slice.read_bytes_range(range).await?;
2004
2005 let decompressed = if let Some(ref dict) = self.dictionary {
2007 crate::compression::decompress_with_dict_limited(
2008 compressed.as_slice(),
2009 dict,
2010 MAX_SSTABLE_BLOCK_BYTES,
2011 )?
2012 } else {
2013 crate::compression::decompress_limited(compressed.as_slice(), MAX_SSTABLE_BLOCK_BYTES)?
2014 };
2015
2016 let block: Arc<[u8]> = Arc::from(decompressed);
2017
2018 {
2020 let mut cache = self.cache.write();
2021 cache.insert(addr.offset, Arc::clone(&block));
2022 }
2023
2024 Ok(block)
2025 }
2026
2027 #[cfg(feature = "sync")]
2029 fn load_block_sync(&self, block_idx: usize) -> io::Result<Arc<[u8]>> {
2030 let addr = self.block_index.get_addr(block_idx).ok_or_else(|| {
2031 io::Error::new(io::ErrorKind::InvalidInput, "Block index out of range")
2032 })?;
2033
2034 {
2036 if let Some(block) = self.cache.read().peek(addr.offset) {
2037 return Ok(block);
2038 }
2039 }
2040
2041 let range = addr.byte_range();
2043 let compressed = self.data_slice.read_bytes_range_sync(range)?;
2044
2045 let decompressed = if let Some(ref dict) = self.dictionary {
2047 crate::compression::decompress_with_dict_limited(
2048 compressed.as_slice(),
2049 dict,
2050 MAX_SSTABLE_BLOCK_BYTES,
2051 )?
2052 } else {
2053 crate::compression::decompress_limited(compressed.as_slice(), MAX_SSTABLE_BLOCK_BYTES)?
2054 };
2055
2056 let block: Arc<[u8]> = Arc::from(decompressed);
2057
2058 {
2060 let mut cache = self.cache.write();
2061 cache.insert(addr.offset, Arc::clone(&block));
2062 }
2063
2064 Ok(block)
2065 }
2066
2067 #[cfg(feature = "sync")]
2069 pub fn get_sync(&self, key: &[u8]) -> io::Result<Option<V>> {
2070 if let Some(ref bloom) = self.bloom_filter
2072 && !bloom.may_contain(key)
2073 {
2074 return Ok(None);
2075 }
2076
2077 let block_idx = match self.block_index.locate(key) {
2079 Some(idx) => idx,
2080 None => {
2081 return Ok(None);
2082 }
2083 };
2084
2085 let block_data = self.load_block_sync(block_idx)?;
2086 self.search_block(&block_data, key)
2087 }
2088
2089 fn block_entries<'b>(&self, block_data: &'b [u8]) -> io::Result<&'b [u8]> {
2091 Ok(BlockParts::split(block_data)?.entries)
2092 }
2093
2094 fn search_block(&self, block_data: &[u8], target_key: &[u8]) -> io::Result<Option<V>> {
2095 let parts = BlockParts::split(block_data)?;
2096
2097 let start = if parts.num_restarts() > 0 {
2100 let (mut lo, mut hi) = (0usize, parts.num_restarts());
2101 while lo < hi {
2102 let mid = lo + (hi - lo) / 2;
2103 if parts.restart_key(mid)? <= target_key {
2104 lo = mid + 1;
2105 } else {
2106 hi = mid;
2107 }
2108 }
2109 if lo == 0 {
2110 return Ok(None);
2112 }
2113 parts.restart_offset(lo - 1)?
2114 } else {
2115 0
2116 };
2117
2118 let mut reader = &parts.entries[start..];
2119 let mut current_key = Vec::new();
2120
2121 while !reader.is_empty() {
2122 let value = decode_block_entry(&mut reader, &mut current_key)?;
2123
2124 match current_key.as_slice().cmp(target_key) {
2125 std::cmp::Ordering::Equal => return Ok(Some(value)),
2126 std::cmp::Ordering::Greater => return Ok(None),
2127 std::cmp::Ordering::Less => continue,
2128 }
2129 }
2130
2131 Ok(None)
2132 }
2133
2134 pub async fn prefetch_range(&self, start_key: &[u8], end_key: &[u8]) -> io::Result<()> {
2136 let start_block = self.block_index.locate(start_key).unwrap_or(0);
2137 let end_block = self
2138 .block_index
2139 .locate(end_key)
2140 .unwrap_or(self.block_index.len().saturating_sub(1));
2141
2142 for block_idx in start_block..=end_block.min(self.block_index.len().saturating_sub(1)) {
2143 let _ = self.load_block(block_idx).await?;
2144 }
2145
2146 Ok(())
2147 }
2148
2149 pub fn iter(&self) -> AsyncSSTableIterator<'_, V> {
2151 AsyncSSTableIterator::new(self)
2152 }
2153
2154 pub async fn all_entries(&self) -> io::Result<Vec<(Vec<u8>, V)>> {
2156 let mut results = Vec::new();
2157
2158 for block_idx in 0..self.block_index.len() {
2159 let block_data = self.load_block(block_idx).await?;
2160 let mut reader = self.block_entries(&block_data)?;
2161 let mut current_key = Vec::new();
2162
2163 while !reader.is_empty() {
2164 let value = decode_block_entry(&mut reader, &mut current_key)?;
2165 results.push((current_key.clone(), value));
2166 }
2167 }
2168
2169 Ok(results)
2170 }
2171
2172 pub async fn prefix_scan(&self, prefix: &[u8]) -> io::Result<Vec<(Vec<u8>, V)>> {
2178 let (results, _) = self.prefix_scan_limited(prefix, usize::MAX).await?;
2179 Ok(results)
2180 }
2181
2182 pub async fn prefix_scan_limited(
2185 &self,
2186 prefix: &[u8],
2187 max_results: usize,
2188 ) -> io::Result<PrefixScanResult<V>> {
2189 self.prefix_scan_filtered(prefix, max_results, usize::MAX, |_| true)
2190 .await
2191 }
2192
2193 pub(crate) async fn prefix_scan_filtered(
2194 &self,
2195 prefix: &[u8],
2196 max_results: usize,
2197 max_scanned: usize,
2198 accepts: impl FnMut(&[u8]) -> bool + Send,
2199 ) -> io::Result<PrefixScanResult<V>> {
2200 self.prefix_scan_projected(prefix, max_results, max_scanned, accepts, |key, value| {
2201 (key.to_vec(), value)
2202 })
2203 .await
2204 }
2205
2206 pub(crate) async fn prefix_scan_values(
2207 &self,
2208 prefix: &[u8],
2209 max_results: usize,
2210 max_scanned: usize,
2211 accepts: impl FnMut(&[u8]) -> bool + Send,
2212 ) -> io::Result<(Vec<V>, bool)> {
2213 self.prefix_scan_projected(prefix, max_results, max_scanned, accepts, |_, value| value)
2214 .await
2215 }
2216
2217 async fn prefix_scan_projected<T: Send>(
2218 &self,
2219 prefix: &[u8],
2220 max_results: usize,
2221 max_scanned: usize,
2222 mut accepts: impl FnMut(&[u8]) -> bool + Send,
2223 mut project: impl FnMut(&[u8], V) -> T + Send,
2224 ) -> io::Result<(Vec<T>, bool)> {
2225 if self.block_index.is_empty() || prefix.is_empty() {
2226 return Ok((Vec::new(), false));
2227 }
2228
2229 let start_block = self.block_index.locate(prefix).unwrap_or(0);
2233
2234 let mut results = Vec::new();
2235 let mut scanned = 0usize;
2236
2237 for block_idx in start_block..self.block_index.len() {
2238 let block_data = self.load_block(block_idx).await?;
2239 let mut reader = self.block_entries(&block_data)?;
2240 let mut current_key = Vec::new();
2241
2242 while !reader.is_empty() {
2243 let value = decode_block_entry(&mut reader, &mut current_key)?;
2244
2245 if current_key.starts_with(prefix) {
2246 if scanned == max_scanned {
2247 return Err(io::Error::other(format!(
2248 "term dictionary scan exceeds {max_scanned} terms"
2249 )));
2250 }
2251 scanned += 1;
2252 if !accepts(¤t_key) {
2253 continue;
2254 }
2255 if results.len() >= max_results {
2256 return Ok((results, true));
2257 }
2258 results.push(project(¤t_key, value));
2259 } else if current_key.as_slice() > prefix {
2260 return Ok((results, false));
2262 }
2263 }
2264 }
2265
2266 Ok((results, false))
2267 }
2268
2269 #[cfg(feature = "sync")]
2271 pub fn prefix_scan_sync(&self, prefix: &[u8]) -> io::Result<Vec<(Vec<u8>, V)>> {
2272 let (results, _) = self.prefix_scan_limited_sync(prefix, usize::MAX)?;
2273 Ok(results)
2274 }
2275
2276 #[cfg(feature = "sync")]
2278 pub fn prefix_scan_limited_sync(
2279 &self,
2280 prefix: &[u8],
2281 max_results: usize,
2282 ) -> io::Result<PrefixScanResult<V>> {
2283 self.prefix_scan_filtered_sync(prefix, max_results, usize::MAX, |_| true)
2284 }
2285
2286 #[cfg(feature = "sync")]
2287 pub(crate) fn prefix_scan_filtered_sync(
2288 &self,
2289 prefix: &[u8],
2290 max_results: usize,
2291 max_scanned: usize,
2292 accepts: impl FnMut(&[u8]) -> bool,
2293 ) -> io::Result<PrefixScanResult<V>> {
2294 self.prefix_scan_projected_sync(prefix, max_results, max_scanned, accepts, |key, value| {
2295 (key.to_vec(), value)
2296 })
2297 }
2298
2299 #[cfg(feature = "sync")]
2300 pub(crate) fn prefix_scan_values_sync(
2301 &self,
2302 prefix: &[u8],
2303 max_results: usize,
2304 max_scanned: usize,
2305 accepts: impl FnMut(&[u8]) -> bool,
2306 ) -> io::Result<(Vec<V>, bool)> {
2307 self.prefix_scan_projected_sync(prefix, max_results, max_scanned, accepts, |_, value| value)
2308 }
2309
2310 #[cfg(feature = "sync")]
2311 fn prefix_scan_projected_sync<T>(
2312 &self,
2313 prefix: &[u8],
2314 max_results: usize,
2315 max_scanned: usize,
2316 mut accepts: impl FnMut(&[u8]) -> bool,
2317 mut project: impl FnMut(&[u8], V) -> T,
2318 ) -> io::Result<(Vec<T>, bool)> {
2319 if self.block_index.is_empty() || prefix.is_empty() {
2320 return Ok((Vec::new(), false));
2321 }
2322
2323 let start_block = self.block_index.locate(prefix).unwrap_or(0);
2325
2326 let mut results = Vec::new();
2327 let mut scanned = 0usize;
2328
2329 for block_idx in start_block..self.block_index.len() {
2330 let block_data = self.load_block_sync(block_idx)?;
2331 let mut reader = self.block_entries(&block_data)?;
2332 let mut current_key = Vec::new();
2333
2334 while !reader.is_empty() {
2335 let value = decode_block_entry(&mut reader, &mut current_key)?;
2336
2337 if current_key.starts_with(prefix) {
2338 if scanned == max_scanned {
2339 return Err(io::Error::other(format!(
2340 "term dictionary scan exceeds {max_scanned} terms"
2341 )));
2342 }
2343 scanned += 1;
2344 if !accepts(¤t_key) {
2345 continue;
2346 }
2347 if results.len() >= max_results {
2348 return Ok((results, true));
2349 }
2350 results.push(project(¤t_key, value));
2351 } else if current_key.as_slice() > prefix {
2352 return Ok((results, false));
2353 }
2354 }
2355 }
2356
2357 Ok((results, false))
2358 }
2359}
2360
2361pub struct AsyncSSTableIterator<'a, V: SSTableValue> {
2363 reader: &'a AsyncSSTableReader<V>,
2364 current_block: usize,
2365 block_data: Option<Arc<[u8]>>,
2366 block_offset: usize,
2367 block_entries_end: usize,
2369 current_key: Vec<u8>,
2370 finished: bool,
2371}
2372
2373impl<'a, V: SSTableValue> AsyncSSTableIterator<'a, V> {
2374 fn new(reader: &'a AsyncSSTableReader<V>) -> Self {
2375 Self {
2376 reader,
2377 current_block: 0,
2378 block_data: None,
2379 block_offset: 0,
2380 block_entries_end: 0,
2381 current_key: Vec::new(),
2382 finished: reader.block_index.is_empty(),
2383 }
2384 }
2385
2386 async fn load_next_block(&mut self) -> io::Result<bool> {
2387 if self.current_block >= self.reader.block_index.len() {
2388 self.finished = true;
2389 return Ok(false);
2390 }
2391
2392 let block = self.reader.load_block(self.current_block).await?;
2393 self.block_entries_end = self.reader.block_entries(&block)?.len();
2394 self.block_data = Some(block);
2395 self.block_offset = 0;
2396 self.current_key.clear();
2397 self.current_block += 1;
2398 Ok(true)
2399 }
2400
2401 pub async fn next(&mut self) -> io::Result<Option<(Vec<u8>, V)>> {
2403 if self.finished {
2404 return Ok(None);
2405 }
2406
2407 if self.block_data.is_none() && !self.load_next_block().await? {
2408 return Ok(None);
2409 }
2410
2411 loop {
2412 let block = self.block_data.as_ref().unwrap();
2413 if self.block_offset >= self.block_entries_end {
2414 if !self.load_next_block().await? {
2415 return Ok(None);
2416 }
2417 continue;
2418 }
2419
2420 let mut reader = &block[self.block_offset..self.block_entries_end];
2421 let start_len = reader.len();
2422
2423 let value = decode_block_entry(&mut reader, &mut self.current_key)?;
2424
2425 self.block_offset += start_len - reader.len();
2426
2427 return Ok(Some((self.current_key.clone(), value)));
2428 }
2429 }
2430}
2431
2432#[cfg(test)]
2433mod tests {
2434 use super::*;
2435
2436 #[tokio::test]
2437 async fn value_projection_preserves_filtered_scan_budgets_order_and_errors() {
2438 let (bytes, _) = keyed_table(4096);
2439 let reader =
2440 AsyncSSTableReader::<u64>::open(FileHandle::from_bytes(OwnedBytes::new(bytes)), 8)
2441 .await
2442 .unwrap();
2443 for prefix in [b"field".as_slice(), b"field02", b"missing", b""] {
2444 for results in [0, 1, 17, 4096] {
2445 for scanned in [0, 1, 12, 8192] {
2446 let accepts = |key: &[u8]| key.last().is_some_and(|byte| byte % 2 == 0);
2447 let expected = reader
2448 .prefix_scan_filtered(prefix, results, scanned, accepts)
2449 .await
2450 .map(|(rows, more)| {
2451 (
2452 rows.into_iter().map(|(_, value)| value).collect::<Vec<_>>(),
2453 more,
2454 )
2455 })
2456 .map_err(|error| (error.kind(), error.to_string()));
2457 let actual = reader
2458 .prefix_scan_values(prefix, results, scanned, accepts)
2459 .await
2460 .map_err(|error| (error.kind(), error.to_string()));
2461 assert_eq!(actual, expected);
2462 #[cfg(feature = "sync")]
2463 assert_eq!(
2464 reader
2465 .prefix_scan_values_sync(prefix, results, scanned, accepts)
2466 .map_err(|error| (error.kind(), error.to_string())),
2467 expected
2468 );
2469 }
2470 }
2471 }
2472 }
2473
2474 #[test]
2475 fn fixed_inline_decode_preserves_values_and_serialized_bytes() {
2476 for (docs, frequencies) in [
2477 (vec![0], vec![1]),
2478 (vec![1, 128], vec![255, 128]),
2479 (vec![0, 1, u32::MAX], vec![1, 129, 10]),
2480 ] {
2481 let info = TermInfo::try_inline(&docs, &frequencies).unwrap();
2482 let mut before = Vec::new();
2483 info.serialize(&mut before).unwrap();
2484 let fixed = info.decode_inline_fixed().unwrap();
2485 assert_eq!(fixed.docs(), docs);
2486 assert_eq!(fixed.frequencies(), frequencies);
2487 assert_eq!(info.decode_inline().unwrap(), (docs, frequencies));
2488 let restored = TermInfo::try_inline(fixed.docs(), fixed.frequencies()).unwrap();
2489 let mut after = Vec::new();
2490 restored.serialize(&mut after).unwrap();
2491 assert_eq!(before, after);
2492 }
2493 }
2494
2495 #[test]
2496 fn inline_values_above_u32_are_rejected_without_truncation() {
2497 for pair in [(u64::from(u32::MAX) + 1, 1), (1, u64::from(u32::MAX) + 1)] {
2498 let mut encoded = Vec::new();
2499 write_vint(&mut encoded, pair.0).unwrap();
2500 write_vint(&mut encoded, pair.1).unwrap();
2501 let mut data = [0; 16];
2502 data[..encoded.len()].copy_from_slice(&encoded);
2503 let info = TermInfo::Inline {
2504 doc_freq: 1,
2505 data,
2506 data_len: encoded.len() as u8,
2507 };
2508 assert!(info.decode_inline_fixed().is_none());
2509 assert!(info.decode_inline().is_none());
2510 }
2511 }
2512
2513 #[test]
2514 fn test_bloom_filter_basic() {
2515 let mut bloom = BloomFilter::new(100, 10);
2516
2517 bloom.insert(b"hello");
2518 bloom.insert(b"world");
2519 bloom.insert(b"test");
2520
2521 assert!(bloom.may_contain(b"hello"));
2522 assert!(bloom.may_contain(b"world"));
2523 assert!(bloom.may_contain(b"test"));
2524
2525 assert!(!bloom.may_contain(b"notfound"));
2527 assert!(!bloom.may_contain(b"missing"));
2528 }
2529
2530 #[test]
2531 fn test_bloom_filter_serialization() {
2532 let mut bloom = BloomFilter::new(100, 10);
2533 bloom.insert(b"key1");
2534 bloom.insert(b"key2");
2535
2536 let bytes = bloom.to_bytes();
2537 let restored = BloomFilter::from_owned_bytes(OwnedBytes::new(bytes)).unwrap();
2538
2539 assert!(restored.may_contain(b"key1"));
2540 assert!(restored.may_contain(b"key2"));
2541 assert!(!restored.may_contain(b"key3"));
2542 }
2543
2544 #[test]
2545 fn bloom_header_preserves_bit_counts_above_u32() {
2546 let num_bits = u32::MAX as usize + 1;
2547 let mut header = Vec::new();
2548 write_bloom_header(&mut header, num_bits, BLOOM_HASH_COUNT, 1).unwrap();
2549 assert_eq!(header.len(), BLOOM_FILTER_HEADER_SIZE);
2550 assert_eq!(
2551 u64::from_le_bytes(header[0..8].try_into().unwrap()),
2552 num_bits as u64
2553 );
2554 }
2555
2556 #[test]
2557 fn test_bloom_filter_false_positive_rate() {
2558 let num_keys = 10000;
2559 let mut bloom = BloomFilter::new(num_keys, BLOOM_BITS_PER_KEY);
2560
2561 for i in 0..num_keys {
2563 let key = format!("key_{}", i);
2564 bloom.insert(key.as_bytes());
2565 }
2566
2567 for i in 0..num_keys {
2569 let key = format!("key_{}", i);
2570 assert!(bloom.may_contain(key.as_bytes()));
2571 }
2572
2573 let mut false_positives = 0;
2575 let test_count = 10000;
2576 for i in 0..test_count {
2577 let key = format!("nonexistent_{}", i);
2578 if bloom.may_contain(key.as_bytes()) {
2579 false_positives += 1;
2580 }
2581 }
2582
2583 let fp_rate = false_positives as f64 / test_count as f64;
2586 assert!(
2587 fp_rate < 0.03,
2588 "False positive rate {} is too high",
2589 fp_rate
2590 );
2591 }
2592
2593 #[test]
2594 fn test_sstable_writer_config() {
2595 use crate::structures::IndexOptimization;
2596
2597 let config = SSTableWriterConfig::default();
2599 assert_eq!(config.compression_level.0, 9); assert!(config.use_bloom_filter); assert!(!config.use_dictionary);
2602
2603 let adaptive = SSTableWriterConfig::from_optimization(IndexOptimization::Adaptive);
2605 assert_eq!(adaptive.compression_level.0, 9);
2606 assert!(adaptive.use_bloom_filter);
2607 assert!(!adaptive.use_dictionary);
2608
2609 let size = SSTableWriterConfig::from_optimization(IndexOptimization::SizeOptimized);
2611 assert_eq!(size.compression_level.0, 22); assert!(size.use_bloom_filter);
2613 assert!(size.use_dictionary);
2614
2615 let perf = SSTableWriterConfig::from_optimization(IndexOptimization::PerformanceOptimized);
2617 assert_eq!(perf.compression_level.0, 1); assert!(perf.use_bloom_filter); assert!(!perf.use_dictionary);
2620
2621 let fast = SSTableWriterConfig::fast();
2623 assert_eq!(fast.compression_level.0, 1);
2624
2625 let max = SSTableWriterConfig::max_compression();
2626 assert_eq!(max.compression_level.0, 22);
2627 }
2628
2629 #[test]
2630 fn test_vint_roundtrip() {
2631 let test_values = [0u64, 1, 127, 128, 255, 256, 16383, 16384, u64::MAX];
2632
2633 for &val in &test_values {
2634 let mut buf = Vec::new();
2635 write_vint(&mut buf, val).unwrap();
2636 let mut reader = buf.as_slice();
2637 let decoded = read_vint(&mut reader).unwrap();
2638 assert_eq!(val, decoded, "Failed for value {}", val);
2639 }
2640 }
2641
2642 fn keyed_table(num_keys: usize) -> (Vec<u8>, Vec<Vec<u8>>) {
2643 let mut keys: Vec<Vec<u8>> = (0..num_keys)
2644 .map(|i| format!("field{:02}/term{:07}", i % 7, i * 7919 % 100_003).into_bytes())
2645 .collect();
2646 keys.sort();
2647 keys.dedup();
2648 let mut writer = SSTableWriter::<_, u64>::new(Vec::new());
2649 for (i, key) in keys.iter().enumerate() {
2650 writer.insert(key, &(i as u64)).unwrap();
2651 }
2652 (writer.finish().unwrap(), keys)
2653 }
2654
2655 async fn check_every_key_and_scan(bytes: Vec<u8>, keys: &[Vec<u8>]) {
2656 let handle = FileHandle::from_bytes(OwnedBytes::new(bytes));
2657 let reader = AsyncSSTableReader::<u64>::open(handle, 8).await.unwrap();
2658 assert!(reader.block_index.len() > 3, "test needs several blocks");
2659
2660 for (i, key) in keys.iter().enumerate() {
2662 assert_eq!(reader.get(key).await.unwrap(), Some(i as u64), "key {i}");
2663 let mut before = key.clone();
2664 *before.last_mut().unwrap() -= 1;
2665 let mut after = key.clone();
2666 after.push(0);
2667 assert!(
2668 reader.get(&before).await.unwrap().is_none() || keys.binary_search(&before).is_ok()
2669 );
2670 assert!(
2671 reader.get(&after).await.unwrap().is_none() || keys.binary_search(&after).is_ok()
2672 );
2673 }
2674 assert!(reader.get(b"").await.unwrap().is_none());
2675 assert!(reader.get(b"zzz").await.unwrap().is_none());
2676
2677 let mut it = reader.iter();
2679 let mut seen = Vec::new();
2680 while let Some((k, v)) = it.next().await.unwrap() {
2681 assert_eq!(v as usize, seen.len());
2682 seen.push(k);
2683 }
2684 assert_eq!(seen, keys);
2685 let scanned = reader.prefix_scan(b"field03/").await.unwrap();
2686 let expected: Vec<&Vec<u8>> = keys.iter().filter(|k| k.starts_with(b"field03/")).collect();
2687 assert_eq!(scanned.len(), expected.len());
2688 assert!(scanned.iter().zip(expected).all(|((k, _), e)| k == e));
2689 assert_eq!(reader.all_entries().await.unwrap().len(), keys.len());
2690 let batch: Vec<&[u8]> = keys.iter().step_by(97).map(|k| k.as_slice()).collect();
2691 let got = reader.get_batch(&batch).await.unwrap();
2692 assert!(got.iter().all(|v| v.is_some()));
2693 }
2694
2695 #[cfg(feature = "native")]
2698 #[tokio::test]
2699 async fn v5_restart_points_find_every_key_across_blocks() {
2700 let (bytes, keys) = keyed_table(20_000);
2701 check_every_key_and_scan(bytes, &keys).await;
2702 }
2703
2704 #[cfg(feature = "native")]
2706 #[tokio::test]
2707 async fn unknown_sstable_magic_is_refused() {
2708 let (mut bytes, _) = keyed_table(100);
2709 let n = bytes.len();
2710 bytes[n - 4..].copy_from_slice(&0x5354_4236u32.to_le_bytes()); let handle = FileHandle::from_bytes(OwnedBytes::new(bytes));
2712 let err = match AsyncSSTableReader::<u64>::open(handle, 8).await {
2713 Ok(_) => panic!("unknown magic must be refused"),
2714 Err(err) => err,
2715 };
2716 assert!(err.to_string().contains("incompatible Summa"), "{err}");
2717 }
2718
2719 #[cfg(feature = "native")]
2722 #[tokio::test]
2723 async fn prefix_scan_matches_prefix_of_smallest_key() {
2724 let mut writer = SSTableWriter::<_, u64>::new(Vec::new());
2725 writer.insert(b"apple", &1).unwrap();
2726 writer.insert(b"apricot", &2).unwrap();
2727 writer.insert(b"banana", &3).unwrap();
2728 let bytes = writer.finish().unwrap();
2729 let handle = FileHandle::from_bytes(OwnedBytes::new(bytes));
2730 let reader = AsyncSSTableReader::<u64>::open(handle, 4).await.unwrap();
2731
2732 let keys = |entries: Vec<(Vec<u8>, u64)>| -> Vec<Vec<u8>> {
2733 entries.into_iter().map(|(k, _)| k).collect()
2734 };
2735
2736 assert_eq!(
2739 keys(reader.prefix_scan(b"ap").await.unwrap()),
2740 vec![b"apple".to_vec(), b"apricot".to_vec()]
2741 );
2742 assert_eq!(
2743 keys(reader.prefix_scan(b"a").await.unwrap()),
2744 vec![b"apple".to_vec(), b"apricot".to_vec()]
2745 );
2746 assert_eq!(
2747 keys(reader.prefix_scan(b"ba").await.unwrap()),
2748 vec![b"banana".to_vec()]
2749 );
2750 assert!(reader.prefix_scan(b"0").await.unwrap().is_empty());
2752
2753 #[cfg(feature = "sync")]
2754 {
2755 assert_eq!(
2756 keys(reader.prefix_scan_sync(b"ap").unwrap()),
2757 vec![b"apple".to_vec(), b"apricot".to_vec()]
2758 );
2759 assert!(reader.prefix_scan_sync(b"0").unwrap().is_empty());
2760 }
2761 }
2762
2763 #[cfg(feature = "native")]
2764 #[tokio::test]
2765 async fn filtered_prefix_scans_bound_examined_terms_and_matched_results_independently() {
2766 let mut writer = SSTableWriter::<_, u64>::new(Vec::new());
2767 for (key, value) in [(b"aa", 1), (b"ab", 2), (b"ac", 3), (b"ba", 4)] {
2768 writer.insert(key, &value).unwrap();
2769 }
2770 let reader = AsyncSSTableReader::<u64>::open(
2771 FileHandle::from_bytes(OwnedBytes::new(writer.finish().unwrap())),
2772 4,
2773 )
2774 .await
2775 .unwrap();
2776 let accepts = |key: &[u8]| key.ends_with(b"c");
2777 assert_eq!(
2778 reader
2779 .prefix_scan_filtered(b"a", 1, 3, accepts)
2780 .await
2781 .unwrap(),
2782 (vec![(b"ac".to_vec(), 3)], false)
2783 );
2784 assert!(
2785 reader
2786 .prefix_scan_filtered(b"a", 1, 2, accepts)
2787 .await
2788 .unwrap_err()
2789 .to_string()
2790 .contains("scan exceeds 2")
2791 );
2792 assert_eq!(
2793 reader
2794 .prefix_scan_filtered(b"a", 0, 3, accepts)
2795 .await
2796 .unwrap(),
2797 (vec![], true)
2798 );
2799 assert_eq!(
2800 reader
2801 .prefix_scan_filtered(b"a", 1, 3, |_| false)
2802 .await
2803 .unwrap(),
2804 (vec![], false)
2805 );
2806 #[cfg(feature = "sync")]
2807 {
2808 assert_eq!(
2809 reader
2810 .prefix_scan_filtered_sync(b"a", 1, 3, accepts)
2811 .unwrap(),
2812 (vec![(b"ac".to_vec(), 3)], false)
2813 );
2814 assert!(
2815 reader
2816 .prefix_scan_filtered_sync(b"a", 1, 2, accepts)
2817 .unwrap_err()
2818 .to_string()
2819 .contains("scan exceeds 2")
2820 );
2821 assert_eq!(
2822 reader
2823 .prefix_scan_filtered_sync(b"a", 0, 3, accepts)
2824 .unwrap(),
2825 (vec![], true)
2826 );
2827 }
2828 }
2829
2830 #[test]
2831 fn test_common_prefix_len() {
2832 assert_eq!(common_prefix_len(b"hello", b"hello"), 5);
2833 assert_eq!(common_prefix_len(b"hello", b"help"), 3);
2834 assert_eq!(common_prefix_len(b"hello", b"world"), 0);
2835 assert_eq!(common_prefix_len(b"", b"hello"), 0);
2836 assert_eq!(common_prefix_len(b"hello", b""), 0);
2837 }
2838
2839 #[test]
2840 fn decode_block_entry_reconstructs_keys_and_consumes_one_entry() {
2841 let mut encoded = Vec::new();
2842
2843 write_vint(&mut encoded, 0).unwrap();
2844 write_vint(&mut encoded, 5).unwrap();
2845 encoded.extend_from_slice(b"alpha");
2846 7_u64.serialize(&mut encoded).unwrap();
2847
2848 write_vint(&mut encoded, 3).unwrap();
2849 write_vint(&mut encoded, 3).unwrap();
2850 encoded.extend_from_slice(b"ine");
2851 11_u64.serialize(&mut encoded).unwrap();
2852
2853 let mut reader = encoded.as_slice();
2854 let mut key = Vec::new();
2855
2856 assert_eq!(decode_block_entry::<u64>(&mut reader, &mut key).unwrap(), 7);
2857 assert_eq!(key, b"alpha");
2858 assert!(!reader.is_empty(), "the second entry must remain unread");
2859
2860 assert_eq!(
2861 decode_block_entry::<u64>(&mut reader, &mut key).unwrap(),
2862 11
2863 );
2864 assert_eq!(key, b"alpine");
2865 assert!(reader.is_empty());
2866 }
2867
2868 #[test]
2869 fn decode_block_entry_rejects_truncated_suffix() {
2870 let encoded = [0, 4, b'o', b'n'];
2871 let mut reader = encoded.as_slice();
2872 let mut key = Vec::new();
2873
2874 let error = decode_block_entry::<u64>(&mut reader, &mut key).unwrap_err();
2875
2876 assert_eq!(error.kind(), io::ErrorKind::UnexpectedEof);
2877 assert_eq!(error.to_string(), "SSTable block suffix truncated");
2878 }
2879}