typed_store/rocks/
sst_file.rs

1// Copyright (c) 2026 IOTA Stiftung
2// SPDX-License-Identifier: Apache-2.0
3
4use std::{marker::PhantomData, path::Path};
5
6use serde::Serialize;
7use typed_store_error::TypedStoreError;
8
9use crate::{
10    be_fix_int_ser,
11    rocks::{
12        TaggedDBMap,
13        errors::{typed_store_err_from_bcs_err, typed_store_err_from_rocks_err},
14    },
15};
16
17/// Rows of one table, held in memory until [`Self::write_file`] writes them,
18/// sorted, to an SST file for [`DBMap::ingest_sst_files`].
19///
20/// Writing rows this way bypasses the memtables, which RocksDB flushes one at
21/// a time per column family, so many threads can each fill their own files
22/// at once. Keys and values are encoded as [`DBMap`] encodes them.
23///
24/// [`DBMap`]: crate::rocks::DBMap
25/// [`DBMap::ingest_sst_files`]: crate::rocks::DBMap::ingest_sst_files
26pub struct SstFileBuffer<K, V> {
27    /// Every row's key followed by its value.
28    bytes: Vec<u8>,
29    /// Where each row's key starts, where its value starts and where it ends
30    /// in `bytes`.
31    rows: Vec<(usize, usize, usize)>,
32    _phantom: PhantomData<fn(K, V)>,
33}
34
35impl<K: Serialize, V: Serialize> SstFileBuffer<K, V> {
36    pub fn new() -> Self {
37        Self {
38            bytes: Vec::new(),
39            rows: Vec::new(),
40            _phantom: PhantomData,
41        }
42    }
43
44    pub fn insert(&mut self, key: &K, value: &V) -> Result<(), TypedStoreError> {
45        self.push(key, value)
46    }
47
48    fn push<S: Serialize + ?Sized>(&mut self, key: &S, value: &V) -> Result<(), TypedStoreError> {
49        let start = self.bytes.len();
50        self.bytes.extend_from_slice(&be_fix_int_ser(key));
51        let value_start = self.bytes.len();
52        bcs::serialize_into(&mut self.bytes, value).map_err(typed_store_err_from_bcs_err)?;
53        self.rows.push((start, value_start, self.bytes.len()));
54        Ok(())
55    }
56
57    pub fn is_empty(&self) -> bool {
58        self.rows.is_empty()
59    }
60
61    /// The memory the buffered rows take.
62    pub fn size_in_bytes(&self) -> usize {
63        self.bytes.len() + self.rows.len() * size_of::<(usize, usize, usize)>()
64    }
65
66    /// Writes the buffered rows, sorted by key, to a new SST file at `path`
67    /// and empties the buffer. A key inserted more than once keeps its last
68    /// value.
69    ///
70    /// `options` must be those of the column family the file is ingested
71    /// into. The buffer must not be empty, since RocksDB cannot write an SST
72    /// file without rows.
73    pub fn write_file(
74        &mut self,
75        path: &Path,
76        options: &rocksdb::Options,
77    ) -> Result<(), TypedStoreError> {
78        let bytes = &self.bytes;
79        let key = |&(start, value_start, _): &(usize, usize, usize)| &bytes[start..value_start];
80        // Stable, so that of equal keys the last inserted stays last.
81        self.rows.sort_by(|a, b| key(a).cmp(key(b)));
82
83        let mut writer = rocksdb::SstFileWriter::create(options);
84        writer.open(path).map_err(typed_store_err_from_rocks_err)?;
85        for (index, row) in self.rows.iter().enumerate() {
86            if self
87                .rows
88                .get(index + 1)
89                .is_some_and(|next| key(next) == key(row))
90            {
91                continue;
92            }
93            let (_, value_start, end) = *row;
94            writer
95                .put(key(row), &bytes[value_start..end])
96                .map_err(typed_store_err_from_rocks_err)?;
97        }
98        writer.finish().map_err(typed_store_err_from_rocks_err)?;
99
100        self.bytes.clear();
101        self.rows.clear();
102        Ok(())
103    }
104}
105
106impl<K: Serialize, V: Serialize> SstFileBuffer<(u8, K), V> {
107    /// Adds a row of `map`, keyed as `map` keys its rows, for a file ingested
108    /// into `map`'s column family.
109    pub fn insert_tagged(
110        &mut self,
111        map: &TaggedDBMap<K, V>,
112        key: &K,
113        value: &V,
114    ) -> Result<(), TypedStoreError> {
115        self.push(&(map.tag(), key), value)
116    }
117}
118
119impl<K: Serialize, V: Serialize> Default for SstFileBuffer<K, V> {
120    fn default() -> Self {
121        Self::new()
122    }
123}
typed_store/rocks/sst_file.rs

typed_store/rocks/
sst_file.rs
