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

lance_table/system_index/
mem_wal.rs
