lance_table/system_index/
mem_wal.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::collections::HashMap;
5
6use lance_core::Error;
7use lance_core::deepsize::DeepSizeOf;
8use serde::{Deserialize, Serialize};
9use uuid::Uuid;
10
11use crate::format::pb;
12
13pub const MEM_WAL_INDEX_NAME: &str = "__lance_mem_wal";
14
15/// Type alias for shard identifier (UUID v4).
16pub type ShardId = Uuid;
17
18/// An SSTable: the immutable result of flushing a MemTable, stored as a Lance dataset.
19#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, DeepSizeOf)]
20pub struct SsTable {
21    pub generation: u64,
22    pub path: String,
23}
24
25impl From<&SsTable> for pb::SsTable {
26    fn from(sstable: &SsTable) -> Self {
27        Self {
28            generation: sstable.generation,
29            path: sstable.path.clone(),
30        }
31    }
32}
33
34impl From<pb::SsTable> for SsTable {
35    fn from(sstable: pb::SsTable) -> Self {
36        Self {
37            generation: sstable.generation,
38            path: sstable.path,
39        }
40    }
41}
42
43/// A pointer to the latest SSTable compacted for a shard.
44#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash, Serialize, Deserialize)]
45pub struct CompactedSsTable {
46    pub shard_id: Uuid,
47    pub generation: u64,
48}
49
50impl DeepSizeOf for CompactedSsTable {
51    fn deep_size_of_children(&self, _context: &mut lance_core::deepsize::Context) -> usize {
52        0 // UUID is 16 bytes fixed size, no heap allocations
53    }
54}
55
56impl CompactedSsTable {
57    pub fn new(shard_id: Uuid, generation: u64) -> Self {
58        Self {
59            shard_id,
60            generation,
61        }
62    }
63}
64
65impl From<&CompactedSsTable> for pb::CompactedSsTable {
66    fn from(sstable: &CompactedSsTable) -> Self {
67        Self {
68            shard_id: Some((&sstable.shard_id).into()),
69            generation: sstable.generation,
70        }
71    }
72}
73
74impl TryFrom<pb::CompactedSsTable> for CompactedSsTable {
75    type Error = Error;
76
77    fn try_from(sstable: pb::CompactedSsTable) -> lance_core::Result<Self> {
78        let shard_id = sstable
79            .shard_id
80            .as_ref()
81            .map(Uuid::try_from)
82            .ok_or_else(|| Error::invalid_input("Missing shard_id in CompactedSsTable"))??;
83        Ok(Self {
84            shard_id,
85            generation: sstable.generation,
86        })
87    }
88}
89
90/// Tracks which compacted SSTable generation a base table index covers.
91/// Used to determine whether to read from SSTable indexes or base table.
92#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize, DeepSizeOf)]
93pub struct IndexCatchupProgress {
94    pub index_name: String,
95    pub caught_up_generations: Vec<CompactedSsTable>,
96}
97
98impl IndexCatchupProgress {
99    pub fn new(index_name: String, caught_up_generations: Vec<CompactedSsTable>) -> Self {
100        Self {
101            index_name,
102            caught_up_generations,
103        }
104    }
105
106    /// Get the caught up generation for a specific shard.
107    /// Returns None if the shard is not present (assumed fully caught up).
108    pub fn caught_up_generation_for_shard(&self, shard_id: &Uuid) -> Option<u64> {
109        self.caught_up_generations
110            .iter()
111            .find(|sstable| &sstable.shard_id == shard_id)
112            .map(|sstable| sstable.generation)
113    }
114}
115
116impl From<&IndexCatchupProgress> for pb::IndexCatchupProgress {
117    fn from(icp: &IndexCatchupProgress) -> Self {
118        Self {
119            index_name: icp.index_name.clone(),
120            caught_up_generations: icp
121                .caught_up_generations
122                .iter()
123                .map(|sstable| sstable.into())
124                .collect(),
125        }
126    }
127}
128
129impl TryFrom<pb::IndexCatchupProgress> for IndexCatchupProgress {
130    type Error = Error;
131
132    fn try_from(icp: pb::IndexCatchupProgress) -> lance_core::Result<Self> {
133        Ok(Self {
134            index_name: icp.index_name,
135            caught_up_generations: icp
136                .caught_up_generations
137                .into_iter()
138                .map(CompactedSsTable::try_from)
139                .collect::<lance_core::Result<_>>()?,
140        })
141    }
142}
143
144/// Lifecycle status of a WAL shard, persisted in [`ShardManifest`].
145///
146/// `Sealed` is the durable in-doubt record for drop-table two-phase
147/// commit: a sealed shard refuses new writer claims (enforced in
148/// `claim_epoch`) but is reversible back to `Active` on rollback.
149#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
150pub enum ShardStatus {
151    /// Normal: the shard accepts writer claims.
152    #[default]
153    Active,
154    /// A drop is in flight: claims are refused. Reversible.
155    Sealed,
156}
157
158impl ShardStatus {
159    /// Map to the protobuf enum discriminant (`pb::ShardStatus`).
160    fn to_i32(self) -> i32 {
161        match self {
162            Self::Active => 0,
163            Self::Sealed => 1,
164        }
165    }
166
167    /// Map from the protobuf enum discriminant; unknown values decode as
168    /// `Active` (forward-compatible default).
169    fn from_i32(v: i32) -> Self {
170        match v {
171            1 => Self::Sealed,
172            _ => Self::Active,
173        }
174    }
175}
176
177/// Shard manifest containing epoch-based fencing and WAL state.
178/// Each shard has exactly one active writer at any time.
179#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
180pub struct ShardManifest {
181    pub shard_id: Uuid,
182    pub version: u64,
183    pub shard_spec_id: u32,
184    /// Computed shard field values as raw Arrow scalar bytes, keyed by field id.
185    /// The byte encoding follows Arrow's little-endian convention: int32 is 4 LE
186    /// bytes, utf8 is raw UTF-8 bytes, etc. The result_type in the corresponding
187    /// ShardingField from the ShardingSpec determines how to interpret each value.
188    pub shard_field_values: HashMap<String, Vec<u8>>,
189    pub writer_epoch: u64,
190    /// The most recent WAL entry position flushed to a MemTable.
191    /// Recovery replays from `replay_after_wal_entry_position + 1`. The
192    /// default value 0 means "no flush has ever stamped this shard" — WAL
193    /// positions themselves are 1-based, so 0 is never a valid covered
194    /// position.
195    pub replay_after_wal_entry_position: u64,
196    /// The most recent WAL entry position observed at manifest write time.
197    /// Default 0 means "no entry has been written yet"; WAL positions are
198    /// 1-based.
199    pub wal_entry_position_last_seen: u64,
200    pub current_generation: u64,
201    pub sstables: Vec<SsTable>,
202    /// Lifecycle status (drop-table 2PC). Defaults to `Active`; preserved
203    /// across claims via `..base` so only fresh constructions set it.
204    pub status: ShardStatus,
205}
206
207impl DeepSizeOf for ShardManifest {
208    fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize {
209        self.shard_field_values.deep_size_of_children(context)
210            + self.sstables.deep_size_of_children(context)
211    }
212}
213
214impl From<&ShardManifest> for pb::ShardManifest {
215    fn from(rm: &ShardManifest) -> Self {
216        Self {
217            shard_id: Some((&rm.shard_id).into()),
218            version: rm.version,
219            shard_spec_id: rm.shard_spec_id,
220            shard_field_entries: rm
221                .shard_field_values
222                .iter()
223                .map(|(k, v)| pb::ShardFieldEntry {
224                    field_id: k.clone(),
225                    value: v.clone(),
226                })
227                .collect(),
228            writer_epoch: rm.writer_epoch,
229            replay_after_wal_entry_position: rm.replay_after_wal_entry_position,
230            wal_entry_position_last_seen: rm.wal_entry_position_last_seen,
231            current_generation: rm.current_generation,
232            sstables: rm.sstables.iter().map(|sstable| sstable.into()).collect(),
233            status: rm.status.to_i32(),
234        }
235    }
236}
237
238impl TryFrom<pb::ShardManifest> for ShardManifest {
239    type Error = Error;
240
241    fn try_from(rm: pb::ShardManifest) -> lance_core::Result<Self> {
242        let shard_id = rm
243            .shard_id
244            .as_ref()
245            .map(Uuid::try_from)
246            .ok_or_else(|| Error::invalid_input("Missing shard_id in ShardManifest"))??;
247        let shard_field_values = rm
248            .shard_field_entries
249            .into_iter()
250            .map(|e| (e.field_id, e.value))
251            .collect();
252        Ok(Self {
253            shard_id,
254            version: rm.version,
255            shard_spec_id: rm.shard_spec_id,
256            shard_field_values,
257            writer_epoch: rm.writer_epoch,
258            replay_after_wal_entry_position: rm.replay_after_wal_entry_position,
259            wal_entry_position_last_seen: rm.wal_entry_position_last_seen,
260            current_generation: rm.current_generation,
261            sstables: rm.sstables.into_iter().map(SsTable::from).collect(),
262            status: ShardStatus::from_i32(rm.status),
263        })
264    }
265}
266
267/// Sharding field definition.
268#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, DeepSizeOf)]
269pub struct ShardingField {
270    pub field_id: String,
271    pub source_ids: Vec<i32>,
272    pub transform: Option<String>,
273    pub expression: Option<String>,
274    pub result_type: String,
275    pub parameters: HashMap<String, String>,
276}
277
278impl From<&ShardingField> for pb::ShardingField {
279    fn from(rf: &ShardingField) -> Self {
280        Self {
281            field_id: rf.field_id.clone(),
282            source_ids: rf.source_ids.clone(),
283            transform: rf.transform.clone(),
284            expression: rf.expression.clone(),
285            result_type: rf.result_type.clone(),
286            parameters: rf.parameters.clone(),
287        }
288    }
289}
290
291impl From<pb::ShardingField> for ShardingField {
292    fn from(rf: pb::ShardingField) -> Self {
293        Self {
294            field_id: rf.field_id,
295            source_ids: rf.source_ids,
296            transform: rf.transform,
297            expression: rf.expression,
298            result_type: rf.result_type,
299            parameters: rf.parameters,
300        }
301    }
302}
303
304/// Sharding spec definition.
305#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, DeepSizeOf)]
306pub struct ShardingSpec {
307    pub spec_id: u32,
308    pub fields: Vec<ShardingField>,
309}
310
311impl From<&ShardingSpec> for pb::ShardingSpec {
312    fn from(rs: &ShardingSpec) -> Self {
313        Self {
314            spec_id: rs.spec_id,
315            fields: rs.fields.iter().map(|f| f.into()).collect(),
316        }
317    }
318}
319
320impl From<pb::ShardingSpec> for ShardingSpec {
321    fn from(rs: pb::ShardingSpec) -> Self {
322        Self {
323            spec_id: rs.spec_id,
324            fields: rs.fields.into_iter().map(ShardingField::from).collect(),
325        }
326    }
327}
328
329/// Index details for MemWAL Index, stored in IndexMetadata.index_details.
330#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize, DeepSizeOf)]
331pub struct MemWalIndexDetails {
332    pub snapshot_ts_millis: i64,
333    pub num_shards: u32,
334    pub inline_snapshots: Option<Vec<u8>>,
335    pub sharding_specs: Vec<ShardingSpec>,
336    pub maintained_indexes: Vec<String>,
337    pub compacted_sstables: Vec<CompactedSsTable>,
338    pub index_catchup: Vec<IndexCatchupProgress>,
339    /// Default `ShardWriter` configuration values for this MemWAL index.
340    ///
341    /// Persisted so every writer — across processes and restarts — starts
342    /// from the same default writer configuration. These are defaults only;
343    /// an individual writer may still override any value at runtime in its
344    /// own (non-persisted) `ShardWriterConfig`.
345    pub writer_config_defaults: HashMap<String, String>,
346}
347
348impl From<&MemWalIndexDetails> for pb::MemWalIndexDetails {
349    fn from(details: &MemWalIndexDetails) -> Self {
350        Self {
351            snapshot_ts_millis: details.snapshot_ts_millis,
352            num_shards: details.num_shards,
353            inline_snapshots: details.inline_snapshots.clone(),
354            sharding_specs: details.sharding_specs.iter().map(|rs| rs.into()).collect(),
355            maintained_indexes: details.maintained_indexes.clone(),
356            compacted_sstables: details
357                .compacted_sstables
358                .iter()
359                .map(|sstable| sstable.into())
360                .collect(),
361            index_catchup: details.index_catchup.iter().map(|icp| icp.into()).collect(),
362            writer_config_defaults: details.writer_config_defaults.clone(),
363        }
364    }
365}
366
367impl TryFrom<pb::MemWalIndexDetails> for MemWalIndexDetails {
368    type Error = Error;
369
370    fn try_from(details: pb::MemWalIndexDetails) -> lance_core::Result<Self> {
371        Ok(Self {
372            snapshot_ts_millis: details.snapshot_ts_millis,
373            num_shards: details.num_shards,
374            inline_snapshots: details.inline_snapshots,
375            sharding_specs: details
376                .sharding_specs
377                .into_iter()
378                .map(ShardingSpec::from)
379                .collect(),
380            maintained_indexes: details.maintained_indexes,
381            compacted_sstables: details
382                .compacted_sstables
383                .into_iter()
384                .map(CompactedSsTable::try_from)
385                .collect::<lance_core::Result<_>>()?,
386            index_catchup: details
387                .index_catchup
388                .into_iter()
389                .map(IndexCatchupProgress::try_from)
390                .collect::<lance_core::Result<_>>()?,
391            writer_config_defaults: details.writer_config_defaults,
392        })
393    }
394}
395
396/// MemWAL Index provides access to MemWAL configuration and state.
397#[derive(Debug, Clone, PartialEq, Eq, DeepSizeOf)]
398pub struct MemWalIndex {
399    pub details: MemWalIndexDetails,
400}
401
402impl MemWalIndex {
403    pub fn new(details: MemWalIndexDetails) -> Self {
404        Self { details }
405    }
406
407    pub fn compacted_generation_for_shard(&self, shard_id: &Uuid) -> Option<u64> {
408        self.details
409            .compacted_sstables
410            .iter()
411            .find(|sstable| &sstable.shard_id == shard_id)
412            .map(|sstable| sstable.generation)
413    }
414
415    /// Get the caught up generation for a specific index and shard.
416    /// Returns None if the index is not tracked (assumed fully caught up).
417    pub fn index_caught_up_generation(&self, index_name: &str, shard_id: &Uuid) -> Option<u64> {
418        self.details
419            .index_catchup
420            .iter()
421            .find(|icp| icp.index_name == index_name)
422            .and_then(|icp| icp.caught_up_generation_for_shard(shard_id))
423    }
424
425    /// Check if an index is fully caught up for a shard.
426    /// Returns true if the index covers all compacted data for the shard.
427    pub fn is_index_caught_up(&self, index_name: &str, shard_id: &Uuid) -> bool {
428        let compacted_gen = self.compacted_generation_for_shard(shard_id).unwrap_or(0);
429        let caught_up_gen = self.index_caught_up_generation(index_name, shard_id);
430
431        // If not tracked in index_catchup, assumed fully caught up
432        caught_up_gen.is_none_or(|generation| generation >= compacted_gen)
433    }
434}
lance_table/system_index/mem_wal.rs

lance_table/system_index/
mem_wal.rs
