Skip to main content

lance_table/rowids/
version.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Row version tracking for cross-version diff functionality
5//!
6//! This module provides data structures and functionality to track the latest
7//! update version for each row in a Lance dataset, enabling efficient
8//! cross-version diff operations.
9
10use std::collections::HashMap;
11use std::{ops::Range, sync::Arc};
12
13use lance_core::Error;
14use lance_core::Result;
15use lance_core::deepsize::DeepSizeOf;
16use prost::Message;
17use serde::de::Deserializer;
18use serde::ser::Serializer;
19use serde::{Deserialize, Serialize};
20
21use crate::format::{Fragment, pb};
22use crate::rowids::segment::U64Segment;
23use crate::rowids::{RowIdSequence, read_row_ids};
24
25/// One fragment's row lineage sequences that live outside the manifest, read
26/// ahead of a commit by the caller.
27///
28/// Building a manifest is synchronous and has no object store, so it cannot
29/// read a sequence spilled to a data file column (`RowIdMeta::Column`,
30/// `RowDatasetVersionMeta::Column`). A caller that can do IO loads them first
31/// and hands them over in [`ManifestBuildConfig`](crate::format::ManifestBuildConfig).
32/// Only the spilled sequences need to be present; inline ones are decoded on
33/// the spot.
34#[derive(Debug, Clone, Default)]
35pub struct LoadedRowLineage {
36    pub row_ids: Option<Arc<RowIdSequence>>,
37    pub created_at: Option<Arc<RowDatasetVersionSequence>>,
38    pub last_updated_at: Option<Arc<RowDatasetVersionSequence>>,
39}
40
41/// [`LoadedRowLineage`] per fragment id.
42pub type SpilledRowLineage = HashMap<u64, LoadedRowLineage>;
43
44/// A run of identical versions over a contiguous span of row positions.
45///
46/// Span is expressed as a U64Segment over row offsets (0..N within a fragment),
47/// not over row IDs. This keeps the encoding aligned with RowIdSequence order
48/// and enables zipped iteration without building a map.
49#[derive(Debug, Clone, PartialEq, Eq, DeepSizeOf)]
50pub struct RowDatasetVersionRun {
51    pub span: U64Segment,
52    pub version: u64,
53}
54
55impl RowDatasetVersionRun {
56    /// Number of rows covered by this run.
57    pub fn len(&self) -> usize {
58        self.span.len()
59    }
60
61    /// Whether this run covers no rows.
62    pub fn is_empty(&self) -> bool {
63        self.span.is_empty()
64    }
65
66    /// The version value of this run.
67    pub fn version(&self) -> u64 {
68        self.version
69    }
70}
71
72/// Sequence of dataset versions
73///
74/// Stores version runs aligned to the positional order of RowIdSequence.
75/// Provides sequential iterators and optional lightweight indexing for
76/// efficient random access.
77#[derive(Debug, Clone, PartialEq, Eq, DeepSizeOf, Default)]
78pub struct RowDatasetVersionSequence {
79    pub runs: Vec<RowDatasetVersionRun>,
80}
81
82/// A reusable cursor for reading ranges from a version sequence in one pass.
83///
84/// The cursor caches the current run length. Readers normally request adjacent
85/// batches, so this avoids rebuilding all run offsets and rescanning the run
86/// prefix for every batch. A backwards selection lazily builds a run-offset
87/// index. Once built, non-adjacent selections use it in either direction.
88#[derive(Debug, Default)]
89pub(crate) struct RowDatasetVersionCursor {
90    run_index: usize,
91    offset_in_run: usize,
92    position: usize,
93    run_len: Option<usize>,
94    run_offsets: Option<Vec<usize>>,
95    indexed_total_len: usize,
96    #[cfg(test)]
97    indexed_seek_count: usize,
98}
99
100impl RowDatasetVersionCursor {
101    fn seek_indexed(
102        &mut self,
103        sequence: &RowDatasetVersionSequence,
104        position: usize,
105    ) -> Result<()> {
106        if self.run_offsets.is_none() {
107            let mut total_len = 0;
108            let run_offsets = sequence
109                .runs
110                .iter()
111                .map(|run| {
112                    let offset = total_len;
113                    total_len += run.len();
114                    offset
115                })
116                .collect();
117            self.run_offsets = Some(run_offsets);
118            self.indexed_total_len = total_len;
119        }
120
121        if position >= self.indexed_total_len {
122            return Err(Error::internal(format!(
123                "version column position {} out of range (total_len={})",
124                position, self.indexed_total_len
125            )));
126        }
127
128        let run_offsets = self.run_offsets.as_ref().unwrap();
129        let mut run_index = match run_offsets.binary_search(&position) {
130            Ok(run_index) => run_index,
131            Err(run_index) => run_index - 1,
132        };
133        while run_index + 1 < run_offsets.len() && run_offsets[run_index + 1] <= position {
134            run_index += 1;
135        }
136        self.run_index = run_index;
137        self.offset_in_run = position - run_offsets[run_index];
138        self.position = position;
139        self.run_len = None;
140        #[cfg(test)]
141        {
142            self.indexed_seek_count += 1;
143        }
144        Ok(())
145    }
146
147    fn current_run<'a>(
148        &mut self,
149        sequence: &'a RowDatasetVersionSequence,
150    ) -> Option<(&'a RowDatasetVersionRun, usize)> {
151        loop {
152            let run = sequence.runs.get(self.run_index)?;
153            let run_len = *self.run_len.get_or_insert_with(|| match &run.span {
154                // Version runs are normally positional ranges. Keep this hot
155                // path local instead of using the general segment length path.
156                U64Segment::Range(range) => (range.end - range.start) as usize,
157                span => span.len(),
158            });
159            if self.offset_in_run < run_len {
160                return Some((run, run_len));
161            }
162            self.run_index += 1;
163            self.offset_in_run = 0;
164            self.run_len = None;
165        }
166    }
167
168    /// Append the versions in `selection` to `versions`.
169    pub(crate) fn extend_range(
170        &mut self,
171        sequence: &RowDatasetVersionSequence,
172        selection: Range<usize>,
173        versions: &mut Vec<u64>,
174    ) -> Result<()> {
175        if selection.is_empty() {
176            return Ok(());
177        }
178        if selection.start < self.position
179            || (self.run_offsets.is_some() && selection.start != self.position)
180        {
181            self.seek_indexed(sequence, selection.start)?;
182        }
183
184        while self.position < selection.start {
185            let Some((_, run_len)) = self.current_run(sequence) else {
186                return Err(Error::internal(format!(
187                    "version column position {} out of range (total_len={})",
188                    selection.start, self.position
189                )));
190            };
191            let advance = (selection.start - self.position).min(run_len - self.offset_in_run);
192            self.offset_in_run += advance;
193            self.position += advance;
194        }
195
196        while self.position < selection.end {
197            let Some((run, run_len)) = self.current_run(sequence) else {
198                return Err(Error::internal(format!(
199                    "version column position {} out of range (total_len={})",
200                    self.position, self.position
201                )));
202            };
203            let count = (selection.end - self.position).min(run_len - self.offset_in_run);
204            versions.extend(std::iter::repeat_n(run.version(), count));
205            self.offset_in_run += count;
206            self.position += count;
207        }
208        Ok(())
209    }
210}
211
212impl RowDatasetVersionSequence {
213    /// Create a new empty version sequence
214    pub fn new() -> Self {
215        Self { runs: Vec::new() }
216    }
217
218    /// Create a version sequence with a single uniform run of `row_count` rows.
219    pub fn from_uniform_row_count(row_count: u64, version: u64) -> Self {
220        if row_count == 0 {
221            return Self::new();
222        }
223        let run = RowDatasetVersionRun {
224            span: U64Segment::Range(0..row_count),
225            version,
226        };
227        Self { runs: vec![run] }
228    }
229
230    /// Run-length encode one version per row, in row offset order.
231    pub fn from_versions(versions: &[u64]) -> Self {
232        let mut runs = Vec::new();
233        let mut run_start = 0u64;
234        for (i, window) in versions.windows(2).enumerate() {
235            if window[0] != window[1] {
236                runs.push(RowDatasetVersionRun {
237                    span: U64Segment::Range(run_start..i as u64 + 1),
238                    version: window[0],
239                });
240                run_start = i as u64 + 1;
241            }
242        }
243        if let Some(last) = versions.last() {
244            runs.push(RowDatasetVersionRun {
245                span: U64Segment::Range(run_start..versions.len() as u64),
246                version: *last,
247            });
248        }
249        Self { runs }
250    }
251
252    /// Number of rows tracked by this sequence (sum of run lengths).
253    pub fn len(&self) -> u64 {
254        self.runs.iter().map(|s| s.len() as u64).sum()
255    }
256
257    /// Empty if there are no runs or all runs are empty.
258    pub fn is_empty(&self) -> bool {
259        self.runs.is_empty() || self.runs.iter().all(|s| s.is_empty())
260    }
261
262    /// Returns a forward iterator over versions, expanding runs lazily.
263    pub fn versions(&self) -> VersionsIter<'_> {
264        VersionsIter::new(&self.runs)
265    }
266
267    /// Create a reusable cursor for sequential range reads.
268    pub(crate) fn cursor(&self) -> RowDatasetVersionCursor {
269        RowDatasetVersionCursor::default()
270    }
271
272    /// Random access: get the version at global row position `index`.
273    pub fn version_at(&self, index: usize) -> Option<u64> {
274        let mut offset = 0usize;
275        for run in &self.runs {
276            let len = run.len();
277            if index < offset + len {
278                return Some(run.version());
279            }
280            offset += len;
281        }
282        None
283    }
284
285    /// Get the version associated with a specific row id.
286    /// This reconstructs the positional offset from RowIdSequence and then
287    /// performs `version_at` lookup.
288    pub fn get_version_for_row_id(&self, row_ids: &RowIdSequence, row_id: u64) -> Option<u64> {
289        let mut offset = 0usize;
290        for seg in &row_ids.0 {
291            if seg.range().is_some_and(|r| r.contains(&row_id))
292                && let Some(local) = seg.position(row_id)
293            {
294                return self.version_at(offset + local);
295            }
296            offset += seg.len();
297        }
298        None
299    }
300
301    /// Convenience: collect row IDs with version strictly greater than `threshold`.
302    pub fn rows_with_version_greater_than(
303        &self,
304        row_ids: &RowIdSequence,
305        threshold: u64,
306    ) -> Vec<u64> {
307        row_ids
308            .iter()
309            .zip(self.versions())
310            .filter_map(|(rid, v)| if v > threshold { Some(rid) } else { None })
311            .collect()
312    }
313
314    /// Delete rows by positional offsets (e.g., from a deletion vector)
315    pub fn mask(&mut self, positions: impl IntoIterator<Item = u32>) -> Result<()> {
316        let mut local_positions: Vec<u32> = Vec::new();
317        let mut positions_iter = positions.into_iter();
318        let mut curr_position = positions_iter.next();
319        let mut offset: usize = 0;
320        let mut cutoff: usize = 0;
321
322        for run in self.runs.iter_mut() {
323            cutoff += run.span.len();
324            while let Some(position) = curr_position {
325                if position as usize >= cutoff {
326                    break;
327                }
328                local_positions.push(position - offset as u32);
329                curr_position = positions_iter.next();
330            }
331
332            if !local_positions.is_empty() {
333                run.span.mask(local_positions.as_slice());
334                local_positions.clear();
335            }
336            offset = cutoff;
337        }
338
339        self.runs.retain(|r| !r.span.is_empty());
340        Ok(())
341    }
342}
343
344/// Iterator over versions expanding runs lazily.
345pub struct VersionsIter<'a> {
346    runs: &'a [RowDatasetVersionRun],
347    run_idx: usize,
348    remaining_in_run: usize,
349    current_version: u64,
350}
351
352impl<'a> VersionsIter<'a> {
353    fn new(runs: &'a [RowDatasetVersionRun]) -> Self {
354        let mut it = Self {
355            runs,
356            run_idx: 0,
357            remaining_in_run: 0,
358            current_version: 0,
359        };
360        it.advance_run();
361        it
362    }
363
364    fn advance_run(&mut self) {
365        if self.run_idx < self.runs.len() {
366            let run = &self.runs[self.run_idx];
367            self.remaining_in_run = run.len();
368            self.current_version = run.version();
369        } else {
370            self.remaining_in_run = 0;
371        }
372    }
373}
374
375impl<'a> Iterator for VersionsIter<'a> {
376    type Item = u64;
377
378    fn next(&mut self) -> Option<Self::Item> {
379        if self.remaining_in_run == 0 {
380            // Move to next run
381            self.run_idx += 1;
382            if self.run_idx >= self.runs.len() {
383                return None;
384            }
385            self.advance_run();
386        }
387        self.remaining_in_run = self.remaining_in_run.saturating_sub(1);
388        Some(self.current_version)
389    }
390}
391
392/// Metadata about the location of dataset version sequence data
393/// Following the same pattern as RowIdMeta
394///
395/// When stored inline, identical byte sequences are shared across fragments
396/// via `Arc<[u8]>` to reduce manifest memory for large tables.
397#[derive(Debug, Clone, PartialEq, Eq, DeepSizeOf)]
398pub enum RowDatasetVersionMeta {
399    /// Small sequences stored inline in the fragment metadata
400    Inline(Arc<[u8]>),
401    /// The sequence is spilled to a hidden column of one of the fragment's
402    /// data files, one version per physical row in offset order, at
403    /// [`ROW_CREATED_AT_VERSION_FIELD_ID`](crate::format::ROW_CREATED_AT_VERSION_FIELD_ID)
404    /// or
405    /// [`ROW_LAST_UPDATED_AT_VERSION_FIELD_ID`](crate::format::ROW_LAST_UPDATED_AT_VERSION_FIELD_ID)
406    /// depending on which sequence this is; see
407    /// [`Fragment::row_lineage_file`](crate::format::Fragment::row_lineage_file).
408    /// Reading it needs IO, so it goes through the dataset's loader rather
409    /// than [`Self::load_sequence`].
410    Column,
411}
412
413/// The JSON form of [`RowDatasetVersionMeta::Column`]: `{"column": {}}`, an
414/// empty object under the arm's name, mirroring the empty `RowLineageColumn`
415/// protobuf message. The name alone tells the arms apart; there is nothing to
416/// carry.
417#[derive(Serialize, Deserialize)]
418struct ColumnMarker {}
419
420// Custom Serialize: convert Arc<[u8]> to slice for transparent JSON output
421impl Serialize for RowDatasetVersionMeta {
422    fn serialize<S: Serializer>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error> {
423        #[derive(Serialize)]
424        #[serde(untagged)]
425        enum Helper<'a> {
426            Inline { inline: &'a [u8] },
427            Column { column: ColumnMarker },
428        }
429
430        match self {
431            Self::Inline(data) => Helper::Inline {
432                inline: data.as_ref(),
433            }
434            .serialize(serializer),
435            Self::Column => Helper::Column {
436                column: ColumnMarker {},
437            }
438            .serialize(serializer),
439        }
440    }
441}
442
443// Custom Deserialize: read Vec<u8> and convert to Arc<[u8]>
444impl<'de> Deserialize<'de> for RowDatasetVersionMeta {
445    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> std::result::Result<Self, D::Error> {
446        #[derive(Deserialize)]
447        #[serde(untagged)]
448        enum Helper {
449            Inline { inline: Vec<u8> },
450            Column { column: ColumnMarker },
451        }
452
453        match Helper::deserialize(deserializer)? {
454            Helper::Inline { inline } => Ok(Self::Inline(Arc::from(inline))),
455            Helper::Column { column } => {
456                let ColumnMarker {} = column;
457                Ok(Self::Column)
458            }
459        }
460    }
461}
462
463impl RowDatasetVersionMeta {
464    /// Create inline metadata from a version sequence
465    pub fn from_sequence(sequence: &RowDatasetVersionSequence) -> lance_core::Result<Self> {
466        let bytes = write_dataset_versions(sequence);
467        Ok(Self::Inline(Arc::from(bytes)))
468    }
469
470    /// Decode the version sequence stored inline in this metadata.
471    ///
472    /// A sequence stored outside the manifest needs IO to read, which this
473    /// synchronous accessor cannot do; it is an error here and is loaded
474    /// through the dataset instead.
475    pub fn load_sequence(&self) -> lance_core::Result<RowDatasetVersionSequence> {
476        match self {
477            Self::Inline(data) => read_dataset_versions(data),
478            Self::Column => Err(Error::not_supported(
479                "row version sequence spilled to a data file column cannot be decoded from \
480                 fragment metadata alone",
481            )),
482        }
483    }
484}
485
486/// Helper function to convert RowDatasetVersionMeta to protobuf format for last_updated_at
487pub fn last_updated_at_version_meta_to_pb(
488    meta: &Option<RowDatasetVersionMeta>,
489) -> Option<pb::data_fragment::LastUpdatedAtVersionSequence> {
490    meta.as_ref().map(|m| match m {
491        RowDatasetVersionMeta::Inline(data) => {
492            pb::data_fragment::LastUpdatedAtVersionSequence::InlineLastUpdatedAtVersions(
493                data.to_vec(),
494            )
495        }
496        RowDatasetVersionMeta::Column => {
497            pb::data_fragment::LastUpdatedAtVersionSequence::ColumnLastUpdatedAtVersions(
498                pb::RowLineageColumn {},
499            )
500        }
501    })
502}
503
504/// Helper function to convert RowDatasetVersionMeta to protobuf format for created_at
505pub fn created_at_version_meta_to_pb(
506    meta: &Option<RowDatasetVersionMeta>,
507) -> Option<pb::data_fragment::CreatedAtVersionSequence> {
508    meta.as_ref().map(|m| match m {
509        RowDatasetVersionMeta::Inline(data) => {
510            pb::data_fragment::CreatedAtVersionSequence::InlineCreatedAtVersions(data.to_vec())
511        }
512        RowDatasetVersionMeta::Column => {
513            pb::data_fragment::CreatedAtVersionSequence::ColumnCreatedAtVersions(
514                pb::RowLineageColumn {},
515            )
516        }
517    })
518}
519
520/// Serialize a dataset version sequence to a buffer (following RowIdSequence pattern)
521pub fn write_dataset_versions(sequence: &RowDatasetVersionSequence) -> Vec<u8> {
522    // Convert to protobuf sequence
523    let pb_sequence = pb::RowDatasetVersionSequence {
524        runs: sequence
525            .runs
526            .iter()
527            .map(|run| pb::RowDatasetVersionRun {
528                span: Some(pb::U64Segment::from(run.span.clone())),
529                version: run.version,
530            })
531            .collect(),
532    };
533
534    pb_sequence.encode_to_vec()
535}
536
537/// Deserialize a dataset version sequence from bytes (following RowIdSequence pattern)
538pub fn read_dataset_versions(data: &[u8]) -> lance_core::Result<RowDatasetVersionSequence> {
539    let pb_sequence = pb::RowDatasetVersionSequence::decode(data).map_err(|e| {
540        Error::internal(format!("Failed to decode RowDatasetVersionSequence: {}", e))
541    })?;
542
543    let segments = pb_sequence
544        .runs
545        .into_iter()
546        .map(|pb_run| {
547            let positions_pb = pb_run.span.ok_or_else(|| {
548                Error::internal("Missing positions in RowDatasetVersionRun".to_string())
549            })?;
550            let segment = U64Segment::try_from(positions_pb)?;
551            Ok(RowDatasetVersionRun {
552                span: segment,
553                version: pb_run.version,
554            })
555        })
556        .collect::<Result<Vec<_>>>()?;
557
558    Ok(RowDatasetVersionSequence { runs: segments })
559}
560
561/// Re-chunk a sequence of dataset version runs into new chunk sizes (aligned with RowIdSequence rechunking)
562pub fn rechunk_version_sequences(
563    sequences: impl IntoIterator<Item = RowDatasetVersionSequence>,
564    chunk_sizes: impl IntoIterator<Item = u64>,
565    allow_incomplete: bool,
566) -> Result<Vec<RowDatasetVersionSequence>> {
567    let chunk_sizes_vec: Vec<u64> = chunk_sizes.into_iter().collect();
568    let total_chunks = chunk_sizes_vec.len();
569    let mut chunked_sequences: Vec<RowDatasetVersionSequence> = Vec::with_capacity(total_chunks);
570
571    let mut run_iter = sequences
572        .into_iter()
573        .flat_map(|sequence| sequence.runs.into_iter())
574        .peekable();
575
576    let too_few_segments_error = |chunk_index: usize, expected_chunk_size: u64, remaining: u64| {
577        Error::invalid_input(format!(
578            "Got too few version runs for chunk {}. Expected chunk size: {}, remaining needed: {}",
579            chunk_index, expected_chunk_size, remaining
580        ))
581    };
582
583    let too_many_segments_error = |processed_chunks: usize, total_chunk_sizes: usize| {
584        Error::invalid_input(format!(
585            "Got too many version runs for the provided chunk lengths. Processed {} chunks out of {} expected",
586            processed_chunks, total_chunk_sizes
587        ))
588    };
589
590    let mut segment_offset = 0_u64;
591
592    for (chunk_index, chunk_size) in chunk_sizes_vec.iter().enumerate() {
593        let chunk_size = *chunk_size;
594        let mut out_seq = RowDatasetVersionSequence::new();
595        let mut remaining = chunk_size;
596
597        while remaining > 0 {
598            let remaining_in_segment = run_iter
599                .peek()
600                .map_or(0, |run| run.span.len() as u64 - segment_offset);
601
602            if remaining_in_segment == 0 {
603                if run_iter.next().is_some() {
604                    segment_offset = 0;
605                    continue;
606                } else if allow_incomplete {
607                    break;
608                } else {
609                    return Err(too_few_segments_error(chunk_index, chunk_size, remaining));
610                }
611            }
612
613            match remaining_in_segment.cmp(&remaining) {
614                std::cmp::Ordering::Greater => {
615                    let run = run_iter.peek().unwrap();
616                    let seg = run.span.slice(segment_offset as usize, remaining as usize);
617                    out_seq.runs.push(RowDatasetVersionRun {
618                        span: seg,
619                        version: run.version,
620                    });
621                    segment_offset += remaining;
622                    remaining = 0;
623                }
624                std::cmp::Ordering::Equal | std::cmp::Ordering::Less => {
625                    let run = run_iter.next().ok_or_else(|| {
626                        too_few_segments_error(chunk_index, chunk_size, remaining)
627                    })?;
628                    let seg = run
629                        .span
630                        .slice(segment_offset as usize, remaining_in_segment as usize);
631                    out_seq.runs.push(RowDatasetVersionRun {
632                        span: seg,
633                        version: run.version,
634                    });
635                    segment_offset = 0;
636                    remaining -= remaining_in_segment;
637                }
638            }
639        }
640
641        chunked_sequences.push(out_seq);
642    }
643
644    if run_iter.peek().is_some() {
645        return Err(too_many_segments_error(
646            chunked_sequences.len(),
647            total_chunks,
648        ));
649    }
650
651    Ok(chunked_sequences)
652}
653
654/// Build version metadata for a fragment if it has physical rows and no existing metadata.
655pub fn build_version_meta(
656    fragment: &Fragment,
657    current_version: u64,
658) -> Option<RowDatasetVersionMeta> {
659    if let Some(physical_rows) = fragment.physical_rows
660        && physical_rows > 0
661    {
662        // Verify row_id_meta exists (sanity check for stable row IDs)
663        if fragment.row_id_meta.is_none() {
664            panic!("Can not find row id meta, please make sure you have enabled stable row id.")
665        }
666
667        // Use physical_rows directly as the authoritative row count
668        // This is correct even for compacted fragments where row_id_meta might
669        // have been partially copied
670        let version_sequence = RowDatasetVersionSequence::from_uniform_row_count(
671            physical_rows as u64,
672            current_version,
673        );
674
675        return Some(RowDatasetVersionMeta::from_sequence(&version_sequence).unwrap());
676    }
677    None
678}
679
680/// Refresh row-level latest update version metadata for a full fragment rewrite-column update.
681///
682/// This sets a uniform version sequence for all rows in the fragment to `current_version`.
683pub fn refresh_row_latest_update_meta_for_full_frag_rewrite_cols(
684    fragment: &mut Fragment,
685    current_version: u64,
686) -> Result<()> {
687    let row_count = if let Some(pr) = fragment.physical_rows {
688        pr as u64
689    } else if let Some(row_id_meta) = fragment.row_id_meta.as_ref() {
690        match row_id_meta {
691            crate::format::RowIdMeta::Inline(data) => {
692                let sequence = read_row_ids(data).unwrap();
693                sequence.len()
694            }
695            // Follow existing behavior: a sequence that is not inline needs IO
696            // to read, which this synchronous path cannot do.
697            crate::format::RowIdMeta::Column => 0,
698        }
699    } else {
700        0
701    };
702
703    if row_count > 0 {
704        let version_seq =
705            RowDatasetVersionSequence::from_uniform_row_count(row_count, current_version);
706        let version_meta = RowDatasetVersionMeta::from_sequence(&version_seq)?;
707        fragment.last_updated_at_version_meta = Some(version_meta);
708    }
709
710    Ok(())
711}
712
713/// Refresh row-level latest update version metadata for a partial fragment rewrite-column update.
714///
715/// `updated_offsets` are local row offsets (within the fragment) that have been updated.
716/// Existing version metadata is preserved and only the updated positions are set to `current_version`.
717/// If no existing metadata is present, positions default to `prev_version`.
718///
719/// A fragment whose existing versions are spilled has them read from
720/// `spilled`; the refreshed sequence is placed inline.
721pub fn refresh_row_latest_update_meta_for_partial_frag_rewrite_cols(
722    fragment: &mut Fragment,
723    updated_offsets: &[usize],
724    current_version: u64,
725    prev_version: u64,
726    spilled: &SpilledRowLineage,
727) -> Result<()> {
728    // Determine row count for fragment
729    let row_count_u64: u64 = if let Some(pr) = fragment.physical_rows {
730        pr as u64
731    } else if let Some(row_id_meta) = fragment.row_id_meta.as_ref() {
732        match row_id_meta {
733            crate::format::RowIdMeta::Inline(data) => {
734                let sequence = read_row_ids(data).unwrap();
735                sequence.len()
736            }
737            // Reading these needs IO, which this synchronous path cannot do.
738            // Reachable only for a fragment that also has no `physical_rows`.
739            crate::format::RowIdMeta::Column => {
740                return Err(Error::not_supported(
741                    "refreshing row update versions for a fragment whose row id \
742                     sequence is stored outside the manifest",
743                ));
744            }
745        }
746    } else {
747        0
748    };
749
750    if row_count_u64 > 0 {
751        // Build base version vector from existing meta or previous dataset version
752        let mut base_versions: Vec<u64> = Vec::with_capacity(row_count_u64 as usize);
753        if let Some(meta) = fragment.last_updated_at_version_meta.as_ref() {
754            let base_seq = if matches!(meta, RowDatasetVersionMeta::Column) {
755                // The existing versions of the rows this update leaves alone
756                // live in a data file, which this path cannot read. Defaulting
757                // them would silently rewrite their lineage, so the caller has
758                // to have read them ahead of time.
759                let loaded = spilled
760                    .get(&fragment.id)
761                    .and_then(|lineage| lineage.last_updated_at.clone())
762                    .ok_or_else(|| {
763                        Error::not_supported(format!(
764                            "fragment {} stores its last-updated-at versions outside the \
765                             manifest and they were not loaded ahead of the commit",
766                            fragment.id
767                        ))
768                    })?;
769                Some(loaded.as_ref().clone())
770            } else {
771                meta.load_sequence().ok()
772            };
773            if let Some(base_seq) = base_seq {
774                base_versions.extend(base_seq.versions().take(row_count_u64 as usize));
775                base_versions.resize(row_count_u64 as usize, prev_version);
776            } else {
777                base_versions.resize(row_count_u64 as usize, prev_version);
778            }
779        } else {
780            base_versions.resize(row_count_u64 as usize, prev_version);
781        }
782
783        // Apply updates to updated positions
784        for &pos in updated_offsets {
785            if pos < base_versions.len() {
786                base_versions[pos] = current_version;
787            }
788        }
789
790        // Compress into runs
791        let mut runs: Vec<RowDatasetVersionRun> = Vec::new();
792        if !base_versions.is_empty() {
793            let mut start = 0usize;
794            let mut curr_ver = base_versions[0];
795            for (idx, &ver) in base_versions.iter().enumerate().skip(1) {
796                if ver != curr_ver {
797                    runs.push(RowDatasetVersionRun {
798                        span: U64Segment::Range(start as u64..idx as u64),
799                        version: curr_ver,
800                    });
801                    start = idx;
802                    curr_ver = ver;
803                }
804            }
805            runs.push(RowDatasetVersionRun {
806                span: U64Segment::Range(start as u64..base_versions.len() as u64),
807                version: curr_ver,
808            });
809        }
810        let new_seq = RowDatasetVersionSequence { runs };
811        let new_meta = RowDatasetVersionMeta::from_sequence(&new_seq)?;
812        fragment.last_updated_at_version_meta = Some(new_meta);
813    }
814
815    Ok(())
816}
817
818// Protobuf conversion implementations
819impl TryFrom<pb::data_fragment::LastUpdatedAtVersionSequence> for RowDatasetVersionMeta {
820    type Error = Error;
821
822    fn try_from(value: pb::data_fragment::LastUpdatedAtVersionSequence) -> Result<Self> {
823        match value {
824            pb::data_fragment::LastUpdatedAtVersionSequence::InlineLastUpdatedAtVersions(data) => {
825                Ok(Self::Inline(Arc::from(data)))
826            }
827            pb::data_fragment::LastUpdatedAtVersionSequence::ColumnLastUpdatedAtVersions(_) => {
828                Ok(Self::Column)
829            }
830        }
831    }
832}
833
834impl TryFrom<pb::data_fragment::CreatedAtVersionSequence> for RowDatasetVersionMeta {
835    type Error = Error;
836
837    fn try_from(value: pb::data_fragment::CreatedAtVersionSequence) -> Result<Self> {
838        match value {
839            pb::data_fragment::CreatedAtVersionSequence::InlineCreatedAtVersions(data) => {
840                Ok(Self::Inline(Arc::from(data)))
841            }
842            pb::data_fragment::CreatedAtVersionSequence::ColumnCreatedAtVersions(_) => {
843                Ok(Self::Column)
844            }
845        }
846    }
847}
848
849#[cfg(test)]
850mod tests {
851    use super::*;
852
853    #[test]
854    fn column_meta_json_is_an_empty_object_under_the_arm_name() {
855        let json = serde_json::to_string(&RowDatasetVersionMeta::Column).unwrap();
856        assert_eq!(json, r#"{"column":{}}"#);
857        let parsed: RowDatasetVersionMeta = serde_json::from_str(&json).unwrap();
858        assert_eq!(parsed, RowDatasetVersionMeta::Column);
859        let inline = RowDatasetVersionMeta::Inline(Arc::from(vec![1u8, 2, 3]));
860        let parsed: RowDatasetVersionMeta =
861            serde_json::from_str(&serde_json::to_string(&inline).unwrap()).unwrap();
862        assert_eq!(parsed, inline);
863    }
864
865    #[test]
866    fn from_versions_run_length_encodes() {
867        assert!(
868            RowDatasetVersionSequence::from_versions(&[])
869                .runs
870                .is_empty()
871        );
872
873        let single = RowDatasetVersionSequence::from_versions(&[3, 3, 3]);
874        assert_eq!(single.runs.len(), 1);
875        assert_eq!(single.runs[0].version, 3);
876        assert_eq!(single.len(), 3);
877
878        let alternating = RowDatasetVersionSequence::from_versions(&[1, 2, 1, 2]);
879        assert_eq!(
880            alternating
881                .runs
882                .iter()
883                .map(|run| run.version)
884                .collect::<Vec<_>>(),
885            [1, 2, 1, 2]
886        );
887        assert_eq!(alternating.versions().collect::<Vec<_>>(), [1, 2, 1, 2]);
888    }
889
890    #[test]
891    fn test_version_random_access() {
892        let seq = RowDatasetVersionSequence {
893            runs: vec![
894                RowDatasetVersionRun {
895                    span: U64Segment::Range(0..3),
896                    version: 1,
897                },
898                RowDatasetVersionRun {
899                    span: U64Segment::Range(0..2),
900                    version: 2,
901                },
902                RowDatasetVersionRun {
903                    span: U64Segment::Range(0..1),
904                    version: 3,
905                },
906            ],
907        };
908        assert_eq!(seq.version_at(0), Some(1));
909        assert_eq!(seq.version_at(2), Some(1));
910        assert_eq!(seq.version_at(3), Some(2));
911        assert_eq!(seq.version_at(4), Some(2));
912        assert_eq!(seq.version_at(5), Some(3));
913        assert_eq!(seq.version_at(6), None);
914    }
915
916    #[test]
917    fn test_partial_refresh_streams_many_lineage_runs() {
918        const ROWS: usize = 10_000;
919        let prior_sequence = RowDatasetVersionSequence {
920            runs: (0..ROWS)
921                .map(|position| RowDatasetVersionRun {
922                    span: U64Segment::Range(position as u64..position as u64 + 1),
923                    version: (position % 2 + 1) as u64,
924                })
925                .collect(),
926        };
927        let mut fragment = Fragment::new(1);
928        fragment.physical_rows = Some(ROWS);
929        fragment.last_updated_at_version_meta =
930            Some(RowDatasetVersionMeta::from_sequence(&prior_sequence).unwrap());
931
932        refresh_row_latest_update_meta_for_partial_frag_rewrite_cols(
933            &mut fragment,
934            &[ROWS - 1],
935            3,
936            1,
937            &Default::default(),
938        )
939        .unwrap();
940
941        let refreshed = fragment
942            .last_updated_at_version_meta
943            .unwrap()
944            .load_sequence()
945            .unwrap();
946        assert_eq!(refreshed.len(), ROWS as u64);
947        assert_eq!(refreshed.version_at(0), Some(1));
948        assert_eq!(refreshed.version_at(1), Some(2));
949        assert_eq!(refreshed.version_at(ROWS - 2), Some(1));
950        assert_eq!(refreshed.version_at(ROWS - 1), Some(3));
951    }
952
953    #[test]
954    fn test_serialization_round_trip() {
955        let seq = RowDatasetVersionSequence {
956            runs: vec![
957                RowDatasetVersionRun {
958                    span: U64Segment::Range(0..4),
959                    version: 42,
960                },
961                RowDatasetVersionRun {
962                    span: U64Segment::Range(0..3),
963                    version: 99,
964                },
965            ],
966        };
967        let bytes = write_dataset_versions(&seq);
968        let seq2 = read_dataset_versions(&bytes).unwrap();
969        assert_eq!(seq2.runs.len(), 2);
970        assert_eq!(seq2.len(), 7);
971        assert_eq!(seq2.version_at(0), Some(42));
972        assert_eq!(seq2.version_at(5), Some(99));
973    }
974
975    #[test]
976    fn test_get_version_for_row_id() {
977        let seq = RowDatasetVersionSequence {
978            runs: vec![
979                RowDatasetVersionRun {
980                    span: U64Segment::Range(0..2),
981                    version: 8,
982                },
983                RowDatasetVersionRun {
984                    span: U64Segment::Range(0..2),
985                    version: 9,
986                },
987            ],
988        };
989        let rows = RowIdSequence::from(10..14); // row ids: 10,11,12,13
990        assert_eq!(seq.get_version_for_row_id(&rows, 10), Some(8));
991        assert_eq!(seq.get_version_for_row_id(&rows, 11), Some(8));
992        assert_eq!(seq.get_version_for_row_id(&rows, 12), Some(9));
993        assert_eq!(seq.get_version_for_row_id(&rows, 13), Some(9));
994        assert_eq!(seq.get_version_for_row_id(&rows, 99), None);
995    }
996
997    #[test]
998    fn test_version_cursor_ranges_gaps_and_rewind() {
999        let seq = RowDatasetVersionSequence {
1000            runs: vec![
1001                RowDatasetVersionRun {
1002                    span: U64Segment::Range(0..3),
1003                    version: 10,
1004                },
1005                RowDatasetVersionRun {
1006                    span: U64Segment::Range(3..3),
1007                    version: 99,
1008                },
1009                RowDatasetVersionRun {
1010                    span: U64Segment::Range(3..5),
1011                    version: 20,
1012                },
1013                RowDatasetVersionRun {
1014                    span: U64Segment::Range(5..9),
1015                    version: 30,
1016                },
1017            ],
1018        };
1019        let expected = [10, 10, 10, 20, 20, 30, 30, 30, 30];
1020        let mut cursor = seq.cursor();
1021        let mut actual = Vec::new();
1022
1023        cursor.extend_range(&seq, 0..2, &mut actual).unwrap();
1024        cursor.extend_range(&seq, 2..5, &mut actual).unwrap();
1025        cursor.extend_range(&seq, 7..9, &mut actual).unwrap();
1026        assert_eq!(actual, [10, 10, 10, 20, 20, 30, 30]);
1027
1028        actual.clear();
1029        cursor.extend_range(&seq, 1..6, &mut actual).unwrap();
1030        assert_eq!(actual, expected[1..6]);
1031
1032        cursor.extend_range(&seq, 6..6, &mut actual).unwrap();
1033        assert_eq!(actual, expected[1..6]);
1034    }
1035
1036    #[test]
1037    fn test_version_cursor_descending_ranges_use_indexed_seek() {
1038        const RUNS: usize = 10_000;
1039        let seq = RowDatasetVersionSequence {
1040            runs: (0..RUNS)
1041                .map(|position| RowDatasetVersionRun {
1042                    span: U64Segment::Range(position as u64..position as u64 + 1),
1043                    version: position as u64,
1044                })
1045                .collect(),
1046        };
1047        let mut cursor = seq.cursor();
1048        let mut actual = Vec::with_capacity(RUNS);
1049
1050        for position in (0..RUNS).rev() {
1051            cursor
1052                .extend_range(&seq, position..position + 1, &mut actual)
1053                .unwrap();
1054        }
1055
1056        assert_eq!(actual, (0..RUNS as u64).rev().collect::<Vec<_>>());
1057        assert_eq!(cursor.run_offsets.as_ref().unwrap().len(), RUNS);
1058    }
1059
1060    #[test]
1061    fn test_version_cursor_alternating_ranges_use_indexed_seek_both_directions() {
1062        const RUNS: usize = 10_000;
1063        let seq = RowDatasetVersionSequence {
1064            runs: (0..RUNS)
1065                .map(|position| RowDatasetVersionRun {
1066                    span: U64Segment::Range(position as u64..position as u64 + 1),
1067                    version: position as u64,
1068                })
1069                .collect(),
1070        };
1071        let mut cursor = seq.cursor();
1072        let mut actual = Vec::with_capacity(RUNS);
1073
1074        for selection_index in 0..RUNS {
1075            let position = if selection_index % 2 == 0 {
1076                RUNS - 1
1077            } else {
1078                0
1079            };
1080            cursor
1081                .extend_range(&seq, position..position + 1, &mut actual)
1082                .unwrap();
1083        }
1084
1085        let expected = (0..RUNS)
1086            .map(|selection_index| {
1087                if selection_index % 2 == 0 {
1088                    (RUNS - 1) as u64
1089                } else {
1090                    0
1091                }
1092            })
1093            .collect::<Vec<_>>();
1094        assert_eq!(actual, expected);
1095        assert_eq!(cursor.run_offsets.as_ref().unwrap().len(), RUNS);
1096        assert_eq!(cursor.indexed_seek_count, RUNS - 1);
1097    }
1098
1099    #[test]
1100    fn test_version_cursor_reports_out_of_bounds() {
1101        let seq = RowDatasetVersionSequence::from_uniform_row_count(4, 7);
1102        let mut cursor = seq.cursor();
1103        let mut actual = Vec::new();
1104        let error = cursor.extend_range(&seq, 4..5, &mut actual).unwrap_err();
1105        assert!(matches!(error, Error::Internal { .. }));
1106        assert!(
1107            error
1108                .to_string()
1109                .contains("position 4 out of range (total_len=4)")
1110        );
1111        assert!(actual.is_empty());
1112    }
1113
1114    #[test]
1115    fn test_version_cursor_non_range_span() {
1116        let non_range_span = U64Segment::from_slice(&[0, 2, 4, 6, 8]);
1117        assert!(!matches!(non_range_span, U64Segment::Range(_)));
1118        let seq = RowDatasetVersionSequence {
1119            runs: vec![
1120                RowDatasetVersionRun {
1121                    span: non_range_span,
1122                    version: 4,
1123                },
1124                RowDatasetVersionRun {
1125                    span: U64Segment::Range(0..3),
1126                    version: 8,
1127                },
1128            ],
1129        };
1130        let mut actual = Vec::new();
1131        seq.cursor().extend_range(&seq, 0..8, &mut actual).unwrap();
1132        assert_eq!(actual, [4, 4, 4, 4, 4, 8, 8, 8]);
1133    }
1134}