1use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
22use parking_lot::RwLock;
23use rustc_hash::FxHashMap;
24use std::io::{self, Read, Write};
25use std::sync::Arc;
26
27#[cfg(feature = "fst-index")]
28use super::sstable_index::FstBlockIndex;
29use super::sstable_index::{BlockAddr, BlockIndex, MmapBlockIndex};
30use super::vint::{read_vint, write_vint};
31use crate::compression::{CompressionDict, CompressionLevel};
32use crate::directories::{FileHandle, OwnedBytes};
33
34pub const SSTABLE_MAGIC: u32 = 0x53544234; pub const BLOCK_SIZE: usize = 16 * 1024;
40
41pub const DEFAULT_DICT_SIZE: usize = 64 * 1024;
43
44pub const BLOOM_BITS_PER_KEY: usize = 10;
46
47pub const BLOOM_HASH_COUNT: usize = 7;
49const BLOOM_FILTER_HEADER_SIZE: usize = 16;
50
51const MAX_SSTABLE_BLOCK_BYTES: usize = 64 * 1024 * 1024;
52const MAX_SSTABLE_DICTIONARY_BYTES: u64 = 16 * 1024 * 1024;
53
54pub type PrefixScanResult<V> = (Vec<(Vec<u8>, V)>, bool);
56
57#[derive(Debug, Clone)]
63pub struct BloomFilter {
64 bits: BloomBits,
65 num_bits: usize,
66 num_hashes: usize,
67}
68
69#[derive(Debug, Clone)]
71enum BloomBits {
72 Vec(Vec<u64>),
74 Bytes(OwnedBytes),
76}
77
78impl BloomBits {
79 #[inline]
80 fn len(&self) -> usize {
81 match self {
82 BloomBits::Vec(v) => v.len(),
83 BloomBits::Bytes(b) => b.len() / 8,
84 }
85 }
86
87 #[inline]
88 fn get(&self, word_idx: usize) -> u64 {
89 match self {
90 BloomBits::Vec(v) => v[word_idx],
91 BloomBits::Bytes(b) => {
92 let off = word_idx * 8;
93 u64::from_le_bytes([
94 b[off],
95 b[off + 1],
96 b[off + 2],
97 b[off + 3],
98 b[off + 4],
99 b[off + 5],
100 b[off + 6],
101 b[off + 7],
102 ])
103 }
104 }
105 }
106
107 #[inline]
108 fn set_bit(&mut self, word_idx: usize, bit_idx: usize) {
109 match self {
110 BloomBits::Vec(v) => v[word_idx] |= 1u64 << bit_idx,
111 BloomBits::Bytes(_) => panic!("cannot mutate read-only bloom filter"),
112 }
113 }
114
115 fn size_bytes(&self) -> usize {
116 match self {
117 BloomBits::Vec(v) => v.len() * 8,
118 BloomBits::Bytes(b) => b.len(),
119 }
120 }
121
122 fn write_to(&self, writer: &mut (impl Write + ?Sized)) -> io::Result<()> {
123 match self {
124 BloomBits::Vec(words) => {
125 #[cfg(target_endian = "little")]
126 {
127 let bytes = unsafe {
130 std::slice::from_raw_parts(
131 words.as_ptr().cast::<u8>(),
132 words.len().saturating_mul(8),
133 )
134 };
135 writer.write_all(bytes)
136 }
137 #[cfg(target_endian = "big")]
138 {
139 for &word in words {
140 writer.write_u64::<LittleEndian>(word)?;
141 }
142 Ok(())
143 }
144 }
145 BloomBits::Bytes(bytes) => writer.write_all(bytes.as_slice()),
146 }
147 }
148}
149
150impl BloomFilter {
151 pub(crate) const SERIALIZED_HEADER_SIZE: usize = BLOOM_FILTER_HEADER_SIZE;
152
153 pub fn new(expected_keys: usize, bits_per_key: usize) -> Self {
155 let num_bits = expected_keys.saturating_mul(bits_per_key).max(64);
156 let num_words = num_bits.div_ceil(64);
157 Self {
158 bits: BloomBits::Vec(vec![0u64; num_words]),
159 num_bits,
160 num_hashes: BLOOM_HASH_COUNT,
161 }
162 }
163
164 pub fn from_bytes_mutable(data: &[u8]) -> io::Result<Self> {
168 if data.len() < BLOOM_FILTER_HEADER_SIZE {
169 return Err(io::Error::new(
170 io::ErrorKind::InvalidData,
171 "Bloom filter data too short",
172 ));
173 }
174 let num_bits = usize::try_from(u64::from_le_bytes(data[0..8].try_into().unwrap()))
175 .map_err(|_| {
176 io::Error::new(
177 io::ErrorKind::InvalidData,
178 "Bloom filter bit count exceeds addressable memory",
179 )
180 })?;
181 let num_hashes = u32::from_le_bytes(data[8..12].try_into().unwrap()) as usize;
182 let num_words = u32::from_le_bytes(data[12..16].try_into().unwrap()) as usize;
183
184 validate_bloom_header(data.len(), num_bits, num_hashes, num_words)?;
185 let expected_len = BLOOM_FILTER_HEADER_SIZE + num_words * 8;
186 if data.len() != expected_len {
187 return Err(io::Error::new(
188 io::ErrorKind::InvalidData,
189 "Bloom filter data truncated",
190 ));
191 }
192
193 let mut vec = vec![0u64; num_words];
194 for (i, v) in vec.iter_mut().enumerate() {
195 let off = BLOOM_FILTER_HEADER_SIZE + i * 8;
196 *v = u64::from_le_bytes(data[off..off + 8].try_into().unwrap());
197 }
198
199 Ok(Self {
200 bits: BloomBits::Vec(vec),
201 num_bits,
202 num_hashes,
203 })
204 }
205
206 pub fn from_owned_bytes(data: OwnedBytes) -> io::Result<Self> {
208 if data.len() < BLOOM_FILTER_HEADER_SIZE {
209 return Err(io::Error::new(
210 io::ErrorKind::InvalidData,
211 "Bloom filter data too short",
212 ));
213 }
214 let d = data.as_slice();
215 let num_bits =
216 usize::try_from(u64::from_le_bytes(d[0..8].try_into().unwrap())).map_err(|_| {
217 io::Error::new(
218 io::ErrorKind::InvalidData,
219 "Bloom filter bit count exceeds addressable memory",
220 )
221 })?;
222 let num_hashes = u32::from_le_bytes(d[8..12].try_into().unwrap()) as usize;
223 let num_words = u32::from_le_bytes(d[12..16].try_into().unwrap()) as usize;
224
225 validate_bloom_header(d.len(), num_bits, num_hashes, num_words)?;
226 let expected_len = BLOOM_FILTER_HEADER_SIZE + num_words * 8;
227 if d.len() != expected_len {
228 return Err(io::Error::new(
229 io::ErrorKind::InvalidData,
230 "Bloom filter data truncated",
231 ));
232 }
233
234 let bits_bytes =
236 data.slice(BLOOM_FILTER_HEADER_SIZE..BLOOM_FILTER_HEADER_SIZE + num_words * 8);
237
238 Ok(Self {
239 bits: BloomBits::Bytes(bits_bytes),
240 num_bits,
241 num_hashes,
242 })
243 }
244
245 pub fn serialized_len(&self) -> usize {
247 BLOOM_FILTER_HEADER_SIZE + self.bits.len() * 8
248 }
249
250 pub fn write_to(&self, writer: &mut (impl Write + ?Sized)) -> io::Result<()> {
252 let num_words = self.bits.len();
253 write_bloom_header(writer, self.num_bits, self.num_hashes, num_words)?;
254 self.bits.write_to(writer)
255 }
256
257 pub fn to_bytes(&self) -> Vec<u8> {
259 let mut data = Vec::with_capacity(self.serialized_len());
260 self.write_to(&mut data)
261 .expect("writing a bloom filter to Vec cannot fail");
262 data
263 }
264
265 pub fn insert(&mut self, key: &[u8]) {
267 let (h1, h2) = self.hash_pair(key);
268 for i in 0..self.num_hashes {
269 let bit_pos = self.get_bit_pos(h1, h2, i);
270 let word_idx = bit_pos / 64;
271 let bit_idx = bit_pos % 64;
272 if word_idx < self.bits.len() {
273 self.bits.set_bit(word_idx, bit_idx);
274 }
275 }
276 }
277
278 pub fn may_contain(&self, key: &[u8]) -> bool {
281 let (h1, h2) = self.hash_pair(key);
282 for i in 0..self.num_hashes {
283 let bit_pos = self.get_bit_pos(h1, h2, i);
284 let word_idx = bit_pos / 64;
285 let bit_idx = bit_pos % 64;
286 if word_idx >= self.bits.len() || (self.bits.get(word_idx) & (1u64 << bit_idx)) == 0 {
287 return false;
288 }
289 }
290 true
291 }
292
293 pub fn size_bytes(&self) -> usize {
295 BLOOM_FILTER_HEADER_SIZE + self.bits.size_bytes()
296 }
297
298 pub fn insert_hashed(&mut self, h1: u64, h2: u64) {
300 for i in 0..self.num_hashes {
301 let bit_pos = self.get_bit_pos(h1, h2, i);
302 let word_idx = bit_pos / 64;
303 let bit_idx = bit_pos % 64;
304 if word_idx < self.bits.len() {
305 self.bits.set_bit(word_idx, bit_idx);
306 }
307 }
308 }
309
310 #[inline]
312 fn hash_pair(&self, key: &[u8]) -> (u64, u64) {
313 let mut h1: u64 = 0xcbf29ce484222325;
314 let mut h2: u64 = 0x84222325cbf29ce4;
315 for &byte in key {
316 h1 ^= byte as u64;
317 h1 = h1.wrapping_mul(0x100000001b3);
318 h2 = h2.wrapping_mul(0x100000001b3);
319 h2 ^= byte as u64;
320 }
321 (h1, h2)
322 }
323
324 #[inline]
326 fn get_bit_pos(&self, h1: u64, h2: u64, i: usize) -> usize {
327 (h1.wrapping_add((i as u64).wrapping_mul(h2)) % (self.num_bits as u64)) as usize
328 }
329}
330
331fn write_bloom_header(
332 writer: &mut (impl Write + ?Sized),
333 num_bits: usize,
334 num_hashes: usize,
335 num_words: usize,
336) -> io::Result<()> {
337 writer.write_u64::<LittleEndian>(u64::try_from(num_bits).map_err(|_| {
338 io::Error::new(
339 io::ErrorKind::InvalidInput,
340 "Bloom filter bit count exceeds u64",
341 )
342 })?)?;
343 writer.write_u32::<LittleEndian>(u32::try_from(num_hashes).map_err(|_| {
344 io::Error::new(
345 io::ErrorKind::InvalidInput,
346 "Bloom filter hash count exceeds u32",
347 )
348 })?)?;
349 writer.write_u32::<LittleEndian>(u32::try_from(num_words).map_err(|_| {
350 io::Error::new(
351 io::ErrorKind::InvalidInput,
352 "Bloom filter word count exceeds u32",
353 )
354 })?)?;
355 Ok(())
356}
357
358fn validate_bloom_header(
359 data_len: usize,
360 num_bits: usize,
361 num_hashes: usize,
362 num_words: usize,
363) -> io::Result<()> {
364 if num_bits == 0 || num_hashes == 0 || num_hashes > 32 || num_words == 0 {
365 return Err(io::Error::new(
366 io::ErrorKind::InvalidData,
367 "invalid bloom filter parameters",
368 ));
369 }
370 let word_bytes = num_words
371 .checked_mul(8)
372 .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "bloom filter size overflow"))?;
373 let expected_len = BLOOM_FILTER_HEADER_SIZE
374 .checked_add(word_bytes)
375 .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "bloom filter size overflow"))?;
376 let capacity_bits = num_words
377 .checked_mul(64)
378 .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "bloom filter size overflow"))?;
379 if expected_len > data_len
380 || num_bits > capacity_bits
381 || num_bits <= capacity_bits.saturating_sub(64)
382 {
383 return Err(io::Error::new(
384 io::ErrorKind::InvalidData,
385 "inconsistent bloom filter dimensions",
386 ));
387 }
388 Ok(())
389}
390
391#[inline]
394fn bloom_hash_pair(key: &[u8]) -> (u64, u64) {
395 let mut h1: u64 = 0xcbf29ce484222325;
396 let mut h2: u64 = 0x84222325cbf29ce4;
397 for &byte in key {
398 h1 ^= byte as u64;
399 h1 = h1.wrapping_mul(0x100000001b3);
400 h2 = h2.wrapping_mul(0x100000001b3);
401 h2 ^= byte as u64;
402 }
403 (h1, h2)
404}
405
406pub trait SSTableValue: Clone + Send + Sync {
413 fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()>;
415
416 fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self>;
418}
419
420impl SSTableValue for u64 {
422 fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
423 write_vint(writer, *self)
424 }
425
426 fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
427 read_vint(reader)
428 }
429}
430
431impl SSTableValue for Vec<u8> {
433 fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
434 write_vint(writer, self.len() as u64)?;
435 writer.write_all(self)
436 }
437
438 fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
439 let len = usize::try_from(read_vint(reader)?).map_err(|_| {
440 io::Error::new(io::ErrorKind::InvalidData, "SSTable value length too large")
441 })?;
442 if len > MAX_SSTABLE_BLOCK_BYTES {
443 return Err(io::Error::new(
444 io::ErrorKind::InvalidData,
445 "SSTable value exceeds block safety limit",
446 ));
447 }
448 let mut data = vec![0u8; len];
449 reader.read_exact(&mut data)?;
450 Ok(data)
451 }
452}
453
454#[derive(Debug, Clone, Copy, PartialEq, Eq)]
457pub struct SparseDimInfo {
458 pub offset: u64,
460 pub length: u32,
462}
463
464impl SparseDimInfo {
465 pub fn new(offset: u64, length: u32) -> Self {
466 Self { offset, length }
467 }
468}
469
470impl SSTableValue for SparseDimInfo {
471 fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
472 write_vint(writer, self.offset)?;
473 write_vint(writer, self.length as u64)
474 }
475
476 fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
477 let offset = read_vint(reader)?;
478 let length = u32::try_from(read_vint(reader)?).map_err(|_| {
479 io::Error::new(
480 io::ErrorKind::InvalidData,
481 "sparse posting length exceeds u32",
482 )
483 })?;
484 Ok(Self { offset, length })
485 }
486}
487
488pub const MAX_INLINE_POSTINGS: usize = 3;
490
491#[derive(Debug, Clone, PartialEq, Eq)]
499pub enum TermInfo {
500 Inline {
503 doc_freq: u8,
505 data: [u8; 16],
508 data_len: u8,
510 },
511 External {
513 posting_offset: u64,
514 posting_len: u64,
515 doc_freq: u32,
516 position_offset: u64,
518 position_len: u64,
520 },
521}
522
523impl TermInfo {
524 pub fn external(posting_offset: u64, posting_len: u64, doc_freq: u32) -> Self {
526 TermInfo::External {
527 posting_offset,
528 posting_len,
529 doc_freq,
530 position_offset: 0,
531 position_len: 0,
532 }
533 }
534
535 pub fn external_with_positions(
537 posting_offset: u64,
538 posting_len: u64,
539 doc_freq: u32,
540 position_offset: u64,
541 position_len: u64,
542 ) -> Self {
543 TermInfo::External {
544 posting_offset,
545 posting_len,
546 doc_freq,
547 position_offset,
548 position_len,
549 }
550 }
551
552 pub fn try_inline(doc_ids: &[u32], term_freqs: &[u32]) -> Option<Self> {
555 if doc_ids.len() > MAX_INLINE_POSTINGS
556 || doc_ids.is_empty()
557 || doc_ids.len() != term_freqs.len()
558 {
559 return None;
560 }
561
562 let mut data = [0u8; 16];
563 let mut cursor = std::io::Cursor::new(&mut data[..]);
564 let mut prev_doc_id = 0u32;
565
566 for (i, &doc_id) in doc_ids.iter().enumerate() {
567 let delta = doc_id.checked_sub(prev_doc_id)?;
568 if write_vint(&mut cursor, delta as u64).is_err() {
569 return None;
570 }
571 if write_vint(&mut cursor, term_freqs[i] as u64).is_err() {
572 return None;
573 }
574 prev_doc_id = doc_id;
575 }
576
577 let data_len = cursor.position() as u8;
578 if data_len > 16 {
579 return None;
580 }
581
582 Some(TermInfo::Inline {
583 doc_freq: doc_ids.len() as u8,
584 data,
585 data_len,
586 })
587 }
588
589 pub fn try_inline_iter(count: usize, iter: impl Iterator<Item = (u32, u32)>) -> Option<Self> {
593 if count > MAX_INLINE_POSTINGS || count == 0 {
594 return None;
595 }
596
597 let mut data = [0u8; 16];
598 let mut cursor = std::io::Cursor::new(&mut data[..]);
599 let mut prev_doc_id = 0u32;
600
601 let mut actual_count = 0usize;
602 for (doc_id, tf) in iter {
603 if actual_count >= count {
604 return None;
605 }
606 let delta = doc_id.checked_sub(prev_doc_id)?;
607 if write_vint(&mut cursor, delta as u64).is_err() {
608 return None;
609 }
610 if write_vint(&mut cursor, tf as u64).is_err() {
611 return None;
612 }
613 prev_doc_id = doc_id;
614 actual_count += 1;
615 }
616
617 if actual_count != count {
618 return None;
619 }
620
621 let data_len = cursor.position() as u8;
622
623 Some(TermInfo::Inline {
624 doc_freq: count as u8,
625 data,
626 data_len,
627 })
628 }
629
630 pub fn doc_freq(&self) -> u32 {
632 match self {
633 TermInfo::Inline { doc_freq, .. } => *doc_freq as u32,
634 TermInfo::External { doc_freq, .. } => *doc_freq,
635 }
636 }
637
638 pub fn is_inline(&self) -> bool {
640 matches!(self, TermInfo::Inline { .. })
641 }
642
643 pub fn external_info(&self) -> Option<(u64, u64)> {
645 match self {
646 TermInfo::External {
647 posting_offset,
648 posting_len,
649 ..
650 } => Some((*posting_offset, *posting_len)),
651 TermInfo::Inline { .. } => None,
652 }
653 }
654
655 pub fn position_info(&self) -> Option<(u64, u64)> {
657 match self {
658 TermInfo::External {
659 position_offset,
660 position_len,
661 ..
662 } if *position_len > 0 => Some((*position_offset, *position_len)),
663 _ => None,
664 }
665 }
666
667 pub fn decode_inline(&self) -> Option<(Vec<u32>, Vec<u32>)> {
670 match self {
671 TermInfo::Inline {
672 doc_freq,
673 data,
674 data_len,
675 } => {
676 if *doc_freq == 0
677 || *doc_freq as usize > MAX_INLINE_POSTINGS
678 || *data_len as usize > data.len()
679 {
680 return None;
681 }
682 let mut doc_ids = Vec::with_capacity(*doc_freq as usize);
683 let mut term_freqs = Vec::with_capacity(*doc_freq as usize);
684 let mut reader = &data[..*data_len as usize];
685 let mut prev_doc_id = 0u32;
686
687 for _ in 0..*doc_freq {
688 let delta = read_vint(&mut reader).ok()? as u32;
689 let tf = read_vint(&mut reader).ok()? as u32;
690 let doc_id = prev_doc_id.checked_add(delta)?;
691 doc_ids.push(doc_id);
692 term_freqs.push(tf);
693 prev_doc_id = doc_id;
694 }
695
696 Some((doc_ids, term_freqs))
697 }
698 TermInfo::External { .. } => None,
699 }
700 }
701}
702
703impl SSTableValue for TermInfo {
704 fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
705 match self {
706 TermInfo::Inline {
707 doc_freq,
708 data,
709 data_len,
710 } => {
711 if *doc_freq == 0
712 || *doc_freq as usize > MAX_INLINE_POSTINGS
713 || *data_len as usize > data.len()
714 {
715 return Err(io::Error::new(
716 io::ErrorKind::InvalidInput,
717 "invalid inline TermInfo",
718 ));
719 }
720 writer.write_u8(0xFF)?;
722 writer.write_u8(*doc_freq)?;
723 writer.write_u8(*data_len)?;
724 writer.write_all(&data[..*data_len as usize])?;
725 }
726 TermInfo::External {
727 posting_offset,
728 posting_len,
729 doc_freq,
730 position_offset,
731 position_len,
732 } => {
733 if *position_len > 0 {
736 writer.write_u8(0x01)?;
737 write_vint(writer, *doc_freq as u64)?;
738 write_vint(writer, *posting_offset)?;
739 write_vint(writer, *posting_len)?;
740 write_vint(writer, *position_offset)?;
741 write_vint(writer, *position_len)?;
742 } else {
743 writer.write_u8(0x00)?;
744 write_vint(writer, *doc_freq as u64)?;
745 write_vint(writer, *posting_offset)?;
746 write_vint(writer, *posting_len)?;
747 }
748 }
749 }
750 Ok(())
751 }
752
753 fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
754 let tag = reader.read_u8()?;
755
756 if tag == 0xFF {
757 let doc_freq = reader.read_u8()?;
759 let data_len = reader.read_u8()?;
760 if doc_freq == 0 || doc_freq as usize > MAX_INLINE_POSTINGS || data_len as usize > 16 {
761 return Err(io::Error::new(
762 io::ErrorKind::InvalidData,
763 "invalid inline TermInfo lengths",
764 ));
765 }
766 let mut data = [0u8; 16];
767 reader.read_exact(&mut data[..data_len as usize])?;
768 Ok(TermInfo::Inline {
769 doc_freq,
770 data,
771 data_len,
772 })
773 } else if tag == 0x00 {
774 let doc_freq = read_vint(reader)? as u32;
776 let posting_offset = read_vint(reader)?;
777 let posting_len = read_vint(reader)?;
778 Ok(TermInfo::External {
779 posting_offset,
780 posting_len,
781 doc_freq,
782 position_offset: 0,
783 position_len: 0,
784 })
785 } else if tag == 0x01 {
786 let doc_freq = read_vint(reader)? as u32;
788 let posting_offset = read_vint(reader)?;
789 let posting_len = read_vint(reader)?;
790 let position_offset = read_vint(reader)?;
791 let position_len = read_vint(reader)?;
792 Ok(TermInfo::External {
793 posting_offset,
794 posting_len,
795 doc_freq,
796 position_offset,
797 position_len,
798 })
799 } else {
800 Err(io::Error::new(
801 io::ErrorKind::InvalidData,
802 format!("Invalid TermInfo tag: {}", tag),
803 ))
804 }
805 }
806}
807
808pub fn common_prefix_len(a: &[u8], b: &[u8]) -> usize {
810 a.iter().zip(b.iter()).take_while(|(x, y)| x == y).count()
811}
812
813fn decode_block_entry<V: SSTableValue>(
819 reader: &mut &[u8],
820 current_key: &mut Vec<u8>,
821) -> io::Result<V> {
822 let common_prefix_len = read_vint(reader)? as usize;
823 let suffix_len = read_vint(reader)? as usize;
824
825 if suffix_len > reader.len() {
826 return Err(io::Error::new(
827 io::ErrorKind::UnexpectedEof,
828 "SSTable block suffix truncated",
829 ));
830 }
831
832 current_key.truncate(common_prefix_len);
833 current_key.extend_from_slice(&reader[..suffix_len]);
834 *reader = &reader[suffix_len..];
835
836 V::deserialize(reader)
837}
838
839#[derive(Debug, Clone)]
841pub struct SSTableStats {
842 pub num_blocks: usize,
844 pub num_sparse_entries: usize,
846 pub num_entries: u64,
848 pub has_bloom_filter: bool,
850 pub has_dictionary: bool,
852 pub bloom_filter_size: usize,
854 pub dictionary_size: usize,
856}
857
858#[derive(Debug, Clone)]
860pub struct SSTableWriterConfig {
861 pub compression_level: CompressionLevel,
863 pub use_dictionary: bool,
865 pub dict_size: usize,
867 pub use_bloom_filter: bool,
869 pub bloom_bits_per_key: usize,
871}
872
873impl Default for SSTableWriterConfig {
874 fn default() -> Self {
875 Self::from_optimization(crate::structures::IndexOptimization::default())
876 }
877}
878
879impl SSTableWriterConfig {
880 pub fn from_optimization(optimization: crate::structures::IndexOptimization) -> Self {
882 use crate::structures::IndexOptimization;
883 match optimization {
884 IndexOptimization::Adaptive => Self {
885 compression_level: CompressionLevel::BETTER, use_dictionary: false,
887 dict_size: DEFAULT_DICT_SIZE,
888 use_bloom_filter: true, bloom_bits_per_key: BLOOM_BITS_PER_KEY,
890 },
891 IndexOptimization::SizeOptimized => Self {
892 compression_level: CompressionLevel::MAX, use_dictionary: true,
894 dict_size: DEFAULT_DICT_SIZE,
895 use_bloom_filter: true,
896 bloom_bits_per_key: BLOOM_BITS_PER_KEY,
897 },
898 IndexOptimization::PerformanceOptimized => Self {
899 compression_level: CompressionLevel::FAST, use_dictionary: false,
901 dict_size: DEFAULT_DICT_SIZE,
902 use_bloom_filter: true, bloom_bits_per_key: BLOOM_BITS_PER_KEY,
904 },
905 }
906 }
907
908 pub fn fast() -> Self {
910 Self::from_optimization(crate::structures::IndexOptimization::PerformanceOptimized)
911 }
912
913 pub fn max_compression() -> Self {
915 Self::from_optimization(crate::structures::IndexOptimization::SizeOptimized)
916 }
917}
918
919pub struct SSTableWriter<W: Write, V: SSTableValue> {
925 writer: W,
926 block_buffer: Vec<u8>,
927 prev_key: Vec<u8>,
928 index: Vec<BlockIndexEntry>,
929 current_offset: u64,
930 num_entries: u64,
931 block_first_key: Option<Vec<u8>>,
932 config: SSTableWriterConfig,
933 dictionary: Option<CompressionDict>,
935 bloom_hashes: Vec<(u64, u64)>,
938 _phantom: std::marker::PhantomData<V>,
939}
940
941impl<W: Write, V: SSTableValue> SSTableWriter<W, V> {
942 pub fn new(writer: W) -> Self {
944 Self::with_config(writer, SSTableWriterConfig::default())
945 }
946
947 pub fn with_config(writer: W, config: SSTableWriterConfig) -> Self {
949 Self {
950 writer,
951 block_buffer: Vec::with_capacity(BLOCK_SIZE),
952 prev_key: Vec::new(),
953 index: Vec::new(),
954 current_offset: 0,
955 num_entries: 0,
956 block_first_key: None,
957 config,
958 dictionary: None,
959 bloom_hashes: Vec::new(),
960 _phantom: std::marker::PhantomData,
961 }
962 }
963
964 pub fn with_dictionary(
966 writer: W,
967 config: SSTableWriterConfig,
968 dictionary: CompressionDict,
969 ) -> Self {
970 Self {
971 writer,
972 block_buffer: Vec::with_capacity(BLOCK_SIZE),
973 prev_key: Vec::new(),
974 index: Vec::new(),
975 current_offset: 0,
976 num_entries: 0,
977 block_first_key: None,
978 config,
979 dictionary: Some(dictionary),
980 bloom_hashes: Vec::new(),
981 _phantom: std::marker::PhantomData,
982 }
983 }
984
985 pub fn insert(&mut self, key: &[u8], value: &V) -> io::Result<()> {
986 if self.block_first_key.is_none() {
987 self.block_first_key = Some(key.to_vec());
988 }
989
990 if self.config.use_bloom_filter {
992 self.bloom_hashes.push(bloom_hash_pair(key));
993 }
994
995 let prefix_len = common_prefix_len(&self.prev_key, key);
996 let suffix = &key[prefix_len..];
997
998 write_vint(&mut self.block_buffer, prefix_len as u64)?;
999 write_vint(&mut self.block_buffer, suffix.len() as u64)?;
1000 self.block_buffer.extend_from_slice(suffix);
1001 value.serialize(&mut self.block_buffer)?;
1002
1003 self.prev_key.clear();
1004 self.prev_key.extend_from_slice(key);
1005 self.num_entries += 1;
1006
1007 if self.block_buffer.len() >= BLOCK_SIZE {
1008 self.flush_block()?;
1009 }
1010
1011 Ok(())
1012 }
1013
1014 fn flush_block(&mut self) -> io::Result<()> {
1016 if self.block_buffer.is_empty() {
1017 return Ok(());
1018 }
1019
1020 let compressed = if let Some(ref dict) = self.dictionary {
1022 crate::compression::compress_with_dict(
1023 &self.block_buffer,
1024 self.config.compression_level,
1025 dict,
1026 )?
1027 } else {
1028 crate::compression::compress(&self.block_buffer, self.config.compression_level)?
1029 };
1030
1031 if let Some(first_key) = self.block_first_key.take() {
1032 self.index.push(BlockIndexEntry {
1033 first_key,
1034 offset: self.current_offset,
1035 length: compressed.len() as u32,
1036 });
1037 }
1038
1039 self.writer.write_all(&compressed)?;
1040 self.current_offset += compressed.len() as u64;
1041 self.block_buffer.clear();
1042 self.prev_key.clear();
1043
1044 Ok(())
1045 }
1046
1047 pub fn finish(mut self) -> io::Result<W> {
1048 self.flush_block()?;
1050
1051 let bloom_filter = if self.config.use_bloom_filter && !self.bloom_hashes.is_empty() {
1053 let mut bloom =
1054 BloomFilter::new(self.bloom_hashes.len(), self.config.bloom_bits_per_key);
1055 for (h1, h2) in &self.bloom_hashes {
1056 bloom.insert_hashed(*h1, *h2);
1057 }
1058 Some(bloom)
1059 } else {
1060 None
1061 };
1062
1063 let data_end_offset = self.current_offset;
1064
1065 let entries: Vec<(Vec<u8>, BlockAddr)> = self
1068 .index
1069 .iter()
1070 .map(|e| {
1071 (
1072 e.first_key.clone(),
1073 BlockAddr {
1074 offset: e.offset,
1075 length: e.length,
1076 },
1077 )
1078 })
1079 .collect();
1080
1081 #[cfg(feature = "native")]
1083 let index_bytes = FstBlockIndex::build(&entries)?;
1084 #[cfg(not(feature = "native"))]
1085 let index_bytes = MmapBlockIndex::build(&entries)?;
1086
1087 self.writer
1089 .write_u32::<LittleEndian>(index_bytes.len() as u32)?;
1090 self.writer.write_all(&index_bytes)?;
1091 self.current_offset += 4 + index_bytes.len() as u64;
1092
1093 let bloom_offset = if let Some(ref bloom) = bloom_filter {
1095 let offset = self.current_offset;
1096 bloom.write_to(&mut self.writer)?;
1097 self.current_offset += bloom.serialized_len() as u64;
1098 offset
1099 } else {
1100 0
1101 };
1102
1103 let dict_offset = if let Some(ref dict) = self.dictionary {
1105 let dict_bytes = dict.as_bytes();
1106 let offset = self.current_offset;
1107 self.writer
1108 .write_u32::<LittleEndian>(dict_bytes.len() as u32)?;
1109 self.writer.write_all(dict_bytes)?;
1110 self.current_offset += 4 + dict_bytes.len() as u64;
1111 offset
1112 } else {
1113 0
1114 };
1115
1116 self.writer.write_u64::<LittleEndian>(data_end_offset)?;
1118 self.writer.write_u64::<LittleEndian>(self.num_entries)?;
1119 self.writer.write_u64::<LittleEndian>(bloom_offset)?; self.writer.write_u64::<LittleEndian>(dict_offset)?; self.writer
1122 .write_u8(self.config.compression_level.0 as u8)?;
1123 self.writer.write_u32::<LittleEndian>(SSTABLE_MAGIC)?;
1124
1125 Ok(self.writer)
1126 }
1127}
1128
1129#[derive(Debug, Clone)]
1131struct BlockIndexEntry {
1132 first_key: Vec<u8>,
1133 offset: u64,
1134 length: u32,
1135}
1136
1137pub struct AsyncSSTableReader<V: SSTableValue> {
1144 data_slice: FileHandle,
1146 block_index: BlockIndex,
1148 num_entries: u64,
1149 cache: RwLock<BlockCache>,
1151 bloom_filter: Option<BloomFilter>,
1153 dictionary: Option<CompressionDict>,
1155 #[allow(dead_code)]
1157 compression_level: CompressionLevel,
1158 _phantom: std::marker::PhantomData<V>,
1159}
1160
1161struct BlockCache {
1167 blocks: FxHashMap<u64, Arc<[u8]>>,
1168 lru_order: std::collections::VecDeque<u64>,
1169 max_blocks: usize,
1170}
1171
1172impl BlockCache {
1173 fn new(max_blocks: usize) -> Self {
1174 Self {
1175 blocks: FxHashMap::default(),
1176 lru_order: std::collections::VecDeque::with_capacity(max_blocks),
1177 max_blocks,
1178 }
1179 }
1180
1181 fn get(&mut self, offset: u64) -> Option<Arc<[u8]>> {
1182 if self.blocks.contains_key(&offset) {
1183 self.promote(offset);
1184 self.blocks.get(&offset).map(Arc::clone)
1185 } else {
1186 None
1187 }
1188 }
1189
1190 fn peek(&self, offset: u64) -> Option<Arc<[u8]>> {
1192 self.blocks.get(&offset).map(Arc::clone)
1193 }
1194
1195 fn insert(&mut self, offset: u64, block: Arc<[u8]>) {
1196 if self.max_blocks == 0 {
1197 return;
1198 }
1199 if self.blocks.contains_key(&offset) {
1200 self.promote(offset);
1201 return;
1202 }
1203 while self.blocks.len() >= self.max_blocks {
1204 if let Some(evict_offset) = self.lru_order.pop_front() {
1205 self.blocks.remove(&evict_offset);
1206 } else {
1207 break;
1208 }
1209 }
1210 self.blocks.insert(offset, block);
1211 self.lru_order.push_back(offset);
1212 }
1213
1214 fn promote(&mut self, offset: u64) {
1216 if let Some(pos) = self.lru_order.iter().position(|&k| k == offset) {
1217 self.lru_order.remove(pos);
1218 self.lru_order.push_back(offset);
1219 }
1220 }
1221}
1222
1223impl<V: SSTableValue> AsyncSSTableReader<V> {
1224 pub async fn open(file_handle: FileHandle, cache_blocks: usize) -> io::Result<Self> {
1229 let file_len = file_handle.len();
1230 if file_len < 37 {
1231 return Err(io::Error::new(
1232 io::ErrorKind::InvalidData,
1233 "SSTable too small",
1234 ));
1235 }
1236
1237 let footer_bytes = file_handle
1240 .read_bytes_range(file_len - 37..file_len)
1241 .await?;
1242
1243 let mut reader = footer_bytes.as_slice();
1244 let data_end_offset = reader.read_u64::<LittleEndian>()?;
1245 let num_entries = reader.read_u64::<LittleEndian>()?;
1246 let bloom_offset = reader.read_u64::<LittleEndian>()?;
1247 let dict_offset = reader.read_u64::<LittleEndian>()?;
1248 let compression_level = CompressionLevel(reader.read_u8()? as i32);
1249 let magic = reader.read_u32::<LittleEndian>()?;
1250
1251 if magic != SSTABLE_MAGIC {
1252 return Err(io::Error::new(
1253 io::ErrorKind::InvalidData,
1254 format!("Invalid SSTable magic: 0x{:08X}", magic),
1255 ));
1256 }
1257
1258 let footer_start = file_len - 37;
1259 if data_end_offset > footer_start {
1260 return Err(io::Error::new(
1261 io::ErrorKind::InvalidData,
1262 "SSTable data section extends past its footer",
1263 ));
1264 }
1265 if bloom_offset != 0 && (bloom_offset < data_end_offset || bloom_offset >= footer_start) {
1266 return Err(io::Error::new(
1267 io::ErrorKind::InvalidData,
1268 "SSTable bloom filter offset is out of bounds",
1269 ));
1270 }
1271 if dict_offset != 0
1272 && (dict_offset < data_end_offset
1273 || dict_offset >= footer_start
1274 || (bloom_offset != 0 && dict_offset <= bloom_offset))
1275 {
1276 return Err(io::Error::new(
1277 io::ErrorKind::InvalidData,
1278 "SSTable dictionary offset is out of bounds",
1279 ));
1280 }
1281
1282 let index_start = data_end_offset;
1284 let index_end = if bloom_offset != 0 {
1285 bloom_offset
1286 } else if dict_offset != 0 {
1287 dict_offset
1288 } else {
1289 footer_start
1290 };
1291 let index_bytes = file_handle.read_bytes_range(index_start..index_end).await?;
1292
1293 let mut idx_reader = index_bytes.as_slice();
1295 let index_len = idx_reader.read_u32::<LittleEndian>()? as usize;
1296
1297 if index_len != idx_reader.len() {
1298 return Err(io::Error::new(
1299 io::ErrorKind::InvalidData,
1300 "Index data truncated",
1301 ));
1302 }
1303
1304 let index_data = index_bytes.slice(4..4 + index_len);
1305
1306 #[cfg(feature = "fst-index")]
1308 let block_index = match FstBlockIndex::load(index_data.clone()) {
1309 Ok(fst_idx) => BlockIndex::Fst(fst_idx),
1310 Err(_) => BlockIndex::Mmap(MmapBlockIndex::load(index_data)?),
1311 };
1312 #[cfg(not(feature = "fst-index"))]
1313 let block_index = BlockIndex::Mmap(MmapBlockIndex::load(index_data)?);
1314
1315 let mut expected_offset = 0u64;
1316 for addr in block_index.all_addrs() {
1317 let end = addr.offset.checked_add(addr.length as u64).ok_or_else(|| {
1318 io::Error::new(io::ErrorKind::InvalidData, "SSTable block range overflow")
1319 })?;
1320 if addr.length == 0 || addr.offset != expected_offset || end > data_end_offset {
1321 return Err(io::Error::new(
1322 io::ErrorKind::InvalidData,
1323 "SSTable block addresses are inconsistent",
1324 ));
1325 }
1326 expected_offset = end;
1327 }
1328 if expected_offset != data_end_offset {
1329 return Err(io::Error::new(
1330 io::ErrorKind::InvalidData,
1331 "SSTable block addresses do not cover the data section",
1332 ));
1333 }
1334
1335 let bloom_filter = if bloom_offset > 0 {
1337 let bloom_start = bloom_offset;
1338 let bloom_end = if dict_offset != 0 {
1339 dict_offset
1340 } else {
1341 footer_start
1342 };
1343 let header_size = BloomFilter::SERIALIZED_HEADER_SIZE as u64;
1345 let bloom_header_end = bloom_start.checked_add(header_size).ok_or_else(|| {
1346 io::Error::new(io::ErrorKind::InvalidData, "bloom filter range overflow")
1347 })?;
1348 if bloom_header_end > bloom_end {
1349 return Err(io::Error::new(
1350 io::ErrorKind::UnexpectedEof,
1351 "bloom filter header is truncated",
1352 ));
1353 }
1354 let bloom_header = file_handle
1355 .read_bytes_range(bloom_start..bloom_header_end)
1356 .await?;
1357 let num_words = u32::from_le_bytes([
1358 bloom_header[12],
1359 bloom_header[13],
1360 bloom_header[14],
1361 bloom_header[15],
1362 ]) as u64;
1363 let bloom_size = num_words
1364 .checked_mul(8)
1365 .and_then(|bytes| bytes.checked_add(header_size))
1366 .ok_or_else(|| {
1367 io::Error::new(io::ErrorKind::InvalidData, "bloom filter size overflow")
1368 })?;
1369 let actual_bloom_size = bloom_end - bloom_start;
1370 if bloom_size != actual_bloom_size {
1371 return Err(io::Error::new(
1372 io::ErrorKind::InvalidData,
1373 "bloom filter length is inconsistent",
1374 ));
1375 }
1376 let bloom_data = file_handle.read_bytes_range(bloom_start..bloom_end).await?;
1377 Some(BloomFilter::from_owned_bytes(bloom_data)?)
1378 } else {
1379 None
1380 };
1381
1382 let dictionary = if dict_offset > 0 {
1384 let dict_start = dict_offset;
1385 let dict_header_end = dict_start.checked_add(4).ok_or_else(|| {
1387 io::Error::new(io::ErrorKind::InvalidData, "dictionary range overflow")
1388 })?;
1389 if dict_header_end > footer_start {
1390 return Err(io::Error::new(
1391 io::ErrorKind::UnexpectedEof,
1392 "dictionary header is truncated",
1393 ));
1394 }
1395 let dict_len_bytes = file_handle
1396 .read_bytes_range(dict_start..dict_header_end)
1397 .await?;
1398 let dict_len = u32::from_le_bytes([
1399 dict_len_bytes[0],
1400 dict_len_bytes[1],
1401 dict_len_bytes[2],
1402 dict_len_bytes[3],
1403 ]) as u64;
1404 if dict_len > MAX_SSTABLE_DICTIONARY_BYTES {
1405 return Err(io::Error::new(
1406 io::ErrorKind::InvalidData,
1407 "SSTable dictionary exceeds safety limit",
1408 ));
1409 }
1410 let dict_end = dict_header_end.checked_add(dict_len).ok_or_else(|| {
1411 io::Error::new(io::ErrorKind::InvalidData, "dictionary range overflow")
1412 })?;
1413 if dict_end != footer_start {
1414 return Err(io::Error::new(
1415 io::ErrorKind::InvalidData,
1416 "SSTable dictionary length is inconsistent",
1417 ));
1418 }
1419 let dict_data = file_handle
1420 .read_bytes_range(dict_header_end..dict_end)
1421 .await?;
1422 Some(CompressionDict::from_owned_bytes(dict_data))
1423 } else {
1424 None
1425 };
1426
1427 let data_slice = file_handle.slice(0..data_end_offset);
1429
1430 Ok(Self {
1431 data_slice,
1432 block_index,
1433 num_entries,
1434 cache: RwLock::new(BlockCache::new(cache_blocks)),
1435 bloom_filter,
1436 dictionary,
1437 compression_level,
1438 _phantom: std::marker::PhantomData,
1439 })
1440 }
1441
1442 pub fn num_entries(&self) -> u64 {
1444 self.num_entries
1445 }
1446
1447 pub fn stats(&self) -> SSTableStats {
1449 SSTableStats {
1450 num_blocks: self.block_index.len(),
1451 num_sparse_entries: 0, num_entries: self.num_entries,
1453 has_bloom_filter: self.bloom_filter.is_some(),
1454 has_dictionary: self.dictionary.is_some(),
1455 bloom_filter_size: self
1456 .bloom_filter
1457 .as_ref()
1458 .map(|b| b.size_bytes())
1459 .unwrap_or(0),
1460 dictionary_size: self.dictionary.as_ref().map(|d| d.len()).unwrap_or(0),
1461 }
1462 }
1463
1464 pub fn cached_blocks(&self) -> usize {
1466 self.cache.read().blocks.len()
1467 }
1468
1469 pub fn cached_bytes(&self) -> usize {
1475 self.cache
1476 .read()
1477 .blocks
1478 .values()
1479 .map(|block| block.len())
1480 .sum()
1481 }
1482
1483 pub async fn get(&self, key: &[u8]) -> io::Result<Option<V>> {
1488 log::debug!(
1489 "SSTable::get called, key_len={}, total_blocks={}",
1490 key.len(),
1491 self.block_index.len()
1492 );
1493
1494 if let Some(ref bloom) = self.bloom_filter
1496 && !bloom.may_contain(key)
1497 {
1498 log::debug!("SSTable::get bloom filter negative");
1499 return Ok(None);
1500 }
1501
1502 let block_idx = match self.block_index.locate(key) {
1504 Some(idx) => idx,
1505 None => {
1506 log::debug!("SSTable::get key not found (before first block)");
1507 return Ok(None);
1508 }
1509 };
1510
1511 log::debug!("SSTable::get loading block_idx={}", block_idx);
1512
1513 let block_data = self.load_block(block_idx).await?;
1515 self.search_block(&block_data, key)
1516 }
1517
1518 pub async fn get_batch(&self, keys: &[&[u8]]) -> io::Result<Vec<Option<V>>> {
1524 if keys.is_empty() {
1525 return Ok(Vec::new());
1526 }
1527
1528 let mut key_to_block: Vec<(usize, usize)> = Vec::with_capacity(keys.len());
1530 for (key_idx, key) in keys.iter().enumerate() {
1531 if let Some(ref bloom) = self.bloom_filter
1533 && !bloom.may_contain(key)
1534 {
1535 key_to_block.push((key_idx, usize::MAX)); continue;
1537 }
1538
1539 match self.block_index.locate(key) {
1540 Some(block_idx) => key_to_block.push((key_idx, block_idx)),
1541 None => key_to_block.push((key_idx, usize::MAX)), }
1543 }
1544
1545 let mut blocks_to_load: Vec<usize> = key_to_block
1547 .iter()
1548 .filter(|(_, b)| *b != usize::MAX)
1549 .map(|(_, b)| *b)
1550 .collect();
1551 blocks_to_load.sort_unstable();
1552 blocks_to_load.dedup();
1553
1554 for &block_idx in &blocks_to_load {
1556 let _ = self.load_block(block_idx).await?;
1557 }
1558
1559 let mut results = vec![None; keys.len()];
1561 for (key_idx, block_idx) in key_to_block {
1562 if block_idx == usize::MAX {
1563 continue;
1564 }
1565 let block_data = self.load_block(block_idx).await?; results[key_idx] = self.search_block(&block_data, keys[key_idx])?;
1567 }
1568
1569 Ok(results)
1570 }
1571
1572 pub async fn preload_all_blocks(&self) -> io::Result<()> {
1577 for block_idx in 0..self.block_index.len() {
1578 self.load_block(block_idx).await?;
1579 }
1580 Ok(())
1581 }
1582
1583 pub async fn prefetch_all_data_bulk(&self) -> io::Result<()> {
1589 let num_blocks = self.block_index.len();
1590 if num_blocks == 0 {
1591 return Ok(());
1592 }
1593
1594 let mut max_end: u64 = 0;
1596 for i in 0..num_blocks {
1597 if let Some(addr) = self.block_index.get_addr(i) {
1598 let end = addr.offset.checked_add(addr.length as u64).ok_or_else(|| {
1599 io::Error::new(io::ErrorKind::InvalidData, "SSTable block range overflow")
1600 })?;
1601 max_end = max_end.max(end);
1602 }
1603 }
1604
1605 let all_data = self.data_slice.read_bytes_range(0..max_end).await?;
1607 let buf = all_data.as_slice();
1608
1609 let mut cache = self.cache.write();
1611 cache.max_blocks = cache.max_blocks.max(num_blocks);
1612 for i in 0..num_blocks {
1613 let addr = self.block_index.get_addr(i).unwrap();
1614 if cache.get(addr.offset).is_some() {
1615 continue;
1616 }
1617 let start = usize::try_from(addr.offset).map_err(|_| {
1618 io::Error::new(io::ErrorKind::InvalidData, "SSTable block offset too large")
1619 })?;
1620 let end = addr
1621 .offset
1622 .checked_add(addr.length as u64)
1623 .and_then(|end| usize::try_from(end).ok())
1624 .ok_or_else(|| {
1625 io::Error::new(io::ErrorKind::InvalidData, "SSTable block range overflow")
1626 })?;
1627 let compressed = buf.get(start..end).ok_or_else(|| {
1628 io::Error::new(io::ErrorKind::UnexpectedEof, "SSTable block is truncated")
1629 })?;
1630 let decompressed = if let Some(ref dict) = self.dictionary {
1631 crate::compression::decompress_with_dict_limited(
1632 compressed,
1633 dict,
1634 MAX_SSTABLE_BLOCK_BYTES,
1635 )?
1636 } else {
1637 crate::compression::decompress_limited(compressed, MAX_SSTABLE_BLOCK_BYTES)?
1638 };
1639 cache.insert(addr.offset, Arc::from(decompressed));
1640 }
1641
1642 Ok(())
1643 }
1644
1645 async fn load_block(&self, block_idx: usize) -> io::Result<Arc<[u8]>> {
1648 let addr = self.block_index.get_addr(block_idx).ok_or_else(|| {
1649 io::Error::new(io::ErrorKind::InvalidInput, "Block index out of range")
1650 })?;
1651
1652 {
1654 if let Some(block) = self.cache.read().peek(addr.offset) {
1655 return Ok(block);
1656 }
1657 }
1658
1659 log::debug!(
1660 "SSTable::load_block idx={} CACHE MISS, reading bytes [{}-{}]",
1661 block_idx,
1662 addr.offset,
1663 addr.offset + addr.length as u64
1664 );
1665
1666 let range = addr.byte_range();
1668 let compressed = self.data_slice.read_bytes_range(range).await?;
1669
1670 let decompressed = if let Some(ref dict) = self.dictionary {
1672 crate::compression::decompress_with_dict_limited(
1673 compressed.as_slice(),
1674 dict,
1675 MAX_SSTABLE_BLOCK_BYTES,
1676 )?
1677 } else {
1678 crate::compression::decompress_limited(compressed.as_slice(), MAX_SSTABLE_BLOCK_BYTES)?
1679 };
1680
1681 let block: Arc<[u8]> = Arc::from(decompressed);
1682
1683 {
1685 let mut cache = self.cache.write();
1686 cache.insert(addr.offset, Arc::clone(&block));
1687 }
1688
1689 Ok(block)
1690 }
1691
1692 #[cfg(feature = "sync")]
1694 fn load_block_sync(&self, block_idx: usize) -> io::Result<Arc<[u8]>> {
1695 let addr = self.block_index.get_addr(block_idx).ok_or_else(|| {
1696 io::Error::new(io::ErrorKind::InvalidInput, "Block index out of range")
1697 })?;
1698
1699 {
1701 if let Some(block) = self.cache.read().peek(addr.offset) {
1702 return Ok(block);
1703 }
1704 }
1705
1706 let range = addr.byte_range();
1708 let compressed = self.data_slice.read_bytes_range_sync(range)?;
1709
1710 let decompressed = if let Some(ref dict) = self.dictionary {
1712 crate::compression::decompress_with_dict_limited(
1713 compressed.as_slice(),
1714 dict,
1715 MAX_SSTABLE_BLOCK_BYTES,
1716 )?
1717 } else {
1718 crate::compression::decompress_limited(compressed.as_slice(), MAX_SSTABLE_BLOCK_BYTES)?
1719 };
1720
1721 let block: Arc<[u8]> = Arc::from(decompressed);
1722
1723 {
1725 let mut cache = self.cache.write();
1726 cache.insert(addr.offset, Arc::clone(&block));
1727 }
1728
1729 Ok(block)
1730 }
1731
1732 #[cfg(feature = "sync")]
1734 pub fn get_sync(&self, key: &[u8]) -> io::Result<Option<V>> {
1735 if let Some(ref bloom) = self.bloom_filter
1737 && !bloom.may_contain(key)
1738 {
1739 return Ok(None);
1740 }
1741
1742 let block_idx = match self.block_index.locate(key) {
1744 Some(idx) => idx,
1745 None => {
1746 return Ok(None);
1747 }
1748 };
1749
1750 let block_data = self.load_block_sync(block_idx)?;
1751 self.search_block(&block_data, key)
1752 }
1753
1754 fn search_block(&self, block_data: &[u8], target_key: &[u8]) -> io::Result<Option<V>> {
1755 let mut reader = block_data;
1756 let mut current_key = Vec::new();
1757
1758 while !reader.is_empty() {
1759 let value = decode_block_entry(&mut reader, &mut current_key)?;
1760
1761 match current_key.as_slice().cmp(target_key) {
1762 std::cmp::Ordering::Equal => return Ok(Some(value)),
1763 std::cmp::Ordering::Greater => return Ok(None),
1764 std::cmp::Ordering::Less => continue,
1765 }
1766 }
1767
1768 Ok(None)
1769 }
1770
1771 pub async fn prefetch_range(&self, start_key: &[u8], end_key: &[u8]) -> io::Result<()> {
1773 let start_block = self.block_index.locate(start_key).unwrap_or(0);
1774 let end_block = self
1775 .block_index
1776 .locate(end_key)
1777 .unwrap_or(self.block_index.len().saturating_sub(1));
1778
1779 for block_idx in start_block..=end_block.min(self.block_index.len().saturating_sub(1)) {
1780 let _ = self.load_block(block_idx).await?;
1781 }
1782
1783 Ok(())
1784 }
1785
1786 pub fn iter(&self) -> AsyncSSTableIterator<'_, V> {
1788 AsyncSSTableIterator::new(self)
1789 }
1790
1791 pub async fn all_entries(&self) -> io::Result<Vec<(Vec<u8>, V)>> {
1793 let mut results = Vec::new();
1794
1795 for block_idx in 0..self.block_index.len() {
1796 let block_data = self.load_block(block_idx).await?;
1797 let mut reader = &block_data[..];
1798 let mut current_key = Vec::new();
1799
1800 while !reader.is_empty() {
1801 let value = decode_block_entry(&mut reader, &mut current_key)?;
1802 results.push((current_key.clone(), value));
1803 }
1804 }
1805
1806 Ok(results)
1807 }
1808
1809 pub async fn prefix_scan(&self, prefix: &[u8]) -> io::Result<Vec<(Vec<u8>, V)>> {
1815 let (results, _) = self.prefix_scan_limited(prefix, usize::MAX).await?;
1816 Ok(results)
1817 }
1818
1819 pub async fn prefix_scan_limited(
1822 &self,
1823 prefix: &[u8],
1824 max_results: usize,
1825 ) -> io::Result<PrefixScanResult<V>> {
1826 if self.block_index.is_empty() || prefix.is_empty() {
1827 return Ok((Vec::new(), false));
1828 }
1829
1830 let start_block = match self.block_index.locate(prefix) {
1831 Some(idx) => idx,
1832 None => return Ok((Vec::new(), false)),
1833 };
1834
1835 let mut results = Vec::new();
1836
1837 for block_idx in start_block..self.block_index.len() {
1838 let block_data = self.load_block(block_idx).await?;
1839 let mut reader = &block_data[..];
1840 let mut current_key = Vec::new();
1841
1842 while !reader.is_empty() {
1843 let value = decode_block_entry(&mut reader, &mut current_key)?;
1844
1845 if current_key.starts_with(prefix) {
1846 if results.len() >= max_results {
1847 return Ok((results, true));
1848 }
1849 results.push((current_key.clone(), value));
1850 } else if current_key.as_slice() > prefix {
1851 return Ok((results, false));
1853 }
1854 }
1855 }
1856
1857 Ok((results, false))
1858 }
1859
1860 #[cfg(feature = "sync")]
1862 pub fn prefix_scan_sync(&self, prefix: &[u8]) -> io::Result<Vec<(Vec<u8>, V)>> {
1863 let (results, _) = self.prefix_scan_limited_sync(prefix, usize::MAX)?;
1864 Ok(results)
1865 }
1866
1867 #[cfg(feature = "sync")]
1869 pub fn prefix_scan_limited_sync(
1870 &self,
1871 prefix: &[u8],
1872 max_results: usize,
1873 ) -> io::Result<PrefixScanResult<V>> {
1874 if self.block_index.is_empty() || prefix.is_empty() {
1875 return Ok((Vec::new(), false));
1876 }
1877
1878 let start_block = match self.block_index.locate(prefix) {
1879 Some(idx) => idx,
1880 None => return Ok((Vec::new(), false)),
1881 };
1882
1883 let mut results = Vec::new();
1884
1885 for block_idx in start_block..self.block_index.len() {
1886 let block_data = self.load_block_sync(block_idx)?;
1887 let mut reader = &block_data[..];
1888 let mut current_key = Vec::new();
1889
1890 while !reader.is_empty() {
1891 let value = decode_block_entry(&mut reader, &mut current_key)?;
1892
1893 if current_key.starts_with(prefix) {
1894 if results.len() >= max_results {
1895 return Ok((results, true));
1896 }
1897 results.push((current_key.clone(), value));
1898 } else if current_key.as_slice() > prefix {
1899 return Ok((results, false));
1900 }
1901 }
1902 }
1903
1904 Ok((results, false))
1905 }
1906}
1907
1908pub struct AsyncSSTableIterator<'a, V: SSTableValue> {
1910 reader: &'a AsyncSSTableReader<V>,
1911 current_block: usize,
1912 block_data: Option<Arc<[u8]>>,
1913 block_offset: usize,
1914 current_key: Vec<u8>,
1915 finished: bool,
1916}
1917
1918impl<'a, V: SSTableValue> AsyncSSTableIterator<'a, V> {
1919 fn new(reader: &'a AsyncSSTableReader<V>) -> Self {
1920 Self {
1921 reader,
1922 current_block: 0,
1923 block_data: None,
1924 block_offset: 0,
1925 current_key: Vec::new(),
1926 finished: reader.block_index.is_empty(),
1927 }
1928 }
1929
1930 async fn load_next_block(&mut self) -> io::Result<bool> {
1931 if self.current_block >= self.reader.block_index.len() {
1932 self.finished = true;
1933 return Ok(false);
1934 }
1935
1936 self.block_data = Some(self.reader.load_block(self.current_block).await?);
1937 self.block_offset = 0;
1938 self.current_key.clear();
1939 self.current_block += 1;
1940 Ok(true)
1941 }
1942
1943 pub async fn next(&mut self) -> io::Result<Option<(Vec<u8>, V)>> {
1945 if self.finished {
1946 return Ok(None);
1947 }
1948
1949 if self.block_data.is_none() && !self.load_next_block().await? {
1950 return Ok(None);
1951 }
1952
1953 loop {
1954 let block = self.block_data.as_ref().unwrap();
1955 if self.block_offset >= block.len() {
1956 if !self.load_next_block().await? {
1957 return Ok(None);
1958 }
1959 continue;
1960 }
1961
1962 let mut reader = &block[self.block_offset..];
1963 let start_len = reader.len();
1964
1965 let value = decode_block_entry(&mut reader, &mut self.current_key)?;
1966
1967 self.block_offset += start_len - reader.len();
1968
1969 return Ok(Some((self.current_key.clone(), value)));
1970 }
1971 }
1972}
1973
1974#[cfg(test)]
1975mod tests {
1976 use super::*;
1977
1978 #[test]
1979 fn test_bloom_filter_basic() {
1980 let mut bloom = BloomFilter::new(100, 10);
1981
1982 bloom.insert(b"hello");
1983 bloom.insert(b"world");
1984 bloom.insert(b"test");
1985
1986 assert!(bloom.may_contain(b"hello"));
1987 assert!(bloom.may_contain(b"world"));
1988 assert!(bloom.may_contain(b"test"));
1989
1990 assert!(!bloom.may_contain(b"notfound"));
1992 assert!(!bloom.may_contain(b"missing"));
1993 }
1994
1995 #[test]
1996 fn test_bloom_filter_serialization() {
1997 let mut bloom = BloomFilter::new(100, 10);
1998 bloom.insert(b"key1");
1999 bloom.insert(b"key2");
2000
2001 let bytes = bloom.to_bytes();
2002 let restored = BloomFilter::from_owned_bytes(OwnedBytes::new(bytes)).unwrap();
2003
2004 assert!(restored.may_contain(b"key1"));
2005 assert!(restored.may_contain(b"key2"));
2006 assert!(!restored.may_contain(b"key3"));
2007 }
2008
2009 #[test]
2010 fn bloom_header_preserves_bit_counts_above_u32() {
2011 let num_bits = u32::MAX as usize + 1;
2012 let mut header = Vec::new();
2013 write_bloom_header(&mut header, num_bits, BLOOM_HASH_COUNT, 1).unwrap();
2014 assert_eq!(header.len(), BLOOM_FILTER_HEADER_SIZE);
2015 assert_eq!(
2016 u64::from_le_bytes(header[0..8].try_into().unwrap()),
2017 num_bits as u64
2018 );
2019 }
2020
2021 #[test]
2022 fn test_bloom_filter_false_positive_rate() {
2023 let num_keys = 10000;
2024 let mut bloom = BloomFilter::new(num_keys, BLOOM_BITS_PER_KEY);
2025
2026 for i in 0..num_keys {
2028 let key = format!("key_{}", i);
2029 bloom.insert(key.as_bytes());
2030 }
2031
2032 for i in 0..num_keys {
2034 let key = format!("key_{}", i);
2035 assert!(bloom.may_contain(key.as_bytes()));
2036 }
2037
2038 let mut false_positives = 0;
2040 let test_count = 10000;
2041 for i in 0..test_count {
2042 let key = format!("nonexistent_{}", i);
2043 if bloom.may_contain(key.as_bytes()) {
2044 false_positives += 1;
2045 }
2046 }
2047
2048 let fp_rate = false_positives as f64 / test_count as f64;
2051 assert!(
2052 fp_rate < 0.03,
2053 "False positive rate {} is too high",
2054 fp_rate
2055 );
2056 }
2057
2058 #[test]
2059 fn test_sstable_writer_config() {
2060 use crate::structures::IndexOptimization;
2061
2062 let config = SSTableWriterConfig::default();
2064 assert_eq!(config.compression_level.0, 9); assert!(config.use_bloom_filter); assert!(!config.use_dictionary);
2067
2068 let adaptive = SSTableWriterConfig::from_optimization(IndexOptimization::Adaptive);
2070 assert_eq!(adaptive.compression_level.0, 9);
2071 assert!(adaptive.use_bloom_filter);
2072 assert!(!adaptive.use_dictionary);
2073
2074 let size = SSTableWriterConfig::from_optimization(IndexOptimization::SizeOptimized);
2076 assert_eq!(size.compression_level.0, 22); assert!(size.use_bloom_filter);
2078 assert!(size.use_dictionary);
2079
2080 let perf = SSTableWriterConfig::from_optimization(IndexOptimization::PerformanceOptimized);
2082 assert_eq!(perf.compression_level.0, 1); assert!(perf.use_bloom_filter); assert!(!perf.use_dictionary);
2085
2086 let fast = SSTableWriterConfig::fast();
2088 assert_eq!(fast.compression_level.0, 1);
2089
2090 let max = SSTableWriterConfig::max_compression();
2091 assert_eq!(max.compression_level.0, 22);
2092 }
2093
2094 #[test]
2095 fn test_vint_roundtrip() {
2096 let test_values = [0u64, 1, 127, 128, 255, 256, 16383, 16384, u64::MAX];
2097
2098 for &val in &test_values {
2099 let mut buf = Vec::new();
2100 write_vint(&mut buf, val).unwrap();
2101 let mut reader = buf.as_slice();
2102 let decoded = read_vint(&mut reader).unwrap();
2103 assert_eq!(val, decoded, "Failed for value {}", val);
2104 }
2105 }
2106
2107 #[test]
2108 fn test_common_prefix_len() {
2109 assert_eq!(common_prefix_len(b"hello", b"hello"), 5);
2110 assert_eq!(common_prefix_len(b"hello", b"help"), 3);
2111 assert_eq!(common_prefix_len(b"hello", b"world"), 0);
2112 assert_eq!(common_prefix_len(b"", b"hello"), 0);
2113 assert_eq!(common_prefix_len(b"hello", b""), 0);
2114 }
2115
2116 #[test]
2117 fn decode_block_entry_reconstructs_keys_and_consumes_one_entry() {
2118 let mut encoded = Vec::new();
2119
2120 write_vint(&mut encoded, 0).unwrap();
2121 write_vint(&mut encoded, 5).unwrap();
2122 encoded.extend_from_slice(b"alpha");
2123 7_u64.serialize(&mut encoded).unwrap();
2124
2125 write_vint(&mut encoded, 3).unwrap();
2126 write_vint(&mut encoded, 3).unwrap();
2127 encoded.extend_from_slice(b"ine");
2128 11_u64.serialize(&mut encoded).unwrap();
2129
2130 let mut reader = encoded.as_slice();
2131 let mut key = Vec::new();
2132
2133 assert_eq!(decode_block_entry::<u64>(&mut reader, &mut key).unwrap(), 7);
2134 assert_eq!(key, b"alpha");
2135 assert!(!reader.is_empty(), "the second entry must remain unread");
2136
2137 assert_eq!(
2138 decode_block_entry::<u64>(&mut reader, &mut key).unwrap(),
2139 11
2140 );
2141 assert_eq!(key, b"alpine");
2142 assert!(reader.is_empty());
2143 }
2144
2145 #[test]
2146 fn decode_block_entry_rejects_truncated_suffix() {
2147 let encoded = [0, 4, b'o', b'n'];
2148 let mut reader = encoded.as_slice();
2149 let mut key = Vec::new();
2150
2151 let error = decode_block_entry::<u64>(&mut reader, &mut key).unwrap_err();
2152
2153 assert_eq!(error.kind(), io::ErrorKind::UnexpectedEof);
2154 assert_eq!(error.to_string(), "SSTable block suffix truncated");
2155 }
2156}