1use 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#[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
41pub type SpilledRowLineage = HashMap<u64, LoadedRowLineage>;
43
44#[derive(Debug, Clone, PartialEq, Eq, DeepSizeOf)]
50pub struct RowDatasetVersionRun {
51 pub span: U64Segment,
52 pub version: u64,
53}
54
55impl RowDatasetVersionRun {
56 pub fn len(&self) -> usize {
58 self.span.len()
59 }
60
61 pub fn is_empty(&self) -> bool {
63 self.span.is_empty()
64 }
65
66 pub fn version(&self) -> u64 {
68 self.version
69 }
70}
71
72#[derive(Debug, Clone, PartialEq, Eq, DeepSizeOf, Default)]
78pub struct RowDatasetVersionSequence {
79 pub runs: Vec<RowDatasetVersionRun>,
80}
81
82#[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 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 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 pub fn new() -> Self {
215 Self { runs: Vec::new() }
216 }
217
218 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 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 pub fn len(&self) -> u64 {
254 self.runs.iter().map(|s| s.len() as u64).sum()
255 }
256
257 pub fn is_empty(&self) -> bool {
259 self.runs.is_empty() || self.runs.iter().all(|s| s.is_empty())
260 }
261
262 pub fn versions(&self) -> VersionsIter<'_> {
264 VersionsIter::new(&self.runs)
265 }
266
267 pub(crate) fn cursor(&self) -> RowDatasetVersionCursor {
269 RowDatasetVersionCursor::default()
270 }
271
272 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 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 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 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
344pub 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 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#[derive(Debug, Clone, PartialEq, Eq, DeepSizeOf)]
398pub enum RowDatasetVersionMeta {
399 Inline(Arc<[u8]>),
401 Column,
411}
412
413#[derive(Serialize, Deserialize)]
418struct ColumnMarker {}
419
420impl 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
443impl<'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 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 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
486pub 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
504pub 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
520pub fn write_dataset_versions(sequence: &RowDatasetVersionSequence) -> Vec<u8> {
522 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
537pub 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
561pub 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
654pub 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 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 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
680pub 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 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
713pub 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 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 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 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 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 for &pos in updated_offsets {
785 if pos < base_versions.len() {
786 base_versions[pos] = current_version;
787 }
788 }
789
790 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
818impl 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); 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}