lsm_tree/ingestion.rs
1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2024-present, fjall-rs
3// Copyright (c) 2026-present, Dmitry Prudnikov
4
5use crate::{
6 AnyTree, UserKey, UserValue, blob_tree::ingest::BlobIngestion, tree::ingest::Ingestion,
7};
8
9/// Unified ingestion builder over `AnyTree`
10// Keep zero allocations and direct dispatch; boxing introduces heap indirection and `dyn` adds virtual dispatch.
11// Ingestion calls use `&mut self` in tight loops; the active variant is stable and branch prediction makes the match cheap.
12// Allowing this lint preserves hot-path performance at the cost of a larger enum size.
13#[expect(clippy::large_enum_variant)]
14pub enum AnyIngestion<'a> {
15 /// Ingestion for a standard LSM-tree
16 Standard(Ingestion<'a>),
17
18 /// Ingestion for a [`BlobTree`](crate::BlobTree) with KV separation
19 Blob(BlobIngestion<'a>),
20}
21
22impl AnyIngestion<'_> {
23 /// Tells the ingestion how many entries it will write in all, an upper
24 /// bound. An ingestion streams its entries, so without it the filters of
25 /// its first tables can take the room a filter budget
26 /// ([`Config::filter_advisor`](crate::Config::filter_advisor)) has left
27 /// and push the later tables' filters past it; with it, each filter keeps
28 /// room for the entries still to come, as a flush's do. Without a filter
29 /// advisor it changes nothing.
30 ///
31 /// # Examples
32 ///
33 /// ```
34 /// # use lsm_tree::Config;
35 /// # let folder = tempfile::tempdir()?;
36 /// # let tree = Config::new(folder, Default::default(), Default::default()).open()?;
37 /// #
38 /// let mut ingestion = tree.ingestion()?.expected_entries(2);
39 /// ingestion.write("a", "abc")?;
40 /// ingestion.write("b", "def")?;
41 /// ingestion.finish()?;
42 /// #
43 /// # Ok::<(), lsm_tree::Error>(())
44 /// ```
45 #[must_use]
46 pub fn expected_entries(self, entries: u64) -> Self {
47 let table = match &self {
48 Self::Standard(i) => i,
49 Self::Blob(b) => &b.table,
50 };
51 if let Some(sizing) = table.writer.filter_sizing() {
52 sizing.expect_keys(entries);
53 }
54 self
55 }
56
57 /// Writes a key-value pair.
58 ///
59 /// # Errors
60 ///
61 /// Will return `Err` if an IO error occurs.
62 pub fn write<K: Into<UserKey>, V: Into<UserValue>>(
63 &mut self,
64 key: K,
65 value: V,
66 ) -> crate::Result<()> {
67 match self {
68 Self::Standard(i) => i.write(key.into(), value.into()),
69 Self::Blob(b) => b.write(key.into(), value.into()),
70 }
71 }
72
73 /// Writes a tombstone for a key.
74 ///
75 /// # Errors
76 ///
77 /// Will return `Err` if an IO error occurs.
78 pub fn write_tombstone<K: Into<UserKey>>(&mut self, key: K) -> crate::Result<()> {
79 match self {
80 Self::Standard(i) => i.write_tombstone(key.into()),
81 Self::Blob(b) => b.write_tombstone(key.into()),
82 }
83 }
84
85 /// Writes a weak tombstone for a key.
86 ///
87 /// # Examples
88 ///
89 /// ```
90 /// # use lsm_tree::Config;
91 /// # let folder = tempfile::tempdir()?;
92 /// # let tree = Config::new(folder, Default::default(), Default::default()).open()?;
93 /// #
94 /// let mut ingestion = tree.ingestion()?;
95 /// ingestion.write("a", "abc")?;
96 /// ingestion.write_weak_tombstone("b")?;
97 /// ingestion.finish()?;
98 /// #
99 /// # Ok::<(), lsm_tree::Error>(())
100 /// ```
101 ///
102 /// # Errors
103 ///
104 /// Will return `Err` if an IO error occurs.
105 pub fn write_weak_tombstone<K: Into<UserKey>>(&mut self, key: K) -> crate::Result<()> {
106 match self {
107 Self::Standard(i) => i.write_weak_tombstone(key.into()),
108 Self::Blob(b) => b.write_weak_tombstone(key.into()),
109 }
110 }
111
112 /// Writes a consumer-provided columnar batch (its value sub-columns) as
113 /// columnar row groups, stored directly without re-transposing the value.
114 /// Small batches of one layout share a group; a batch past the row group
115 /// size is cut into groups of that size.
116 ///
117 /// The batch carries the three intrinsic columns (`[key, seqno, value-type]`)
118 /// plus one or more value sub-columns. Its keys must be strictly increasing
119 /// (by the tree comparator) within the batch and after any previously written
120 /// data, and every per-row seqno must be `0`: [`finish`](Self::finish) assigns
121 /// the atomic global sequence number. The columnar layout must be enabled
122 /// (`columnar` in the runtime config) on a standard tree; a row-mode or blob
123 /// tree rejects the batch.
124 ///
125 /// # Errors
126 ///
127 /// Returns an error if the batch shape is invalid, the keys are not strictly
128 /// increasing, any row carries a non-zero seqno, the layout is not columnar,
129 /// a block write fails, or the tree is a blob tree (columnar ingest does not
130 /// support KV separation).
131 ///
132 /// # Examples
133 ///
134 /// ```
135 /// use lsm_tree::table::columnar::{Column, TypeTag, entries_to_column_batch};
136 /// use lsm_tree::{AnyTree, Config, InternalValue, ValueType};
137 ///
138 /// let folder = tempfile::tempdir()?;
139 /// let any = Config::new(folder, Default::default(), Default::default()).open()?;
140 /// if let AnyTree::Standard(tree) = &any {
141 /// tree.update_runtime_config(|cfg| cfg.columnar = true)?;
142 /// }
143 ///
144 /// // One row whose value is a single fixed-4 sub-column (id 3).
145 /// let mut batch = entries_to_column_batch(&[InternalValue::from_components(
146 /// b"k0".to_vec(),
147 /// b"x".to_vec(),
148 /// 0,
149 /// ValueType::Value,
150 /// )])?;
151 /// batch.columns.pop();
152 /// batch.columns.push(Column {
153 /// column_id: 3,
154 /// type_tag: TypeTag::Fixed(4),
155 /// validity: None,
156 /// data: vec![1, 0, 0, 0].into(),
157 /// });
158 ///
159 /// let mut ingestion = any.ingestion()?;
160 /// ingestion.write_columnar_batch(&batch)?;
161 /// ingestion.finish()?;
162 /// # Ok::<(), lsm_tree::Error>(())
163 /// ```
164 #[cfg(feature = "columnar")]
165 pub fn write_columnar_batch(
166 &mut self,
167 batch: &crate::table::columnar::ColumnBatch,
168 ) -> crate::Result<()> {
169 match self {
170 Self::Standard(i) => i.write_columnar_batch(batch),
171 Self::Blob(_) => Err(crate::Error::FeatureUnsupported(
172 "columnar batch ingest is not supported for blob trees",
173 )),
174 }
175 }
176
177 /// Finalizes ingestion and registers created tables (and blob files if present).
178 ///
179 /// # Errors
180 ///
181 /// Will return `Err` if an IO error occurs.
182 pub fn finish(self) -> crate::Result<()> {
183 match self {
184 Self::Standard(i) => i.finish(),
185 Self::Blob(b) => b.finish(),
186 }
187 }
188}
189
190impl AnyTree {
191 /// Starts an ingestion for any tree type (standard or blob).
192 ///
193 /// # Errors
194 ///
195 /// Will return `Err` if an IO error occurs.
196 pub fn ingestion(&self) -> crate::Result<AnyIngestion<'_>> {
197 match self {
198 Self::Standard(t) => Ok(AnyIngestion::Standard(Ingestion::new(t)?)),
199 Self::Blob(b) => Ok(AnyIngestion::Blob(BlobIngestion::new(b)?)),
200 }
201 }
202
203 /// Runs a projected columnar scan across the tree.
204 ///
205 /// Delegates to [`Tree::columnar_scan`](crate::Tree::columnar_scan) on a
206 /// standard tree and to [`BlobTree::columnar_scan`](crate::BlobTree::columnar_scan)
207 /// on a blob tree: reads the memtables and every table intersecting
208 /// `range` at `seqno`, and yields projected
209 /// [`ColumnBatch`](crate::table::columnar::ColumnBatch)es in key order, one
210 /// row per key, its newest visible version. See
211 /// [`Tree::columnar_scan`](crate::Tree::columnar_scan) for the full contract.
212 ///
213 /// # Errors
214 ///
215 /// As [`Tree::columnar_scan`](crate::Tree::columnar_scan) and
216 /// [`BlobTree::columnar_scan`](crate::BlobTree::columnar_scan).
217 ///
218 /// # Examples
219 ///
220 /// ```
221 /// use lsm_tree::table::columnar::{Column, TypeTag, entries_to_column_batch};
222 /// use lsm_tree::{AnyTree, Config, InternalValue, SeqNo, ValueType};
223 ///
224 /// let folder = tempfile::tempdir()?;
225 /// let any = Config::new(folder, Default::default(), Default::default()).open()?;
226 /// if let AnyTree::Standard(tree) = &any {
227 /// tree.update_runtime_config(|cfg| cfg.columnar = true)?;
228 /// }
229 ///
230 /// // Ingest one row whose value is a single fixed-4 sub-column (id 3).
231 /// let mut batch = entries_to_column_batch(&[InternalValue::from_components(
232 /// b"k0".to_vec(),
233 /// b"x".to_vec(),
234 /// 0,
235 /// ValueType::Value,
236 /// )])?;
237 /// batch.columns.pop();
238 /// batch.columns.push(Column {
239 /// column_id: 3,
240 /// type_tag: TypeTag::Fixed(4),
241 /// validity: None,
242 /// data: vec![1, 0, 0, 0].into(),
243 /// });
244 /// let mut ingestion = any.ingestion()?;
245 /// ingestion.write_columnar_batch(&batch)?;
246 /// ingestion.finish()?;
247 ///
248 /// // Scan projecting only sub-column 3 across the whole tree.
249 /// let mut rows = 0;
250 /// for batch in any.columnar_scan(&[3], None, SeqNo::MAX, ..)? {
251 /// rows += batch?.row_count;
252 /// }
253 /// assert_eq!(rows, 1);
254 /// # Ok::<(), lsm_tree::Error>(())
255 /// ```
256 #[cfg(feature = "columnar")]
257 pub fn columnar_scan<R: core::ops::RangeBounds<crate::UserKey>>(
258 &self,
259 projection: impl Into<crate::Projection>,
260 predicate: Option<&crate::table::columnar_predicate::ColumnRangePredicate>,
261 seqno: crate::SeqNo,
262 range: R,
263 ) -> crate::Result<crate::tree::columnar_scan::ColumnarScan> {
264 match self {
265 Self::Standard(t) => t.columnar_scan(projection, predicate, seqno, range),
266 Self::Blob(t) => t.columnar_scan(projection, predicate, seqno, range),
267 }
268 }
269}