use std::path::Path; use anyhow::Result; use bytes::BufMut; use super::{BlockMeta, FileObject, SsTable}; use crate::block::BlockBuilder; /// Builds an SSTable from key-value pairs. pub struct SsTableBuilder { builder: BlockBuilder, first_key: Vec, data: Vec, pub(super) meta: Vec, block_size: usize, } impl SsTableBuilder { /// Create a builder based on target block size. pub fn new(block_size: usize) -> Self { Self { data: Vec::new(), meta: Vec::new(), first_key: Vec::new(), block_size, builder: BlockBuilder::new(block_size), } } /// Adds a key-value pair to SSTable pub fn add(&mut self, key: &[u8], value: &[u8]) { if self.first_key.is_empty() { self.first_key = key.to_vec(); } if self.builder.add(key, value) { return; } // create a new block builder and append block data self.finish_block(); // add the key-value pair to the next block assert!(self.builder.add(key, value)); self.first_key = key.to_vec(); } /// Get the estimated size of the SSTable. pub fn estimated_size(&self) -> usize { self.data.len() } fn finish_block(&mut self) { let builder = std::mem::replace(&mut self.builder, BlockBuilder::new(self.block_size)); let encoded_block = builder.build().encode(); self.meta.push(BlockMeta { offset: self.data.len(), first_key: std::mem::take(&mut self.first_key).into(), }); self.data.extend(encoded_block); } /// Builds the SSTable and writes it to the given path. No need to actually write to disk until /// chapter 4 block cache. pub fn build(mut self, path: impl AsRef) -> Result { self.finish_block(); let mut buf = self.data; let meta_offset = buf.len(); BlockMeta::encode_block_meta(&self.meta, &mut buf); buf.put_u32(meta_offset as u32); let file = FileObject::create(path.as_ref(), buf)?; Ok(SsTable { file, block_metas: self.meta, block_meta_offset: meta_offset, }) } }