lance_table/system_index/
mem_wal.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! MemWAL index data structures and metadata helpers.
5//!
6//! The MemWAL Index stores:
7//! - Configuration (sharding_specs, maintained_indexes)
8//! - SSTable compaction progress
9//! - Shard state snapshots (eventually consistent)
10//!
11//! Writers no longer update the index on every write. Instead, they update
12//! shard manifests directly. This module provides functions to:
13//! - Load the MemWAL index
14//! - Update compacted SSTables (called during merge-insert commits)
15
16use std::collections::{HashMap, HashSet};
17use std::sync::Arc;
18
19use lance_core::deepsize::DeepSizeOf;
20use lance_core::{Error, Result};
21use serde::{Deserialize, Serialize};
22use uuid::Uuid;
23
24use crate::format::{IndexMetadata, pb};
25
26pub const MEM_WAL_INDEX_NAME: &str = "__lance_mem_wal";
27
28/// Type alias for shard identifier (UUID v4).
29pub type ShardId = Uuid;
30
31/// An SSTable: the immutable result of flushing a MemTable, stored as a Lance dataset.
32#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, DeepSizeOf)]
33pub struct SsTable {
34    pub generation: u64,
35    pub path: String,
36}
37
38impl From<&SsTable> for pb::SsTable {
39    fn from(sstable: &SsTable) -> Self {
40        Self {
41            generation: sstable.generation,
42            path: sstable.path.clone(),
43        }
44    }
45}
46
47impl From<pb::SsTable> for SsTable {
48    fn from(sstable: pb::SsTable) -> Self {
49        Self {
50            generation: sstable.generation,
51            path: sstable.path,
52        }
53    }
54}
55
56/// A pointer to the latest SSTable compacted for a shard.
57#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash, Serialize, Deserialize)]
58pub struct CompactedSsTable {
59    pub shard_id: Uuid,
60    pub generation: u64,
61}
62
63impl DeepSizeOf for CompactedSsTable {
64    fn deep_size_of_children(&self, _context: &mut lance_core::deepsize::Context) -> usize {
65        0 // UUID is 16 bytes fixed size, no heap allocations
66    }
67}
68
69impl CompactedSsTable {
70    pub fn new(shard_id: Uuid, generation: u64) -> Self {
71        Self {
72            shard_id,
73            generation,
74        }
75    }
76}
77
78impl From<&CompactedSsTable> for pb::CompactedSsTable {
79    fn from(sstable: &CompactedSsTable) -> Self {
80        Self {
81            shard_id: Some((&sstable.shard_id).into()),
82            generation: sstable.generation,
83        }
84    }
85}
86
87impl TryFrom<pb::CompactedSsTable> for CompactedSsTable {
88    type Error = Error;
89
90    fn try_from(sstable: pb::CompactedSsTable) -> lance_core::Result<Self> {
91        let shard_id = sstable
92            .shard_id
93            .as_ref()
94            .map(Uuid::try_from)
95            .ok_or_else(|| Error::invalid_input("Missing shard_id in CompactedSsTable"))??;
96        Ok(Self {
97            shard_id,
98            generation: sstable.generation,
99        })
100    }
101}
102
103/// Tracks which compacted SSTable generation a base table index covers.
104/// Used to determine whether to read from SSTable indexes or base table.
105#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize, DeepSizeOf)]
106pub struct IndexCatchupProgress {
107    pub index_name: String,
108    pub caught_up_generations: Vec<CompactedSsTable>,
109}
110
111impl IndexCatchupProgress {
112    pub fn new(index_name: String, caught_up_generations: Vec<CompactedSsTable>) -> Self {
113        Self {
114            index_name,
115            caught_up_generations,
116        }
117    }
118
119    /// Get the caught up generation for a specific shard.
120    /// Returns None if the shard is not present (assumed fully caught up).
121    pub fn caught_up_generation_for_shard(&self, shard_id: &Uuid) -> Option<u64> {
122        self.caught_up_generations
123            .iter()
124            .find(|sstable| &sstable.shard_id == shard_id)
125            .map(|sstable| sstable.generation)
126    }
127}
128
129impl From<&IndexCatchupProgress> for pb::IndexCatchupProgress {
130    fn from(icp: &IndexCatchupProgress) -> Self {
131        Self {
132            index_name: icp.index_name.clone(),
133            caught_up_generations: icp
134                .caught_up_generations
135                .iter()
136                .map(|sstable| sstable.into())
137                .collect(),
138        }
139    }
140}
141
142impl TryFrom<pb::IndexCatchupProgress> for IndexCatchupProgress {
143    type Error = Error;
144
145    fn try_from(icp: pb::IndexCatchupProgress) -> lance_core::Result<Self> {
146        Ok(Self {
147            index_name: icp.index_name,
148            caught_up_generations: icp
149                .caught_up_generations
150                .into_iter()
151                .map(CompactedSsTable::try_from)
152                .collect::<lance_core::Result<_>>()?,
153        })
154    }
155}
156
157/// Lifecycle status of a WAL shard, persisted in [`ShardManifest`].
158///
159/// `Sealed` is the durable in-doubt record for drop-table two-phase
160/// commit: a sealed shard refuses new writer claims (enforced in
161/// `claim_epoch`) but is reversible back to `Active` on rollback.
162#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
163pub enum ShardStatus {
164    /// Normal: the shard accepts writer claims.
165    #[default]
166    Active,
167    /// A drop is in flight: claims are refused. Reversible.
168    Sealed,
169}
170
171impl ShardStatus {
172    /// Map to the protobuf enum discriminant (`pb::ShardStatus`).
173    fn to_i32(self) -> i32 {
174        match self {
175            Self::Active => 0,
176            Self::Sealed => 1,
177        }
178    }
179
180    /// Map from the protobuf enum discriminant; unknown values decode as
181    /// `Active` (forward-compatible default).
182    fn from_i32(v: i32) -> Self {
183        match v {
184            1 => Self::Sealed,
185            _ => Self::Active,
186        }
187    }
188}
189
190/// Shard manifest containing epoch-based fencing and WAL state.
191/// Each shard has exactly one active writer at any time.
192#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
193pub struct ShardManifest {
194    pub shard_id: Uuid,
195    pub version: u64,
196    pub shard_spec_id: u32,
197    /// Computed shard field values as raw Arrow scalar bytes, keyed by field id.
198    /// The byte encoding follows Arrow's little-endian convention: int32 is 4 LE
199    /// bytes, utf8 is raw UTF-8 bytes, etc. The result_type in the corresponding
200    /// ShardingField from the ShardingSpec determines how to interpret each value.
201    pub shard_field_values: HashMap<String, Vec<u8>>,
202    pub writer_epoch: u64,
203    /// The most recent WAL entry position flushed to a MemTable.
204    /// Recovery replays from `replay_after_wal_entry_position + 1`. The
205    /// default value 0 means "no flush has ever stamped this shard" — WAL
206    /// positions themselves are 1-based, so 0 is never a valid covered
207    /// position.
208    pub replay_after_wal_entry_position: u64,
209    /// The most recent WAL entry position observed at manifest write time.
210    /// Default 0 means "no entry has been written yet"; WAL positions are
211    /// 1-based.
212    pub wal_entry_position_last_seen: u64,
213    pub current_generation: u64,
214    pub sstables: Vec<SsTable>,
215    /// Lifecycle status (drop-table 2PC). Defaults to `Active`; preserved
216    /// across claims via `..base` so only fresh constructions set it.
217    pub status: ShardStatus,
218}
219
220impl DeepSizeOf for ShardManifest {
221    fn deep_size_of_children(&self, context: &mut lance_core::deepsize::Context) -> usize {
222        self.shard_field_values.deep_size_of_children(context)
223            + self.sstables.deep_size_of_children(context)
224    }
225}
226
227impl From<&ShardManifest> for pb::ShardManifest {
228    fn from(rm: &ShardManifest) -> Self {
229        Self {
230            shard_id: Some((&rm.shard_id).into()),
231            version: rm.version,
232            shard_spec_id: rm.shard_spec_id,
233            shard_field_entries: rm
234                .shard_field_values
235                .iter()
236                .map(|(k, v)| pb::ShardFieldEntry {
237                    field_id: k.clone(),
238                    value: v.clone(),
239                })
240                .collect(),
241            writer_epoch: rm.writer_epoch,
242            replay_after_wal_entry_position: rm.replay_after_wal_entry_position,
243            wal_entry_position_last_seen: rm.wal_entry_position_last_seen,
244            current_generation: rm.current_generation,
245            sstables: rm.sstables.iter().map(|sstable| sstable.into()).collect(),
246            status: rm.status.to_i32(),
247        }
248    }
249}
250
251impl TryFrom<pb::ShardManifest> for ShardManifest {
252    type Error = Error;
253
254    fn try_from(rm: pb::ShardManifest) -> lance_core::Result<Self> {
255        let shard_id = rm
256            .shard_id
257            .as_ref()
258            .map(Uuid::try_from)
259            .ok_or_else(|| Error::invalid_input("Missing shard_id in ShardManifest"))??;
260        let shard_field_values = rm
261            .shard_field_entries
262            .into_iter()
263            .map(|e| (e.field_id, e.value))
264            .collect();
265        Ok(Self {
266            shard_id,
267            version: rm.version,
268            shard_spec_id: rm.shard_spec_id,
269            shard_field_values,
270            writer_epoch: rm.writer_epoch,
271            replay_after_wal_entry_position: rm.replay_after_wal_entry_position,
272            wal_entry_position_last_seen: rm.wal_entry_position_last_seen,
273            current_generation: rm.current_generation,
274            sstables: rm.sstables.into_iter().map(SsTable::from).collect(),
275            status: ShardStatus::from_i32(rm.status),
276        })
277    }
278}
279
280/// Sharding field definition.
281#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, DeepSizeOf)]
282pub struct ShardingField {
283    pub field_id: String,
284    pub source_ids: Vec<i32>,
285    pub transform: Option<String>,
286    pub expression: Option<String>,
287    pub result_type: String,
288    pub parameters: HashMap<String, String>,
289}
290
291impl From<&ShardingField> for pb::ShardingField {
292    fn from(rf: &ShardingField) -> Self {
293        Self {
294            field_id: rf.field_id.clone(),
295            source_ids: rf.source_ids.clone(),
296            transform: rf.transform.clone(),
297            expression: rf.expression.clone(),
298            result_type: rf.result_type.clone(),
299            parameters: rf.parameters.clone(),
300        }
301    }
302}
303
304impl From<pb::ShardingField> for ShardingField {
305    fn from(rf: pb::ShardingField) -> Self {
306        Self {
307            field_id: rf.field_id,
308            source_ids: rf.source_ids,
309            transform: rf.transform,
310            expression: rf.expression,
311            result_type: rf.result_type,
312            parameters: rf.parameters,
313        }
314    }
315}
316
317/// Sharding spec definition.
318#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, DeepSizeOf)]
319pub struct ShardingSpec {
320    pub spec_id: u32,
321    pub fields: Vec<ShardingField>,
322}
323
324impl From<&ShardingSpec> for pb::ShardingSpec {
325    fn from(rs: &ShardingSpec) -> Self {
326        Self {
327            spec_id: rs.spec_id,
328            fields: rs.fields.iter().map(|f| f.into()).collect(),
329        }
330    }
331}
332
333impl From<pb::ShardingSpec> for ShardingSpec {
334    fn from(rs: pb::ShardingSpec) -> Self {
335        Self {
336            spec_id: rs.spec_id,
337            fields: rs.fields.into_iter().map(ShardingField::from).collect(),
338        }
339    }
340}
341
342/// Index details for MemWAL Index, stored in IndexMetadata.index_details.
343#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize, DeepSizeOf)]
344pub struct MemWalIndexDetails {
345    pub snapshot_ts_millis: i64,
346    pub num_shards: u32,
347    pub inline_snapshots: Option<Vec<u8>>,
348    pub sharding_specs: Vec<ShardingSpec>,
349    pub maintained_indexes: Vec<String>,
350    pub compacted_sstables: Vec<CompactedSsTable>,
351    pub index_catchup: Vec<IndexCatchupProgress>,
352    /// Default `ShardWriter` configuration values for this MemWAL index.
353    ///
354    /// Persisted so every writer — across processes and restarts — starts
355    /// from the same default writer configuration. These are defaults only;
356    /// an individual writer may still override any value at runtime in its
357    /// own (non-persisted) `ShardWriterConfig`.
358    pub writer_config_defaults: HashMap<String, String>,
359}
360
361impl From<&MemWalIndexDetails> for pb::MemWalIndexDetails {
362    fn from(details: &MemWalIndexDetails) -> Self {
363        Self {
364            snapshot_ts_millis: details.snapshot_ts_millis,
365            num_shards: details.num_shards,
366            inline_snapshots: details.inline_snapshots.clone(),
367            sharding_specs: details.sharding_specs.iter().map(|rs| rs.into()).collect(),
368            maintained_indexes: details.maintained_indexes.clone(),
369            compacted_sstables: details
370                .compacted_sstables
371                .iter()
372                .map(|sstable| sstable.into())
373                .collect(),
374            index_catchup: details.index_catchup.iter().map(|icp| icp.into()).collect(),
375            writer_config_defaults: details.writer_config_defaults.clone(),
376        }
377    }
378}
379
380impl TryFrom<pb::MemWalIndexDetails> for MemWalIndexDetails {
381    type Error = Error;
382
383    fn try_from(details: pb::MemWalIndexDetails) -> lance_core::Result<Self> {
384        Ok(Self {
385            snapshot_ts_millis: details.snapshot_ts_millis,
386            num_shards: details.num_shards,
387            inline_snapshots: details.inline_snapshots,
388            sharding_specs: details
389                .sharding_specs
390                .into_iter()
391                .map(ShardingSpec::from)
392                .collect(),
393            maintained_indexes: details.maintained_indexes,
394            compacted_sstables: details
395                .compacted_sstables
396                .into_iter()
397                .map(CompactedSsTable::try_from)
398                .collect::<lance_core::Result<_>>()?,
399            index_catchup: details
400                .index_catchup
401                .into_iter()
402                .map(IndexCatchupProgress::try_from)
403                .collect::<lance_core::Result<_>>()?,
404            writer_config_defaults: details.writer_config_defaults,
405        })
406    }
407}
408
409/// MemWAL Index provides access to MemWAL configuration and state.
410#[derive(Debug, Clone, PartialEq, Eq, DeepSizeOf)]
411pub struct MemWalIndex {
412    pub details: MemWalIndexDetails,
413}
414
415impl MemWalIndex {
416    pub fn new(details: MemWalIndexDetails) -> Self {
417        Self { details }
418    }
419
420    pub fn compacted_generation_for_shard(&self, shard_id: &Uuid) -> Option<u64> {
421        self.details
422            .compacted_sstables
423            .iter()
424            .find(|sstable| &sstable.shard_id == shard_id)
425            .map(|sstable| sstable.generation)
426    }
427
428    /// Get the caught up generation for a specific index and shard.
429    /// Returns None if the index is not tracked (assumed fully caught up).
430    pub fn index_caught_up_generation(&self, index_name: &str, shard_id: &Uuid) -> Option<u64> {
431        self.details
432            .index_catchup
433            .iter()
434            .find(|icp| icp.index_name == index_name)
435            .and_then(|icp| icp.caught_up_generation_for_shard(shard_id))
436    }
437}
438
439// Reading and updating the `IndexMetadata` entry that carries the details above.
440
441/// Load MemWalIndexDetails from an IndexMetadata.
442pub fn load_mem_wal_index_details(index: IndexMetadata) -> Result<MemWalIndexDetails> {
443    if let Some(details_any) = index.index_details.as_ref() {
444        if !details_any.type_url.ends_with("MemWalIndexDetails") {
445            return Err(Error::index(format!(
446                "Index details is not for the MemWAL index, but {}",
447                details_any.type_url
448            )));
449        }
450
451        Ok(MemWalIndexDetails::try_from(
452            details_any.to_msg::<pb::MemWalIndexDetails>()?,
453        )?)
454    } else {
455        Err(Error::index("Index details not found for the MemWAL index"))
456    }
457}
458
459/// Open the MemWAL index from its metadata.
460pub fn open_mem_wal_index(index: IndexMetadata) -> Result<Arc<MemWalIndex>> {
461    Ok(Arc::new(MemWalIndex::new(load_mem_wal_index_details(
462        index,
463    )?)))
464}
465
466/// Update `compacted_sstables` in the MemWAL index.
467///
468/// Called from the final data-changing merge-insert commit for a compaction
469/// target, so the rows and the generation that describes them publish
470/// together.
471///
472/// A proposed generation must be **strictly greater** than the one the latest
473/// state records for that shard, and a stale one fails the whole transaction.
474/// Accepting it while keeping the larger marker would publish that worker's row
475/// mutations under a generation it did not produce, and anything reading only
476/// the marker could then stop serving SSTables whose rows were never inserted.
477///
478/// Every other `MemWalIndexDetails` field is carried through untouched.
479pub fn update_mem_wal_index_compacted_sstables(
480    indices: &mut [IndexMetadata],
481    dataset_version: u64,
482    new_compacted_sstables: Vec<CompactedSsTable>,
483) -> Result<()> {
484    if new_compacted_sstables.is_empty() {
485        return Ok(());
486    }
487
488    let mut seen_shards = HashSet::with_capacity(new_compacted_sstables.len());
489    for sstable in &new_compacted_sstables {
490        if !seen_shards.insert(sstable.shard_id) {
491            return Err(Error::invalid_input(format!(
492                "Duplicate shard {} in one SSTable compaction update; each shard \
493                 may advance at most once per transaction",
494                sstable.shard_id
495            )));
496        }
497    }
498
499    // Default details would describe a table with no MemWAL shards at all, so
500    // the recorded generation would name a shard nothing can corroborate.
501    // Refuse instead of inventing metadata.
502    let pos = indices
503        .iter()
504        .position(|idx| idx.name == MEM_WAL_INDEX_NAME)
505        .ok_or_else(|| {
506            Error::invalid_input(format!(
507                "Cannot record SSTable compaction progress: the {} system index \
508                 does not exist on this table",
509                MEM_WAL_INDEX_NAME
510            ))
511        })?;
512
513    // Validated against a copy so a rejected update leaves `indices` exactly as
514    // the caller passed it.
515    let mut details = load_mem_wal_index_details(indices[pos].clone())?;
516
517    for new_sstable in new_compacted_sstables {
518        match details
519            .compacted_sstables
520            .iter_mut()
521            .find(|sstable| sstable.shard_id == new_sstable.shard_id)
522        {
523            Some(existing) if new_sstable.generation <= existing.generation => {
524                return Err(Error::invalid_input(format!(
525                    "Stale SSTable compaction for shard {}: proposed generation {} \
526                     is not greater than the recorded generation {}",
527                    new_sstable.shard_id, new_sstable.generation, existing.generation
528                )));
529            }
530            Some(existing) => existing.generation = new_sstable.generation,
531            None => details.compacted_sstables.push(new_sstable),
532        }
533    }
534
535    // Replaced in place so the index list keeps its order.
536    indices[pos] = new_mem_wal_index_meta(dataset_version, details)?;
537    Ok(())
538}
539
540/// Create a new MemWAL index metadata entry.
541///
542/// A fresh UUID is minted on every rewrite, including metadata-only updates.
543/// The decoded-details cache is keyed on that UUID, so the change of identity
544/// is what invalidates it; holding the UUID steady would leave a warmed reader
545/// answering with the state from before the update.
546pub fn new_mem_wal_index_meta(
547    dataset_version: u64,
548    details: MemWalIndexDetails,
549) -> Result<IndexMetadata> {
550    Ok(IndexMetadata {
551        uuid: Uuid::new_v4(),
552        name: MEM_WAL_INDEX_NAME.to_string(),
553        fields: vec![],
554        covering_fields: vec![],
555        dataset_version,
556        fragment_bitmap: None,
557        index_details: Some(Arc::new(prost_types::Any::from_msg(
558            &pb::MemWalIndexDetails::from(&details),
559        )?)),
560        index_version: 0,
561        created_at: Some(chrono::Utc::now()),
562        base_id: None,
563        // Memory WAL index is inline (no files)
564        files: None,
565    })
566}
lance_table/system_index/mem_wal.rs

lance_table/system_index/
mem_wal.rs
