use std::{ path::Path, time::{Duration, Instant}, }; use anyhow::{Context, Result, bail}; use binary_heap_plus::BinaryHeap; use tantivy::{ Index, collector::DocSetCollector, query::TermQuery, schema::{Field, IndexRecordOption}, }; use crate::{ compactor::obsolete_keys::ObsoleteRowEntryTempFile, reader::frag_reader::FragReaderMeta, utils::{ self, dedup_bitset::FragDocBitSet, key::ShardingKeyRange, metrics::{DEDUP_STAGE_LATENCY, begin_dedup_metrics}, row_entry::{RowEntry, RowEntryReader}, schema::HANDLE_STORE_NAME, }, }; struct RowEntryDedupSource { reader: RowEntryReader, start: Vec, end: Vec, obsolete_entries: ObsoleteRowEntryTempFile, in_range_key_count: usize, obsolete_by_range_count: usize, trim_out_of_range_time: Duration, } impl RowEntryDedupSource { fn open( shard_range: &ShardingKeyRange, entries_path: &str, obsolete_keys_temp_dir: &Path, ) -> Result<(Self, usize)> { let reader = RowEntryReader::open(entries_path)?; let num_rows = reader.num_rows(); let obsolete_entries = ObsoleteRowEntryTempFile::new_in(obsolete_keys_temp_dir)?; Ok(( Self { reader, start: shard_range.start().raw().to_vec(), end: shard_range.end().raw().to_vec(), obsolete_entries, in_range_key_count: 0, obsolete_by_range_count: 0, trim_out_of_range_time: Duration::ZERO, }, num_rows, )) } fn next_in_range(&mut self) -> Result> { let mut trim_start: Option = None; loop { let entry = match self.reader.next_entry()? { Some(entry) => entry, None => { if let Some(trim_start) = trim_start { self.trim_out_of_range_time += trim_start.elapsed(); } return Ok(None); } }; if entry.key.as_slice() < self.start.as_slice() || entry.key.as_slice() >= self.end.as_slice() { if trim_start.is_none() { trim_start = Some(Instant::now()); } self.obsolete_entries.push(&entry)?; self.obsolete_by_range_count += 1; continue; } if let Some(trim_start) = trim_start { self.trim_out_of_range_time += trim_start.elapsed(); } self.in_range_key_count += 1; return Ok(Some(entry)); } } fn mark_obsolete_by_dedup(&mut self, entry: &RowEntry) -> Result<()> { self.obsolete_entries.push(entry) } } trait MergeRowEntrySource { fn next_entry(&mut self) -> Result>; fn mark_obsolete(&mut self, entry: &RowEntry) -> Result<()>; } impl MergeRowEntrySource for RowEntryDedupSource { fn next_entry(&mut self) -> Result> { self.next_in_range() } fn mark_obsolete(&mut self, entry: &RowEntry) -> Result<()> { self.mark_obsolete_by_dedup(entry) } } fn merge_and_collect_obsolete(sources: &mut [S], mut on_retain: R) -> Result<()> where S: MergeRowEntrySource, R: FnMut(RowEntry) -> Result<()>, { struct StreamCursor { src_idx: usize, entry: RowEntry, } let mut heap = BinaryHeap::new_by(|left: &StreamCursor, right: &StreamCursor| { right .entry .key .cmp(&left.entry.key) .then_with(|| left.entry.version.cmp(&right.entry.version)) .then_with(|| left.src_idx.cmp(&right.src_idx)) }); for (src_idx, source) in sources.iter_mut().enumerate() { if let Some(entry) = source.next_entry()? { heap.push(StreamCursor { src_idx, entry }); } } while let Some(cursor) = heap.pop() { // The merge logic assumes each source emits unique keys within the // current in-range stream. We still prefetch before deciding which // duplicate is obsolete, so that prefetch may already trim a later // out-of-range key into the obsolete temp file before we append the // current duplicate key. // Due to the heap's comparator (Min-Key, Max-Version, Max-SrcIdx), // the first popped cursor for any given key is strictly the newest/highest // version. Therefore, it is guaranteed to be the retained entry. if let Some(next_entry) = sources[cursor.src_idx].next_entry()? { heap.push(StreamCursor { src_idx: cursor.src_idx, entry: next_entry, }); } while let Some(top_cursor) = heap.peek() { if top_cursor.entry.key != cursor.entry.key { break; } // Since heap.peek() is Some, heap.pop() is also Some. let next_cursor = heap.pop().context("heap pop failed")?; if let Some(next_entry) = sources[next_cursor.src_idx].next_entry()? { heap.push(StreamCursor { src_idx: next_cursor.src_idx, entry: next_entry, }); } sources[next_cursor.src_idx].mark_obsolete(&next_cursor.entry)?; } on_retain(cursor.entry)?; } Ok(()) } #[cfg(any(test, feature = "bench"))] pub fn merge_row_entries_for_bench( sorted_entries_list: &[Vec], ) -> Result<(usize, usize)> { let mut retained_count = 0; let total_entry_count: usize = sorted_entries_list.iter().map(Vec::len).sum(); #[cfg(any(test, feature = "bench"))] struct SliceRowEntrySource<'a> { entries: &'a [RowEntry], pos: usize, } #[cfg(any(test, feature = "bench"))] impl<'a> SliceRowEntrySource<'a> { fn new(entries: &'a [RowEntry]) -> Self { Self { entries, pos: 0 } } } #[cfg(any(test, feature = "bench"))] impl MergeRowEntrySource for SliceRowEntrySource<'_> { fn next_entry(&mut self) -> Result> { if self.pos >= self.entries.len() { return Ok(None); } let entry = self.entries[self.pos].clone(); self.pos += 1; Ok(Some(entry)) } fn mark_obsolete(&mut self, _entry: &RowEntry) -> Result<()> { Ok(()) } } let mut sources: Vec> = sorted_entries_list .iter() .map(|entries| SliceRowEntrySource::new(entries)) .collect(); merge_and_collect_obsolete(&mut sources, |_entry| { retained_count += 1; Ok(()) })?; let obsolete_count = total_entry_count - retained_count; Ok((retained_count, obsolete_count)) } #[derive(Default, Debug)] pub struct FragsDedupStat { pub input_frag_count: usize, pub input_key_count: usize, pub in_range_key_count: usize, pub retained_key_count: usize, pub retained_delete_mark_count: usize, pub obsolete_by_dedup_count: usize, pub obsolete_by_range_count: usize, pub load_index_time: Duration, pub load_entries_time: Duration, pub trim_out_of_range_time: Duration, pub merge_and_collect_obsolete_time: Duration, pub build_obsolete_doc_bitsets_time: Duration, pub total_dedup_time: Duration, } fn observe_dedup_stage_latency(stage: &str, duration: Duration) { DEDUP_STAGE_LATENCY .with_label_values(&[stage]) .observe(duration.as_secs_f64()); } fn load_indexes(frag_reader_metas: &[FragReaderMeta]) -> Result> { let mut indexes = Vec::with_capacity(frag_reader_metas.len()); for frag in frag_reader_metas { indexes.push(frag.load_index()?); } Ok(indexes) } // Ensure all frags have the same schema. fn check_schema(indexes: &[Index], frag_reader_metas: &[FragReaderMeta]) -> Result { let schema = indexes[0].schema(); for i in 1..indexes.len() { if indexes[i].schema() != schema { bail!( "Schema mismatch, frag1={}, schema1={:?}, frag2={}, schema2={:?}", frag_reader_metas[0].local_frag_path, schema, frag_reader_metas[i].local_frag_path, indexes[i].schema(), ); } } let handle_field = schema.get_field(HANDLE_STORE_NAME).with_context(|| { format!( "Get schema field {} failed, schema={:?}, frag_path={}", HANDLE_STORE_NAME, schema, frag_reader_metas[0].local_frag_path ) })?; Ok(handle_field) } pub struct FragsDedupOutput { pub handle_field: Field, pub dedup_stat: FragsDedupStat, pub obsolete_doc_bitsets: Vec, } fn empty_obsolete_doc_bitsets(indexes: &[Index]) -> Result> { indexes.iter().map(FragDocBitSet::empty_for_index).collect() } fn build_obsolete_doc_bitsets( indexes: &[Index], frag_reader_metas: &[FragReaderMeta], sources: &mut [RowEntryDedupSource], handle_field: Field, dedup_stat: &mut FragsDedupStat, ) -> Result> { if dedup_stat.obsolete_by_dedup_count == 0 && dedup_stat.obsolete_by_range_count == 0 { return empty_obsolete_doc_bitsets(indexes); } let start = Instant::now(); let mut searchers = Vec::with_capacity(indexes.len()); let mut frag_doc_bitsets = Vec::with_capacity(indexes.len()); for (index, frag_meta) in indexes.iter().zip(frag_reader_metas) { let index_reader = index .reader_builder() .reload_policy(tantivy::ReloadPolicy::Manual) .doc_store_cache_num_blocks(0) .try_into() .with_context(|| { format!( "create index reader for obsolete doc bitset failed, frag_path={}", frag_meta.local_frag_path ) })?; let searcher = index_reader.searcher(); frag_doc_bitsets.push(FragDocBitSet::empty_for_index(index)?); searchers.push(searcher); } for (((frag_doc_bitset, searcher), frag_meta), source) in frag_doc_bitsets .iter_mut() .zip(searchers.iter()) .zip(frag_reader_metas.iter()) .zip(sources.iter_mut()) { let mut obsolete_reader = source.obsolete_entries.reader()?; while let Some(entry) = obsolete_reader.next_entry()? { let handle_term = tantivy::Term::from_field_bytes(handle_field, &entry.key); let doc_addrs = searcher .search( &TermQuery::new(handle_term, IndexRecordOption::Basic), &DocSetCollector, ) .with_context(|| { format!( "search obsolete key failed, frag_path={}, key={:?}", frag_meta.local_frag_path, entry.key ) })?; for doc_addr in doc_addrs { frag_doc_bitset.push_deleted_doc(doc_addr)?; } } frag_doc_bitset.sort_and_dedup(); } dedup_stat.build_obsolete_doc_bitsets_time = start.elapsed(); Ok(frag_doc_bitsets) } fn dedup_frags_inner( shard_range: &ShardingKeyRange, frag_reader_metas: &[FragReaderMeta], mut on_retain: R, obsolete_keys_temp_dir: &Path, ) -> Result where R: FnMut(RowEntry) -> Result<()>, { let total_start = Instant::now(); let mut dedup_stat = FragsDedupStat { input_frag_count: frag_reader_metas.len(), ..Default::default() }; if frag_reader_metas.is_empty() { return Ok(FragsDedupOutput { handle_field: Field::from_field_id(0), dedup_stat, obsolete_doc_bitsets: Vec::new(), }); } let start = Instant::now(); let indexes = load_indexes(frag_reader_metas)?; let handle_field = check_schema(&indexes, frag_reader_metas)?; dedup_stat.load_index_time = start.elapsed(); observe_dedup_stage_latency("load_index", dedup_stat.load_index_time); let start = Instant::now(); let mut sources = Vec::with_capacity(frag_reader_metas.len()); for frag_reader_meta in frag_reader_metas { let entries_path = utils::frag_row_entry_path(&frag_reader_meta.local_frag_path); let (source, num_rows) = RowEntryDedupSource::open(shard_range, &entries_path, obsolete_keys_temp_dir)?; dedup_stat.input_key_count += num_rows; sources.push(source); } dedup_stat.load_entries_time = start.elapsed(); observe_dedup_stage_latency("load_entries", dedup_stat.load_entries_time); let start = Instant::now(); let mut retained_key_count = 0; let mut retained_delete_mark_count = 0; merge_and_collect_obsolete(&mut sources, |entry| { retained_key_count += 1; retained_delete_mark_count += entry.del_mark as usize; on_retain(entry) })?; for source in &sources { dedup_stat.in_range_key_count += source.in_range_key_count; dedup_stat.obsolete_by_range_count += source.obsolete_by_range_count; dedup_stat.trim_out_of_range_time += source.trim_out_of_range_time; } dedup_stat.obsolete_by_dedup_count = dedup_stat.in_range_key_count - retained_key_count; dedup_stat.retained_key_count = retained_key_count; dedup_stat.retained_delete_mark_count = retained_delete_mark_count; dedup_stat.merge_and_collect_obsolete_time = start.elapsed(); observe_dedup_stage_latency( "merge_and_collect_obsolete", dedup_stat.merge_and_collect_obsolete_time, ); observe_dedup_stage_latency("trim_out_of_range", dedup_stat.trim_out_of_range_time); for source in &mut sources { source.obsolete_entries.finalize()?; } let obsolete_doc_bitsets = build_obsolete_doc_bitsets( &indexes, frag_reader_metas, &mut sources, handle_field, &mut dedup_stat, )?; dedup_stat.total_dedup_time = total_start.elapsed(); observe_dedup_stage_latency( "build_obsolete_doc_bitsets", dedup_stat.build_obsolete_doc_bitsets_time, ); observe_dedup_stage_latency("total", dedup_stat.total_dedup_time); Ok(FragsDedupOutput { handle_field, dedup_stat, obsolete_doc_bitsets, }) } pub fn dedup_frags( shard_range: &ShardingKeyRange, frag_reader_metas: &[FragReaderMeta], on_retain: R, obsolete_keys_temp_dir: &Path, ) -> Result where R: FnMut(RowEntry) -> Result<()>, { let dedup_metrics = begin_dedup_metrics(); dedup_frags_inner( shard_range, frag_reader_metas, on_retain, obsolete_keys_temp_dir, ) .inspect_err(|_| { dedup_metrics.mark_error(); }) } #[cfg(test)] mod tests { use std::{cell::RefCell, rc::Rc, time::Duration}; use tantivy::{TantivyDocument, collector::DocSetCollector, query::AllQuery, schema::Value}; use tempfile::TempDir; use super::*; use crate::test_utils::create_test_frag_path; struct MockRowEntrySource { entries: Vec, obsolete_entries: Vec, obsolete_events: Rc>>, src_idx: usize, pos: usize, fail_next_at: Option, fail_mark_obsolete: bool, } impl MockRowEntrySource { fn new( entries: Vec, src_idx: usize, obsolete_events: Rc>>, ) -> Self { Self { entries, obsolete_entries: Vec::new(), obsolete_events, src_idx, pos: 0, fail_next_at: None, fail_mark_obsolete: false, } } } impl MergeRowEntrySource for MockRowEntrySource { fn next_entry(&mut self) -> Result> { if self.fail_next_at == Some(self.pos) { bail!("next_entry failed at pos={}", self.pos); } if self.pos >= self.entries.len() { return Ok(None); } let entry = self.entries[self.pos].clone(); self.pos += 1; Ok(Some(entry)) } fn mark_obsolete(&mut self, entry: &RowEntry) -> Result<()> { if self.fail_mark_obsolete { bail!("mark_obsolete failed"); } self.obsolete_entries.push(entry.clone()); self.obsolete_events .borrow_mut() .push((self.src_idx, entry.clone())); Ok(()) } } fn merge_lists(lists: &[&[RowEntry]]) -> (Vec, Vec<(usize, RowEntry)>) { let obsolete_events = Rc::new(RefCell::new(Vec::new())); let mut sources: Vec = lists .iter() .enumerate() .map(|(src_idx, entries)| { MockRowEntrySource::new(entries.to_vec(), src_idx, Rc::clone(&obsolete_events)) }) .collect(); let mut collected = Vec::new(); merge_and_collect_obsolete(&mut sources, |e| { collected.push(e); Ok(()) }) .unwrap(); let obsolete = obsolete_events.borrow().clone(); (collected, obsolete) } #[test] fn test_merge_del_mark_keep_latest() { let list0 = vec![ RowEntry::new(b"a".to_vec(), false, 0), RowEntry::new(b"b".to_vec(), false, 0), RowEntry::new(b"c".to_vec(), true, 0), ]; let list1 = vec![ RowEntry::new(b"a".to_vec(), true, 0), RowEntry::new(b"b".to_vec(), false, 0), RowEntry::new(b"c".to_vec(), false, 0), ]; let list2 = vec![ RowEntry::new(b"a".to_vec(), false, 0), RowEntry::new(b"d".to_vec(), true, 0), ]; let (collected, obsolete) = merge_lists(&[&list0, &list1, &list2]); // "a" appears in all three lists, list2 is the latest (src 2, del_mark=false) // "b" appears in list0 and list1, list1 is the latest (src 1, del_mark=false) // "c" appears in list0 and list1, list1 is the latest (src 1, del_mark=false) // "d" only appears in list2 (src 2, del_mark=true) assert_eq!( collected, vec![ RowEntry::new(b"a".to_vec(), false, 0), RowEntry::new(b"b".to_vec(), false, 0), RowEntry::new(b"c".to_vec(), false, 0), RowEntry::new(b"d".to_vec(), true, 0), ] ); // Obsolete duplicates are emitted in heap order, so the newer duplicate // is marked obsolete before older ones. // "a" in list1 and list0 are obsolete // "b" in list0 is obsolete // "c" in list0 is obsolete assert_eq!( obsolete, vec![ (1, RowEntry::new(b"a".to_vec(), true, 0)), (0, RowEntry::new(b"a".to_vec(), false, 0)), (0, RowEntry::new(b"b".to_vec(), false, 0)), (0, RowEntry::new(b"c".to_vec(), true, 0)), ] ); } #[test] fn test_merge_prefers_higher_version_over_source_order() { let list0 = vec![RowEntry::new(b"a".to_vec(), false, 9)]; let list1 = vec![RowEntry::new(b"a".to_vec(), true, 3)]; let list2 = vec![RowEntry::new(b"a".to_vec(), true, 7)]; let (collected, obsolete) = merge_lists(&[&list0, &list1, &list2]); assert_eq!(collected, vec![RowEntry::new(b"a".to_vec(), false, 9)]); assert_eq!( obsolete, vec![ (2, RowEntry::new(b"a".to_vec(), true, 7)), (1, RowEntry::new(b"a".to_vec(), true, 3)), ] ); } #[test] fn test_empty_lists() { let list0: Vec = vec![]; let list1: Vec = vec![]; let list2: Vec = vec![]; let (collected, obsolete) = merge_lists(&[&list0, &list1, &list2]); assert!(collected.is_empty()); assert!(obsolete.is_empty()); } #[test] fn test_no_duplicates() { let list0 = vec![RowEntry::new(b"apple".to_vec(), false, 0)]; let list1 = vec![RowEntry::new(b"banana".to_vec(), true, 0)]; let list2 = vec![RowEntry::new(b"cherry".to_vec(), false, 0)]; let (collected, obsolete) = merge_lists(&[&list0, &list1, &list2]); assert_eq!( collected, vec![ RowEntry::new(b"apple".to_vec(), false, 0), RowEntry::new(b"banana".to_vec(), true, 0), RowEntry::new(b"cherry".to_vec(), false, 0), ] ); assert!(obsolete.is_empty()); } #[test] fn test_single_list_with_unique_keys() { let list0 = vec![ RowEntry::new(b"key1".to_vec(), false, 0), RowEntry::new(b"key2".to_vec(), false, 0), RowEntry::new(b"key3".to_vec(), true, 0), ]; let (collected, obsolete) = merge_lists(&[&list0]); assert_eq!( collected, vec![ RowEntry::new(b"key1".to_vec(), false, 0), RowEntry::new(b"key2".to_vec(), false, 0), RowEntry::new(b"key3".to_vec(), true, 0), ] ); assert!(obsolete.is_empty()); } #[test] fn test_mixed_empty_and_non_empty() { let list0 = vec![]; let list1 = vec![RowEntry::new(b"key1".to_vec(), false, 0)]; let list2 = vec![]; let list3 = vec![RowEntry::new(b"key1".to_vec(), true, 0)]; let (collected, obsolete) = merge_lists(&[&list0, &list1, &list2, &list3]); assert_eq!(collected, vec![RowEntry::new(b"key1".to_vec(), true, 0)]); assert_eq!( obsolete, vec![(1, RowEntry::new(b"key1".to_vec(), false, 0))] ); } #[test] fn test_binary_keys() { let key1 = vec![0x00, 0x01, 0x02]; let key2 = vec![0x00, 0x01, 0x02, 0x03]; // key1 is prefix of key2 let key3 = vec![0xFF, 0xFE]; let list0 = vec![ RowEntry::new(key1.clone(), false, 0), RowEntry::new(key3.clone(), true, 0), ]; let list1 = vec![RowEntry::new(key2.clone(), false, 0)]; let list2 = vec![RowEntry::new(key1.clone(), true, 0)]; let (collected, obsolete) = merge_lists(&[&list0, &list1, &list2]); assert_eq!( collected, vec![ RowEntry::new(key1.clone(), true, 0), RowEntry::new(key2.clone(), false, 0), RowEntry::new(key3.clone(), true, 0), ] ); assert_eq!(obsolete, vec![(0, RowEntry::new(key1.clone(), false, 0))]); } #[test] fn test_del_mark_transitions() { // Test all possible del_mark transitions: false->false, false->true, // true->false, true->true let list0 = vec![ RowEntry::new(b"k1".to_vec(), false, 0), RowEntry::new(b"k2".to_vec(), false, 0), RowEntry::new(b"k3".to_vec(), true, 0), RowEntry::new(b"k4".to_vec(), true, 0), ]; let list1 = vec![ RowEntry::new(b"k1".to_vec(), false, 0), // false -> false RowEntry::new(b"k2".to_vec(), true, 0), // false -> true RowEntry::new(b"k3".to_vec(), false, 0), // true -> false RowEntry::new(b"k4".to_vec(), true, 0), // true -> true ]; let (collected, obsolete) = merge_lists(&[&list0, &list1]); assert_eq!( collected, vec![ RowEntry::new(b"k1".to_vec(), false, 0), RowEntry::new(b"k2".to_vec(), true, 0), RowEntry::new(b"k3".to_vec(), false, 0), RowEntry::new(b"k4".to_vec(), true, 0), ] ); // All keys from list0 should be obsolete assert_eq!( obsolete, vec![ (0, RowEntry::new(b"k1".to_vec(), false, 0)), (0, RowEntry::new(b"k2".to_vec(), false, 0)), (0, RowEntry::new(b"k3".to_vec(), true, 0)), (0, RowEntry::new(b"k4".to_vec(), true, 0)), ] ); } #[test] fn test_error_propagation() { let list0 = vec![ RowEntry::new(b"a".to_vec(), false, 0), RowEntry::new(b"b".to_vec(), false, 0), ]; let list1 = vec![RowEntry::new(b"c".to_vec(), false, 0)]; let mut sources = vec![ MockRowEntrySource::new(list0, 0, Rc::new(RefCell::new(Vec::new()))), MockRowEntrySource::new(list1, 1, Rc::new(RefCell::new(Vec::new()))), ]; let result = merge_and_collect_obsolete(&mut sources, |_e| Err(anyhow::anyhow!("Test error"))); assert!(result.is_err()); assert_eq!(result.unwrap_err().to_string(), "Test error"); } #[test] fn test_next_entry_error_propagation() { let obsolete_events = Rc::new(RefCell::new(Vec::new())); let mut sources = vec![ MockRowEntrySource::new( vec![RowEntry::new(b"a".to_vec(), false, 0)], 0, Rc::clone(&obsolete_events), ), MockRowEntrySource { entries: vec![RowEntry::new(b"b".to_vec(), false, 0)], obsolete_entries: Vec::new(), obsolete_events, src_idx: 1, pos: 0, fail_next_at: Some(0), fail_mark_obsolete: false, }, ]; let result = merge_and_collect_obsolete(&mut sources, |_e| Ok(())); assert!(result.is_err()); assert_eq!( result.unwrap_err().to_string(), "next_entry failed at pos=0" ); } #[test] fn test_mark_obsolete_error_propagation() { let obsolete_events = Rc::new(RefCell::new(Vec::new())); let mut sources = vec![ MockRowEntrySource { entries: vec![RowEntry::new(b"same".to_vec(), false, 0)], obsolete_entries: Vec::new(), obsolete_events: Rc::clone(&obsolete_events), src_idx: 0, pos: 0, fail_next_at: None, fail_mark_obsolete: true, }, MockRowEntrySource::new( vec![RowEntry::new(b"same".to_vec(), true, 0)], 1, obsolete_events, ), ]; let result = merge_and_collect_obsolete(&mut sources, |_e| Ok(())); assert!(result.is_err()); assert_eq!(result.unwrap_err().to_string(), "mark_obsolete failed"); } #[test] fn test_large_scale_merge() { use std::collections::{HashMap, HashSet}; // Create multiple lists with overlapping keys let num_lists = 10; let keys_per_list = 100; let mut lists = Vec::with_capacity(num_lists); for list_idx in 0..num_lists { let mut list = Vec::new(); for key_idx in 0..keys_per_list { // Create overlapping keys: each list shares some keys with others if key_idx % 3 == 0 || (key_idx + list_idx) % 5 == 0 { let key = format!("key{:04}", key_idx).into_bytes(); let del_mark = (key_idx + list_idx) % 2 == 0; list.push(RowEntry::new(key, del_mark, 0)); } } list.sort_by(|a, b| a.key.cmp(&b.key)); lists.push(list); } let mut all_keys = HashSet::new(); let mut last_seen: HashMap, usize> = HashMap::new(); for (i, list) in lists.iter().enumerate() { for e in list { all_keys.insert(e.key.clone()); last_seen.insert(e.key.clone(), i); } } let list_slices: Vec<&[RowEntry]> = lists.iter().map(|list| list.as_slice()).collect(); let mut collected = Vec::new(); let obsolete_events = Rc::new(RefCell::new(Vec::new())); let mut sources: Vec = list_slices .iter() .enumerate() .map(|(src_idx, entries)| { MockRowEntrySource::new(entries.to_vec(), src_idx, Rc::clone(&obsolete_events)) }) .collect(); merge_and_collect_obsolete(&mut sources, |e| { collected.push(e); Ok(()) }) .unwrap(); let obsolete_map: HashMap> = sources .iter() .enumerate() .filter(|(_, source)| !source.obsolete_entries.is_empty()) .map(|(src, source)| (src, source.obsolete_entries.clone())) .collect(); // Verify that collected keys are in sorted order for i in 1..collected.len() { assert!(collected[i - 1].key < collected[i].key); } // Verify that each key in collected appears in at least one input list for entry in &collected { assert!( all_keys.contains(&entry.key), "Key {:?} not found in any input list", entry.key ); } // Verify that obsolete keys are strictly older than the latest version for (src, entries) in &obsolete_map { for entry in entries { let last_src = last_seen.get(&entry.key).unwrap(); assert!( *last_src > *src, "Obsolete key {:?} from src {} should be older than last seen at {}", String::from_utf8_lossy(&entry.key), src, last_src ); } } } #[test] fn test_multiple_duplicates_across_lists() { let list0 = vec![ RowEntry::new(b"apple".to_vec(), false, 0), RowEntry::new(b"banana".to_vec(), false, 0), RowEntry::new(b"cherry".to_vec(), true, 0), ]; let list1 = vec![ RowEntry::new(b"apple".to_vec(), true, 0), RowEntry::new(b"cherry".to_vec(), false, 0), RowEntry::new(b"date".to_vec(), false, 0), ]; let list2 = vec![ RowEntry::new(b"apple".to_vec(), false, 0), RowEntry::new(b"banana".to_vec(), true, 0), RowEntry::new(b"date".to_vec(), true, 0), ]; let (collected, obsolete) = merge_lists(&[&list0, &list1, &list2]); assert_eq!( collected, vec![ RowEntry::new(b"apple".to_vec(), false, 0), // from list2 (latest) RowEntry::new(b"banana".to_vec(), true, 0), // from list2 (latest) RowEntry::new(b"cherry".to_vec(), false, 0), // from list1 (latest) RowEntry::new(b"date".to_vec(), true, 0), // from list2 (latest) ] ); assert_eq!( obsolete, vec![ (1, RowEntry::new(b"apple".to_vec(), true, 0)), // list1's apple is obsolete (0, RowEntry::new(b"apple".to_vec(), false, 0)), // list0's apple is obsolete (0, RowEntry::new(b"banana".to_vec(), false, 0)), // list0's banana is obsolete (0, RowEntry::new(b"cherry".to_vec(), true, 0)), // list0's cherry is obsolete (1, RowEntry::new(b"date".to_vec(), false, 0)), // list1's date is obsolete ] ); } #[test] fn test_edge_case_all_same_key() { let list0 = vec![RowEntry::new(b"same".to_vec(), false, 0)]; let list1 = vec![RowEntry::new(b"same".to_vec(), true, 0)]; let list2 = vec![RowEntry::new(b"same".to_vec(), false, 0)]; let list3 = vec![RowEntry::new(b"same".to_vec(), true, 0)]; let (collected, obsolete) = merge_lists(&[&list0, &list1, &list2, &list3]); // Should keep only the latest version (from list3) assert_eq!(collected, vec![RowEntry::new(b"same".to_vec(), true, 0)]); // All previous versions are obsolete, from newer duplicates to older // ones. assert_eq!( obsolete, vec![ (2, RowEntry::new(b"same".to_vec(), false, 0)), (1, RowEntry::new(b"same".to_vec(), true, 0)), (0, RowEntry::new(b"same".to_vec(), false, 0)), ] ); } #[test] fn test_obsolete_key_temp_file_round_trip() -> Result<()> { let mut file = ObsoleteRowEntryTempFile::new()?; file.push(&RowEntry::new(Vec::new(), false, 0))?; file.push(&RowEntry::new(vec![0, 1, 2], false, 0))?; file.push(&RowEntry::new(b"a".to_vec(), true, 0))?; file.finalize()?; let mut reader = file.reader()?; let mut entries = Vec::new(); while let Some(entry) = reader.next_entry()? { entries.push(entry); } assert_eq!( entries, vec![ RowEntry::new(Vec::new(), false, 0), RowEntry::new(vec![0, 1, 2], false, 0), RowEntry::new(b"a".to_vec(), true, 0), ] ); Ok(()) } #[test] fn test_row_entry_dedup_source_filters_range() -> Result<()> { let temp_dir = TempDir::new()?; let file_path = temp_dir.path().join("row_entry.parquet"); let entries = vec![ RowEntry::new(b"a".to_vec(), false, 0), RowEntry::new(b"b".to_vec(), false, 0), RowEntry::new(b"c".to_vec(), true, 0), RowEntry::new(b"d".to_vec(), false, 0), RowEntry::new(b"e".to_vec(), true, 0), ]; RowEntry::write_all(&file_path, &entries)?; let shard_range = ShardingKeyRange::half_open(b"b".to_vec(), b"e".to_vec())?; let entries_path = file_path.to_str().context("file_path to_str failed")?; let (mut source, num_rows) = RowEntryDedupSource::open(&shard_range, entries_path, temp_dir.path())?; assert_eq!(num_rows, entries.len()); let mut in_range_entries = Vec::new(); while let Some(entry) = source.next_in_range()? { in_range_entries.push(entry); } assert_eq!( in_range_entries, vec![ RowEntry::new(b"b".to_vec(), false, 0), RowEntry::new(b"c".to_vec(), true, 0), RowEntry::new(b"d".to_vec(), false, 0), ] ); assert_eq!(source.in_range_key_count, 3); assert_eq!(source.obsolete_by_range_count, 2); source.obsolete_entries.finalize()?; let mut obsolete_reader = source.obsolete_entries.reader()?; let mut obsolete_entries = Vec::new(); while let Some(entry) = obsolete_reader.next_entry()? { obsolete_entries.push(entry); } assert_eq!( obsolete_entries, vec![ RowEntry::new(b"a".to_vec(), false, 0), RowEntry::new(b"e".to_vec(), true, 0), ] ); Ok(()) } #[test] fn test_obsolete_entries_can_be_out_of_order_across_range_trim_and_dedup() -> Result<()> { let temp_dir = TempDir::new()?; let source0_path = temp_dir.path().join("source0.parquet"); let source1_path = temp_dir.path().join("source1.parquet"); let source0_entries = vec![ RowEntry::new(b"a".to_vec(), false, 0), RowEntry::new(b"b".to_vec(), false, 0), RowEntry::new(b"e".to_vec(), true, 0), ]; let source1_entries = vec![RowEntry::new(b"b".to_vec(), true, 0)]; RowEntry::write_all(&source0_path, &source0_entries)?; RowEntry::write_all(&source1_path, &source1_entries)?; let shard_range = ShardingKeyRange::half_open(b"b".to_vec(), b"d".to_vec())?; let source0_path_str = source0_path .to_str() .context("source0_path to_str failed")?; let source1_path_str = source1_path .to_str() .context("source1_path to_str failed")?; let (source0, _num_rows0) = RowEntryDedupSource::open(&shard_range, source0_path_str, temp_dir.path())?; let (source1, _num_rows1) = RowEntryDedupSource::open(&shard_range, source1_path_str, temp_dir.path())?; let mut sources = vec![source0, source1]; let mut retained = Vec::new(); merge_and_collect_obsolete(&mut sources, |entry| { retained.push(entry); Ok(()) })?; assert_eq!(retained, vec![RowEntry::new(b"b".to_vec(), true, 0)]); let mut source0_obsolete_reader = sources[0].obsolete_entries.reader()?; let mut source0_obsolete_entries = Vec::new(); while let Some(entry) = source0_obsolete_reader.next_entry()? { source0_obsolete_entries.push(entry); } // `a` is trimmed before merge starts. While processing duplicate key // `b`, source0 is prefetched first, so `e` is trimmed as obsolete // before the older `b` is marked obsolete by dedup. assert_eq!( source0_obsolete_entries, vec![ RowEntry::new(b"a".to_vec(), false, 0), RowEntry::new(b"e".to_vec(), true, 0), RowEntry::new(b"b".to_vec(), false, 0), ] ); let mut source1_obsolete_reader = sources[1].obsolete_entries.reader()?; let mut source1_obsolete_entries = Vec::new(); while let Some(entry) = source1_obsolete_reader.next_entry()? { source1_obsolete_entries.push(entry); } assert!(source1_obsolete_entries.is_empty()); Ok(()) } fn create_test_frag(base_dir: &Path, entries: &[RowEntry]) -> Result { let frag_path = create_test_frag_path(base_dir)?; std::fs::create_dir_all(&frag_path)?; // Write row entries file. let entries_path = crate::utils::frag_row_entry_path(&frag_path); RowEntry::write_all(&entries_path, entries)?; // Create a minimal index with the required handle field. let index_path = crate::utils::frag_index_path(&frag_path); std::fs::create_dir_all(&index_path)?; let mut schema_builder = tantivy::schema::Schema::builder(); let handle_field = schema_builder.add_bytes_field( HANDLE_STORE_NAME, tantivy::schema::INDEXED | tantivy::schema::STORED, ); let schema = schema_builder.build(); let index = Index::create_in_dir(&index_path, schema)?; let mut index_writer = crate::utils::create_index_writer(&index, 1, 16 * 1024 * 1024, &frag_path)?; for entry in entries { let mut doc = TantivyDocument::new(); doc.add_bytes(handle_field, &entry.key); index_writer.add_document(doc)?; } index_writer.commit()?; index_writer.wait_merging_threads()?; let min_handle = entries.first().map(|e| e.key.clone()).unwrap_or_default(); let max_handle = entries.last().map(|e| e.key.clone()).unwrap_or_default(); FragReaderMeta::new(&frag_path, min_handle, max_handle) } fn read_handles_from_index(frag_path: &str) -> Result>> { let index = Index::open_in_dir(crate::utils::frag_index_path(frag_path))?; let searcher = index.reader()?.searcher(); let handle_field = index.schema().get_field(HANDLE_STORE_NAME)?; let doc_addrs = searcher.search(&AllQuery, &DocSetCollector)?; let mut handles = Vec::with_capacity(doc_addrs.len()); for doc_addr in doc_addrs { let doc: TantivyDocument = searcher.doc(doc_addr)?; let handle = doc .get_first(handle_field) .context("missing handle field in doc")? .as_bytes() .context("handle field is not bytes")? .to_vec(); handles.push(handle); } handles.sort(); Ok(handles) } fn read_handles_from_index_with_doc_bitset( frag_path: &str, frag_doc_bitset: &FragDocBitSet, ) -> Result>> { let index = Index::open_in_dir(crate::utils::frag_index_path(frag_path))?; let handle_field = index.schema().get_field(HANDLE_STORE_NAME)?; let segments = index.searchable_segments()?; let custom_alive_bitsets = frag_doc_bitset.open_alive_bitsets()?; assert_eq!(segments.len(), custom_alive_bitsets.len()); let mut handles = Vec::new(); for (segment, custom_alive_bitset) in segments.iter().zip(custom_alive_bitsets) { let segment_reader = tantivy::index::SegmentReader::open_with_custom_alive_set( segment, custom_alive_bitset, )?; let store_reader = segment_reader.get_store_reader(1)?; for doc_id in 0..segment_reader.max_doc() { if segment_reader.is_deleted(doc_id) { continue; } let doc: TantivyDocument = store_reader.get(doc_id)?; let handle = doc .get_first(handle_field) .context("missing handle field in doc")? .as_bytes() .context("handle field is not bytes")? .to_vec(); handles.push(handle); } } handles.sort(); Ok(handles) } #[test] fn test_dedup_frags_retained_entries_two_overlapping_frags() -> Result<()> { let temp_dir = TempDir::new()?; let frag1_entries = vec![ RowEntry::new(b"a".to_vec(), false, 0), RowEntry::new(b"b".to_vec(), false, 0), RowEntry::new(b"c".to_vec(), false, 0), ]; let frag2_entries = vec![ RowEntry::new(b"a".to_vec(), false, 0), RowEntry::new(b"b".to_vec(), false, 0), RowEntry::new(b"d".to_vec(), false, 0), ]; let frag1 = create_test_frag(temp_dir.path(), &frag1_entries)?; let frag2 = create_test_frag(temp_dir.path(), &frag2_entries)?; let frags = vec![frag1, frag2]; let shard_range = ShardingKeyRange::half_open(b"a".to_vec(), b"z".to_vec())?; let mut retained = Vec::new(); dedup_frags( &shard_range, &frags, |e| { retained.push(e); Ok(()) }, temp_dir.path(), )?; assert_eq!( retained, vec![ RowEntry::new(b"a".to_vec(), false, 0), RowEntry::new(b"b".to_vec(), false, 0), RowEntry::new(b"c".to_vec(), false, 0), RowEntry::new(b"d".to_vec(), false, 0), ] ); Ok(()) } #[test] fn test_dedup_marks_obsolete_docs_in_bitsets_without_mutating_index() -> Result<()> { let temp_dir = TempDir::new()?; let frag1_entries = vec![ RowEntry::new(b"a".to_vec(), false, 0), RowEntry::new(b"b".to_vec(), false, 0), RowEntry::new(b"c".to_vec(), false, 0), ]; let frag2_entries = vec![ RowEntry::new(b"a".to_vec(), false, 0), RowEntry::new(b"b".to_vec(), false, 0), RowEntry::new(b"d".to_vec(), false, 0), ]; let frag1 = create_test_frag(temp_dir.path(), &frag1_entries)?; let frag2 = create_test_frag(temp_dir.path(), &frag2_entries)?; let frags = vec![frag1.clone(), frag2.clone()]; let shard_range = ShardingKeyRange::half_open(b"a".to_vec(), b"z".to_vec())?; let output = dedup_frags(&shard_range, &frags, |_e| Ok(()), temp_dir.path())?; let obsolete_doc_bitsets = output.obsolete_doc_bitsets; assert_eq!( read_handles_from_index(&frag1.local_frag_path)?, vec![b"a".to_vec(), b"b".to_vec(), b"c".to_vec()] ); assert_eq!( read_handles_from_index(&frag2.local_frag_path)?, vec![b"a".to_vec(), b"b".to_vec(), b"d".to_vec()] ); assert_eq!( read_handles_from_index_with_doc_bitset( &frag1.local_frag_path, &obsolete_doc_bitsets[0] )?, vec![b"c".to_vec()] ); assert_eq!( read_handles_from_index_with_doc_bitset( &frag2.local_frag_path, &obsolete_doc_bitsets[1] )?, vec![b"a".to_vec(), b"b".to_vec(), b"d".to_vec()] ); Ok(()) } #[test] fn test_dedup_no_obsolete_keys_returns_empty_bitsets() -> Result<()> { let temp_dir = TempDir::new()?; let frag1 = create_test_frag( temp_dir.path(), &[ RowEntry::new(b"a".to_vec(), false, 0), RowEntry::new(b"b".to_vec(), false, 0), ], )?; let frag2 = create_test_frag( temp_dir.path(), &[ RowEntry::new(b"c".to_vec(), false, 0), RowEntry::new(b"d".to_vec(), false, 0), ], )?; let frags = vec![frag1.clone(), frag2.clone()]; let shard_range = ShardingKeyRange::half_open(b"a".to_vec(), b"z".to_vec())?; let mut retained = Vec::new(); let output = dedup_frags( &shard_range, &frags, |entry| { retained.push(entry); Ok(()) }, temp_dir.path(), )?; let obsolete_doc_bitsets = output.obsolete_doc_bitsets; assert_eq!(output.dedup_stat.obsolete_by_dedup_count, 0); assert_eq!(retained.len(), 4); assert_eq!( output.dedup_stat.build_obsolete_doc_bitsets_time, Duration::ZERO ); assert!( obsolete_doc_bitsets .iter() .all(|frag| !frag.has_deleted_docs()) ); assert_eq!( read_handles_from_index_with_doc_bitset( &frag1.local_frag_path, &obsolete_doc_bitsets[0] )?, vec![b"a".to_vec(), b"b".to_vec()] ); assert_eq!( read_handles_from_index_with_doc_bitset( &frag2.local_frag_path, &obsolete_doc_bitsets[1] )?, vec![b"c".to_vec(), b"d".to_vec()] ); Ok(()) } }