use std::{ path::Path, sync::Arc, time::{Duration, Instant}, }; use anyhow::{Context, Result, anyhow, bail, ensure}; use log::{debug, error, info}; use tantivy::index::SegmentId; use tici_grpc::{CompactFragRequest, CompactType, FinishCompactFragRequest}; use uuid::Uuid; use crate::{ compactor::{ frag_utils::{FragsDedupStat, dedup_frags}, merge_utils::{RebuildIndexStat, rebuild_index}, }, config::WorkerNodeConfig, reader::frag_reader::FragReaderMeta, utils::{ dedup_bitset::FragAliveBitSetBytes, key::{ShardingKey, ShardingKeyRange}, metrics::begin_worker_compaction_metrics, row_entry::{RowEntry, RowEntryWriter}, s3_frag, schema::IndexMeta, shard_tracing_id, }, writer::{frag_batch_uploader::FragBatchUploader, tantivy_doc_buffer::HandleStat}, }; const ROW_ENTRY_WRITE_BATCH_SIZE: usize = 4096; // Summary of one retained-entry batch. This lets us estimate global handle // statistics without keeping all retained entries in memory. struct RetainGroupStat { // Global row offset of this group's first retained handle. offset: usize, min_handle: Vec, mid_handle: Vec, max_handle: Vec, } impl RetainGroupStat { fn from_entries(offset: usize, entries: &[RowEntry]) -> Self { debug_assert!(!entries.is_empty()); let rows = entries.len(); Self { offset, min_handle: entries[0].key.clone(), mid_handle: entries[rows / 2].key.clone(), max_handle: entries[rows - 1].key.clone(), } } } // Estimate a shard-level middle handle from per-batch summaries. // Use group offsets to locate the group containing `total_rows / 2` with // binary search, then compare that group's sampled indexes against the target // index, preferring `mid` and using `min/max` as fallbacks. // Rank-error bound: for a group with `r` rows, sampled candidate indexes are // `start`, `start + r / 2`, and `start + r - 1`, so the farthest distance from // any index in this group to its nearest candidate is at most `r / 4` (floor). // Therefore, compared with the exact median handle, this estimate's maximal // row-index error is `max_group_rows / 4`; with batch size 4096, the bound is // 1024 rows. fn estimate_mid_handle(group_stats: &[RetainGroupStat], total_rows: usize) -> Vec { if total_rows == 0 || group_stats.is_empty() { return vec![]; } let mid_index = total_rows / 2; let group_idx = group_stats .partition_point(|stat| stat.offset <= mid_index) .saturating_sub(1); let group = &group_stats[group_idx]; let group_start = group.offset; let group_end_exclusive = group_stats .get(group_idx + 1) .map_or(total_rows, |next| next.offset); assert!( group_end_exclusive > group_start, "invalid group range: start={}, end_exclusive={}", group_start, group_end_exclusive ); let group_rows = group_end_exclusive - group_start; let group_mid = group_start + group_rows / 2; let group_end = group_end_exclusive - 1; let mut best_distance = usize::MAX; let mut best_handle = None; for (idx, key) in [ (group_mid, &group.mid_handle), (group_start, &group.min_handle), (group_end, &group.max_handle), ] { let distance = idx.abs_diff(mid_index); if distance < best_distance { best_distance = distance; best_handle = Some(key.clone()); } } best_handle.unwrap_or_default() } #[derive(Default, Debug)] pub struct ShardCompactorStat { pub input_frag_count: usize, pub input_frag_size: usize, pub input_doc_count: usize, pub dedup_stat: FragsDedupStat, pub merge_deleted_set_time: Duration, pub rebuild_index_stat: RebuildIndexStat, pub calc_handle_stat_time: Duration, pub output_frag_size: usize, pub output_doc_count: usize, pub download_frags_ms: usize, pub upload_frag_ms: usize, } pub struct ShardCompactor { index_meta: IndexMeta, output_frag_path: String, input_frags: Vec, shard_range: ShardingKeyRange, compact_type: CompactType, index_num_threads: usize, index_mem_budget: usize, } impl ShardCompactor { pub fn new( index_meta: IndexMeta, output_frag_path: String, input_frags: Vec, shard_range: ShardingKeyRange, compact_type: CompactType, index_num_threads: usize, index_mem_budget: usize, ) -> Self { assert!(!input_frags.is_empty()); Self { index_meta, output_frag_path, input_frags, shard_range, compact_type, index_num_threads, index_mem_budget, } } pub async fn compact( self, stat: &mut ShardCompactorStat, ) -> Result<(Vec, HandleStat)> { s3_frag::prepare_local_frag_path(&self.output_frag_path).await?; // For Full compaction, filter out deleted entries (del_mark = true) // For Delta/Base compaction, retain all entries including deleted ones let should_retain = |e: &RowEntry| -> bool { match self.compact_type { CompactType::Full => !e.del_mark, CompactType::Base | CompactType::Delta => true, } }; let handle_set_path = crate::utils::frag_row_entry_path(&self.output_frag_path); let mut row_entry_writer = RowEntryWriter::create(&handle_set_path)?; let mut retain_entries_batch = Vec::with_capacity(ROW_ENTRY_WRITE_BATCH_SIZE); let mut retain_group_stats = Vec::new(); let mut retain_handle_count = 0; let mut retain_delete_count = 0; let obsolete_keys_temp_dir = Path::new(&self.output_frag_path).parent().ok_or_else(|| { anyhow!( "Failed to get parent of output frag path: {}", self.output_frag_path ) })?; let dedup_res = dedup_frags( &self.shard_range, &self.input_frags, |e| { if should_retain(&e) { retain_delete_count += e.del_mark as usize; retain_handle_count += 1; retain_entries_batch.push(e); if retain_entries_batch.len() >= ROW_ENTRY_WRITE_BATCH_SIZE { let group_offset = retain_handle_count - retain_entries_batch.len(); retain_group_stats.push(RetainGroupStat::from_entries( group_offset, &retain_entries_batch, )); row_entry_writer.write_batch(&retain_entries_batch)?; retain_entries_batch.clear(); } } Ok(()) }, obsolete_keys_temp_dir, )?; if !retain_entries_batch.is_empty() { let group_offset = retain_handle_count - retain_entries_batch.len(); retain_group_stats.push(RetainGroupStat::from_entries( group_offset, &retain_entries_batch, )); row_entry_writer.write_batch(&retain_entries_batch)?; } row_entry_writer.close()?; let obsolete_doc_bitsets = FragAliveBitSetBytes::from_doc_bitsets(dedup_res.obsolete_doc_bitsets)?; stat.dedup_stat = dedup_res.dedup_stat; stat.merge_deleted_set_time = stat.dedup_stat.merge_and_collect_obsolete_time; let calc_handle_stat_start = Instant::now(); let mid_handle = estimate_mid_handle(&retain_group_stats, retain_handle_count); let min_handle = retain_group_stats .first() .map_or(vec![], |group| group.min_handle.clone()); let max_handle = retain_group_stats .last() .map_or(vec![], |group| group.max_handle.clone()); let handle_stat = HandleStat { handle_count: retain_handle_count, delete_count: retain_delete_count, min_handle, mid_handle, max_handle, }; stat.calc_handle_stat_time = calc_handle_stat_start.elapsed(); // Copy data of old frags into new_frag_path and generate a new meta file. let input_frag_paths: Vec = self .input_frags .iter() .map(|frag| frag.local_frag_path.clone()) .collect(); let (seg_ids, rebuild_index_stat) = rebuild_index( &input_frag_paths, &self.index_meta, &self.output_frag_path, self.index_num_threads, self.index_mem_budget, obsolete_doc_bitsets, )?; stat.rebuild_index_stat = rebuild_index_stat; Ok((seg_ids, handle_stat)) } } pub async fn compact_frags( req: CompactFragRequest, config: WorkerNodeConfig, output_frag_path: String, frag_batch_uploader: Arc, ) -> Result { ensure!(!req.input_frags.is_empty(), "input_frags is empty"); let compact_download_metrics = begin_worker_compaction_metrics("compact_download"); let download_frags_time = Instant::now(); let compact_type = CompactType::try_from(req.compact_type) .with_context(|| format!("compact_type {} is invalid", req.compact_type))?; let s3_location = req .upload_location .as_ref() .context("S3 location is none")?; let shard_header = req.shard.as_ref().context("ShardManifestHeader is none")?; let shard_range = ShardingKeyRange::half_open(shard_header.start_key.clone(), shard_header.end_key.clone())?; let table_info = req.table_info.as_ref().context("TableInfo is none")?; let index_info = req.index_info.as_ref().context("IndexInfo is none")?; let index_meta = IndexMeta::new(table_info, index_info)?; let mut input_frag_metas = Vec::with_capacity(req.input_frags.len()); let mut input_frag_size = 0; let mut input_doc_count = 0; for frag_meta in &req.input_frags { let frag_property = frag_meta .property .as_ref() .context("FragMeta property is none")?; input_frag_size += frag_property.size; input_doc_count += frag_property.count; let start = &frag_property.min_handle; let end = ShardingKey::next_lexicographic(&frag_property.max_handle); let frag_range = ShardingKeyRange::half_open(start.clone(), end)?; if shard_range.is_intersect_with(&frag_range) { input_frag_metas.push(frag_meta.clone()); } else { info!( "Frag {:?}, frag_range={:?}, is not in the shard range {:?}, skip it", frag_meta, frag_range, shard_range ); } } // TODO: just let meta service remove these fragments from this shard? if input_frag_metas.is_empty() { bail!("No frags in the shard range={:?}", req); } let frag_base_path = config.storage.frag_base_path(); let backend_storage = frag_batch_uploader.backend_storage(); let mut input_frag_reader_metas = Vec::with_capacity(input_frag_metas.len()); // Different shards may use the same frag for compaction, so the frag needs to // be isolated. let compaction_local_data_path = format!("{}/{}", frag_base_path, Uuid::new_v4()); debug!("compaction_local_data_path={}", compaction_local_data_path); let download_result = async { for frag_meta in &input_frag_metas { let local_frag_path = s3_frag::download_frag( backend_storage.as_ref(), &s3_location.bucket, frag_meta, &compaction_local_data_path, ) .await?; let frag_reader_meta = FragReaderMeta::from_local_path(frag_meta, local_frag_path)?; input_frag_reader_metas.push(frag_reader_meta); } Ok::<(), anyhow::Error>(()) } .await; compact_download_metrics.mark_result(&download_result); download_result?; drop(compact_download_metrics); let compact_execute_metrics = begin_worker_compaction_metrics("compact_execute"); let mut compact_stat = ShardCompactorStat { input_frag_count: input_frag_reader_metas.len(), input_frag_size: input_frag_size as usize, input_doc_count: input_doc_count as usize, download_frags_ms: download_frags_time.elapsed().as_millis() as usize, ..Default::default() }; let compactor = ShardCompactor::new( index_meta, output_frag_path.clone(), input_frag_reader_metas, shard_range, compact_type, config.compaction.index_num_threads, config.compaction.index_mem_budget.0 as usize, ); let compact_execute_result = compactor.compact(&mut compact_stat).await; compact_execute_metrics.mark_result(&compact_execute_result); let (seg_ids, handle_stat) = compact_execute_result?; drop(compact_execute_metrics); let compact_upload_metrics = begin_worker_compaction_metrics("compact_upload"); let upload_frag_time = Instant::now(); let upload_frag_result = s3_frag::upload_frag( &frag_batch_uploader, &s3_location.bucket, &s3_location.prefix, &frag_base_path, &output_frag_path, handle_stat, &seg_ids, ) .await; compact_upload_metrics.mark_result(&upload_frag_result); let frag_meta = upload_frag_result?; let _ = std::fs::remove_dir_all(&output_frag_path) .inspect_err(|e| error!("Remove {} failed, err={:?}", output_frag_path, e)); // Input frags local data path let _ = std::fs::remove_dir_all(&compaction_local_data_path) .inspect_err(|e| error!("Remove {} failed, err={:?}", compaction_local_data_path, e)); let property = frag_meta.property.as_ref().unwrap(); // Must success. compact_stat.output_doc_count = property.count as usize; compact_stat.output_frag_size = property.size as usize; compact_stat.upload_frag_ms = upload_frag_time.elapsed().as_millis() as usize; drop(compact_upload_metrics); let finish_compact_frag_req = FinishCompactFragRequest { addr: config.server.get_advertise_addr().to_string(), shard: Some(shard_header.clone()), keyspace_id: req.keyspace_id, table_id: req.table_id, index_id: req.index_id, seq: req.seq, timestamp: req.timestamp, leader_term: req.leader_term.clone(), output_frags: vec![frag_meta], input_frags: req.input_frags.clone(), }; info!( "{} compaction finished, stat={:?}", shard_tracing_id(&req.get_index_key(), shard_header.shard_id), compact_stat, ); Ok(finish_compact_frag_req) } #[cfg(test)] mod tests { use std::collections::{BTreeMap, BTreeSet}; use anyhow::Context; use tantivy::{ Index, TantivyDocument, collector::DocSetCollector, query::AllQuery, schema::{Schema, Value}, }; use tempfile::TempDir; use tici_test_utils::test_utils::init_logger_for_ut; use tokio::sync::mpsc; use super::*; use crate::{ config::FragWriterConfig, test_utils::{ TEST_INDEX_META_OVERLAP, TEST_S3_BUCKET, TEST_SHARD_META, is_sorted_frag, test_gen_raw_docs, test_get_end_handle, test_get_frag_batch_uploader, test_get_inf_range, test_get_start_handle, test_new_frag_path_generator, }, utils::{ self, row_entry::RowEntry, schema::{HANDLE_STORE_NAME, VERSION_STORE_NAME}, tantivy_doc::RawDoc, }, writer::cdc_log_frag_writer::CDCLogFragWriter, }; fn get_version(doc: &TantivyDocument, schema: &Schema) -> u64 { let field = schema.get_field(VERSION_STORE_NAME).unwrap(); let value = doc.get_first(field).unwrap(); value.as_u64().unwrap() } fn get_handle(doc: &TantivyDocument, schema: &Schema) -> ShardingKey { let field = schema.get_field(HANDLE_STORE_NAME).unwrap(); let value = doc.get_first(field).unwrap(); ShardingKey::from_raw(value.as_bytes().unwrap().to_vec()) } fn get_handle_versions(frag_path: &str) -> Vec<(ShardingKey, u64)> { let index = Index::open_in_dir(utils::frag_index_path(frag_path)).unwrap(); let schema = index.schema(); let searcher = index.reader().unwrap().searcher(); let docs = searcher.search(&AllQuery, &DocSetCollector).unwrap(); let mut handle_versions = Vec::new(); for addr in docs { let doc: TantivyDocument = searcher.doc(addr).unwrap(); let h = get_handle(&doc, &schema); let v = get_version(&doc, &schema); handle_versions.push((h, v)); } handle_versions } fn copy_dir_all(src: impl AsRef, dst: impl AsRef) { std::fs::create_dir_all(&dst).unwrap(); for entry in std::fs::read_dir(src).unwrap() { let entry = entry.unwrap(); let ty = entry.file_type().unwrap(); if ty.is_dir() { copy_dir_all(entry.path(), dst.as_ref().join(entry.file_name())); } else { std::fs::copy(entry.path(), dst.as_ref().join(entry.file_name())).unwrap(); } } } fn copy_frags( parent_path: &str, frag_reader_metas: &[FragReaderMeta], new_parent_path: &str, ) -> Vec { copy_dir_all(parent_path, new_parent_path); let mut new_frag_reader_metas = Vec::with_capacity(frag_reader_metas.len()); for meta in frag_reader_metas { let mut new_meta = meta.clone(); new_meta.local_frag_path = meta.local_frag_path.replace(parent_path, new_parent_path); new_frag_reader_metas.push(new_meta); } new_frag_reader_metas } fn assert_handle_stat_matches_expected( handle_stat: &HandleStat, expected_handle_set: &BTreeSet, ) { assert_eq!(handle_stat.handle_count, expected_handle_set.len()); if expected_handle_set.is_empty() { assert!(handle_stat.min_handle.is_empty()); assert!(handle_stat.mid_handle.is_empty()); assert!(handle_stat.max_handle.is_empty()); return; } let expected_min = expected_handle_set.iter().next().unwrap().raw(); let expected_max = expected_handle_set.iter().last().unwrap().raw(); assert_eq!(handle_stat.min_handle, expected_min); assert_eq!(handle_stat.max_handle, expected_max); let mid_handle = ShardingKey::from_raw(handle_stat.mid_handle.clone()); assert!( expected_handle_set.contains(&mid_handle), "mid_handle must be one of retained handles" ); let expected_mid_rank = handle_stat.handle_count / 2; let actual_mid_rank = expected_handle_set .iter() .position(|h| h.raw() == handle_stat.mid_handle.as_slice()) .expect("mid_handle should exist in expected_handle_set"); if handle_stat.handle_count <= ROW_ENTRY_WRITE_BATCH_SIZE { assert_eq!( actual_mid_rank, expected_mid_rank, "single-group mid_handle should be exact" ); } else { // Candidates are sampled as {min, mid, max} per retain group. let max_rank_error = ROW_ENTRY_WRITE_BATCH_SIZE / 4; assert!( actual_mid_rank.abs_diff(expected_mid_rank) <= max_rank_error, "mid_handle rank too far from median: actual_rank={}, expected_rank={}, max_error={}", actual_mid_rank, expected_mid_rank, max_rank_error ); } } #[test] fn test_estimate_mid_handle_uses_offset_binary_search() { let middle_group = vec![ RetainGroupStat { offset: 0, min_handle: vec![0], mid_handle: vec![2], max_handle: vec![3], }, RetainGroupStat { offset: 4, min_handle: vec![4], mid_handle: vec![5], max_handle: vec![6], }, RetainGroupStat { offset: 7, min_handle: vec![7], mid_handle: vec![7], max_handle: vec![7], }, ]; // total_rows=8 => mid_index=4, so the second group is selected and // `min_handle` is closest to the global mid. assert_eq!(estimate_mid_handle(&middle_group, 8), vec![4]); let last_group = vec![ RetainGroupStat { offset: 0, min_handle: vec![0], mid_handle: vec![1], max_handle: vec![1], }, RetainGroupStat { offset: 2, min_handle: vec![2], mid_handle: vec![2], max_handle: vec![2], }, RetainGroupStat { offset: 3, min_handle: vec![3], mid_handle: vec![4], max_handle: vec![5], }, ]; // total_rows=6 => mid_index=3, so binary search should select the last // group and return its boundary candidate. assert_eq!(estimate_mid_handle(&last_group, 6), vec![3]); } async fn compact_and_verify( index_meta: &IndexMeta, input_frags: Vec, shard_start_key: Vec, shard_end_key: Vec, compact_type: CompactType, expected_raw_docs: BTreeMap, expected_handle_set: BTreeSet, ) -> HandleStat { let temp_dir = TempDir::new().unwrap(); let local_data_path = temp_dir.path().to_str().unwrap(); let frag_path_gen = test_new_frag_path_generator(local_data_path); let output_frag_path = frag_path_gen .new_path(index_meta.index_key()) .await .unwrap(); let mut compact_stat = ShardCompactorStat { input_frag_count: input_frags.len(), ..Default::default() }; let frag_writer_config = FragWriterConfig::default(); // Ignore input_frag_size, input_doc_count here. let compactor = ShardCompactor::new( index_meta.clone(), output_frag_path.to_string(), input_frags, ShardingKeyRange::half_open(shard_start_key, shard_end_key).unwrap(), compact_type, frag_writer_config.index_num_threads, frag_writer_config.index_mem_budget.0 as usize, ); let (seg_ids, handle_stat) = compactor.compact(&mut compact_stat).await.unwrap(); assert!(is_sorted_frag(&output_frag_path, index_meta)); debug!("CompactStat={:?}", compact_stat); let frag_batch_uploader = test_get_frag_batch_uploader().await.unwrap(); let frag_meta = s3_frag::upload_frag( &frag_batch_uploader, TEST_S3_BUCKET, output_frag_path.strip_prefix("/").unwrap(), local_data_path, &output_frag_path, handle_stat.clone(), &seg_ids, ) .await .unwrap(); assert_eq!(frag_meta.seg_count(), seg_ids.len()); assert!(!frag_meta.storage_refs.is_empty()); assert_handle_stat_matches_expected(&handle_stat, &expected_handle_set); let mut row_entries = RowEntry::read_all(utils::frag_row_entry_path(&output_frag_path)).unwrap(); row_entries.sort_by(|left, right| left.key.cmp(&right.key)); assert_eq!(row_entries.len(), expected_handle_set.len()); for entry in &row_entries { let handle = ShardingKey::from_raw(entry.key.clone()); assert!(expected_handle_set.contains(&handle)); if let Some(raw_doc) = expected_raw_docs.get(&handle) { assert!(!entry.del_mark); assert_eq!(entry.version, raw_doc.version); } else { assert!(entry.del_mark); } } let mut handle_versions = get_handle_versions(&output_frag_path); handle_versions.sort_by_key(|hv| hv.0.clone()); assert_eq!(handle_versions.len(), expected_raw_docs.len()); let has_dup = handle_versions.windows(2).any(|hv| hv[0].0 == hv[1].0); assert!(!has_dup); for (h, v) in &handle_versions { assert_eq!(v, &expected_raw_docs[h].version); } let temp_dir = TempDir::new().unwrap(); let download_local_path = temp_dir.path().to_str().unwrap(); let backend_storage = frag_batch_uploader.backend_storage(); let download_frag_path = s3_frag::download_frag( backend_storage.as_ref(), TEST_S3_BUCKET, &frag_meta, download_local_path, ) .await .unwrap(); let mut download_handle_versions = get_handle_versions(&download_frag_path); download_handle_versions.sort_by_key(|hv| hv.0.clone()); assert_eq!(download_handle_versions.len(), handle_versions.len()); assert_eq!(download_handle_versions, handle_versions); handle_stat } #[tokio::test] async fn test_shard_compactor() { init_logger_for_ut(); let temp_dir = TempDir::new().unwrap(); let local_data_path = temp_dir.path().to_str().unwrap(); let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1); let (frag_writer_sender, _) = mpsc::channel(1); drop(cdc_log_sender); let index_meta = TEST_INDEX_META_OVERLAP.clone(); let mut cdc_log_frag_writer = CDCLogFragWriter::new( &FragWriterConfig::default(), TEST_SHARD_META.clone(), index_meta.clone(), test_new_frag_path_generator(local_data_path), cdc_log_receiver, frag_writer_sender, ) .await .unwrap(); let inserts = 10000; let updates = 3000; let deletes = 1000; let (raw_docs, expected_raw_docs) = test_gen_raw_docs(inserts, updates, deletes); let handle_set = raw_docs .iter() .map(|d| d.handle.clone()) .collect::>(); let frag_count = 10; let docs_per_frag = raw_docs.len() / frag_count; let mut frag_reader_metas = Vec::with_capacity(frag_count); let mut current_docs = 0; for raw_doc in raw_docs { cdc_log_frag_writer.add_raw_doc(raw_doc).unwrap(); current_docs += 1; if current_docs >= docs_per_frag { let mut frag_writer = cdc_log_frag_writer.switch_frag_writer().await.unwrap(); let frag_res = frag_writer.finalize().unwrap(); frag_reader_metas.push( FragReaderMeta::new( &frag_writer.frag_path(), frag_res.handle_stat.min_handle, frag_res.handle_stat.max_handle, ) .unwrap(), ); current_docs = 0; assert!(is_sorted_frag( &frag_reader_metas.iter().last().unwrap().local_frag_path, &index_meta )); } } assert_eq!(frag_reader_metas.len(), frag_count); let handle_stat = { let temp_dir = TempDir::new().unwrap(); let new_path = temp_dir.path().to_str().unwrap(); let frag_infos = copy_frags(local_data_path, &frag_reader_metas, new_path); compact_and_verify( &index_meta, frag_infos, test_get_start_handle().raw().to_vec(), test_get_end_handle().raw().to_vec(), CompactType::Delta, expected_raw_docs.clone(), handle_set.clone(), ) .await }; { let temp_dir = TempDir::new().unwrap(); let new_path = temp_dir.path().to_str().unwrap(); let frag_infos = copy_frags(local_data_path, &frag_reader_metas, new_path); let left_start = ShardingKey::from_raw(handle_stat.min_handle.clone()); let left_end = ShardingKey::from_raw(handle_stat.mid_handle.clone()); let expected_raw_docs_left: BTreeMap<_, _> = expected_raw_docs .range(&left_start..&left_end) .map(|(k, v)| (k.clone(), v.clone())) .collect(); let expected_handle_set_left = handle_set.range(&left_start..&left_end).cloned().collect(); compact_and_verify( &index_meta, frag_infos, handle_stat.min_handle.clone(), handle_stat.mid_handle.clone(), CompactType::Delta, expected_raw_docs_left, expected_handle_set_left, ) .await; } { let temp_dir = TempDir::new().unwrap(); let new_path = temp_dir.path().to_str().unwrap(); let frag_infos = copy_frags(local_data_path, &frag_reader_metas, new_path); let right_start = ShardingKey::from_raw(handle_stat.mid_handle.clone()); let right_end = ShardingKey::from_raw(handle_stat.max_handle.clone()); let expected_raw_docs_right: BTreeMap<_, _> = expected_raw_docs .range(right_start.clone()..right_end.clone()) .map(|(k, v)| (k.clone(), v.clone())) .collect(); let expected_handle_set_right = handle_set.range(right_start..right_end).cloned().collect(); compact_and_verify( &index_meta, frag_infos, handle_stat.mid_handle.clone(), handle_stat.max_handle.clone(), CompactType::Delta, expected_raw_docs_right, expected_handle_set_right, ) .await; } } #[tokio::test] async fn test_deduplication() { init_logger_for_ut(); let temp_dir = TempDir::new().unwrap(); let local_data_path = temp_dir.path().to_str().unwrap(); let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1); let (frag_writer_sender, _) = mpsc::channel(1); drop(cdc_log_sender); let index_meta = TEST_INDEX_META_OVERLAP.clone(); let mut cdc_log_frag_writer = CDCLogFragWriter::new( &FragWriterConfig::default(), TEST_SHARD_META.clone(), index_meta.clone(), test_new_frag_path_generator(local_data_path), cdc_log_receiver, frag_writer_sender, ) .await .unwrap(); let (raw_docs, _expected_raw_docs) = test_gen_raw_docs(1000, 0, 0); let frag_count = 10; let mut frag_reader_metas = Vec::with_capacity(frag_count); for _i in 0..frag_count { for raw_doc in &raw_docs { cdc_log_frag_writer.add_raw_doc(raw_doc.clone()).unwrap(); } let mut frag_writer = cdc_log_frag_writer.switch_frag_writer().await.unwrap(); let frag_res = frag_writer.finalize().unwrap(); frag_reader_metas.push( FragReaderMeta::new( &frag_writer.frag_path(), frag_res.handle_stat.min_handle, frag_res.handle_stat.max_handle, ) .unwrap(), ); assert!(is_sorted_frag( &frag_reader_metas.iter().last().unwrap().local_frag_path, &index_meta )); } assert_eq!(frag_reader_metas.len(), frag_count); let temp_dir = TempDir::new().unwrap(); let output_frag_path = temp_dir.path().to_str().unwrap(); let mut compact_stat = ShardCompactorStat { input_frag_count: frag_reader_metas.len(), ..Default::default() }; let frag_writer_config = FragWriterConfig::default(); let compactor = ShardCompactor::new( index_meta.clone(), output_frag_path.to_string(), frag_reader_metas, test_get_inf_range(), CompactType::Delta, frag_writer_config.index_num_threads, frag_writer_config.index_mem_budget.0 as usize, ); let (_seg_ids, _handle_stat) = compactor.compact(&mut compact_stat).await.unwrap(); assert!(is_sorted_frag(output_frag_path, &index_meta)); let index = Index::open_in_dir(crate::utils::frag_index_path(output_frag_path)).unwrap(); let docs = index .reader() .unwrap() .searcher() .search(&AllQuery, &DocSetCollector) .unwrap(); assert_eq!(docs.len(), 1000); assert_eq!(compact_stat.dedup_stat.input_key_count, 1000 * frag_count); assert_eq!( compact_stat.dedup_stat.in_range_key_count, 1000 * frag_count ); assert_eq!(compact_stat.dedup_stat.retained_key_count, 1000); assert_eq!(compact_stat.dedup_stat.retained_delete_mark_count, 0); assert_eq!( compact_stat.dedup_stat.obsolete_by_dedup_count, 1000 * (frag_count - 1) ); assert_eq!(compact_stat.dedup_stat.obsolete_by_range_count, 0); assert_eq!( compact_stat.merge_deleted_set_time, compact_stat.dedup_stat.merge_and_collect_obsolete_time ); assert!(compact_stat.calc_handle_stat_time.as_nanos() > 0); } #[tokio::test] async fn test_full_compaction_filters_deleted_handles() { init_logger_for_ut(); let temp_dir = TempDir::new().unwrap(); let local_data_path = temp_dir.path().to_str().unwrap(); let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1); let (frag_writer_sender, _) = mpsc::channel(1); drop(cdc_log_sender); let index_meta = TEST_INDEX_META_OVERLAP.clone(); let mut cdc_log_frag_writer = CDCLogFragWriter::new( &FragWriterConfig::default(), TEST_SHARD_META.clone(), index_meta.clone(), test_new_frag_path_generator(local_data_path), cdc_log_receiver, frag_writer_sender, ) .await .unwrap(); // Generate test data: 1000 inserts, 300 updates, 100 deletes let inserts = 1000; let updates = 300; let deletes = 100; let (raw_docs, expected_raw_docs) = test_gen_raw_docs(inserts, updates, deletes); // For Full compaction, we should filter out deleted handles // So expected docs should exclude the deleted ones let expected_raw_docs_full: BTreeMap<_, _> = expected_raw_docs .iter() .filter(|(_, doc)| doc.op_type != crate::writer::cdc_log::OperationType::Delete) .map(|(k, v)| (k.clone(), v.clone())) .collect(); let expected_handle_set_full: BTreeSet<_> = expected_raw_docs_full.keys().cloned().collect(); // For Delta compaction, keep all handles let handle_set_delta = raw_docs .iter() .map(|d| d.handle.clone()) .collect::>(); let frag_count = 5; let docs_per_frag = raw_docs.len() / frag_count; let mut frag_reader_metas = Vec::with_capacity(frag_count); let mut current_docs = 0; for raw_doc in raw_docs { cdc_log_frag_writer.add_raw_doc(raw_doc).unwrap(); current_docs += 1; if current_docs >= docs_per_frag { let mut frag_writer = cdc_log_frag_writer.switch_frag_writer().await.unwrap(); let frag_res = frag_writer.finalize().unwrap(); frag_reader_metas.push( FragReaderMeta::new( &frag_writer.frag_path(), frag_res.handle_stat.min_handle, frag_res.handle_stat.max_handle, ) .unwrap(), ); current_docs = 0; assert!(is_sorted_frag( &frag_reader_metas.iter().last().unwrap().local_frag_path, &index_meta )); } } assert_eq!(frag_reader_metas.len(), frag_count); // Test Delta compaction - should keep all handles including deleted ones { let temp_dir = TempDir::new().unwrap(); let new_path = temp_dir.path().to_str().unwrap(); let frag_infos = copy_frags(local_data_path, &frag_reader_metas, new_path); compact_and_verify( &index_meta, frag_infos, test_get_start_handle().raw().to_vec(), test_get_end_handle().raw().to_vec(), CompactType::Delta, expected_raw_docs.clone(), handle_set_delta.clone(), ) .await; } // Test Full compaction - should filter out deleted handles { let temp_dir = TempDir::new().unwrap(); let new_path = temp_dir.path().to_str().unwrap(); let frag_infos = copy_frags(local_data_path, &frag_reader_metas, new_path); compact_and_verify( &index_meta, frag_infos, test_get_start_handle().raw().to_vec(), test_get_end_handle().raw().to_vec(), CompactType::Full, expected_raw_docs_full.clone(), expected_handle_set_full.clone(), ) .await; } } #[tokio::test] async fn test_delete_all_full_compaction() -> Result<()> { init_logger_for_ut(); let temp_dir = TempDir::new()?; let local_data_path = temp_dir.path().to_str().context("temp dir path to_str")?; let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1); let (frag_writer_sender, _) = mpsc::channel(1); drop(cdc_log_sender); let index_meta = TEST_INDEX_META_OVERLAP.clone(); let mut cdc_log_frag_writer = CDCLogFragWriter::new( &FragWriterConfig::default(), TEST_SHARD_META.clone(), index_meta.clone(), test_new_frag_path_generator(local_data_path), cdc_log_receiver, frag_writer_sender, ) .await?; let inserts = 1000; let updates = 0; let deletes = 1000; let (raw_docs, expected_raw_docs) = test_gen_raw_docs(inserts, updates, deletes); assert!(expected_raw_docs.is_empty()); let frag_count = 5; let docs_per_frag = raw_docs.len() / frag_count; let mut frag_reader_metas = Vec::with_capacity(frag_count); let mut current_docs = 0; for raw_doc in raw_docs { cdc_log_frag_writer.add_raw_doc(raw_doc)?; current_docs += 1; if current_docs >= docs_per_frag { let mut frag_writer = cdc_log_frag_writer.switch_frag_writer().await?; let frag_res = frag_writer.finalize()?; frag_reader_metas.push(FragReaderMeta::new( &frag_writer.frag_path(), frag_res.handle_stat.min_handle, frag_res.handle_stat.max_handle, )?); current_docs = 0; } } assert_eq!(frag_reader_metas.len(), frag_count); let temp_dir2 = TempDir::new()?; let new_path = temp_dir2.path().to_str().context("temp dir 2 path")?; let frag_infos = copy_frags(local_data_path, &frag_reader_metas, new_path); let frag_path_gen = test_new_frag_path_generator(new_path); let output_frag_path = frag_path_gen.new_path(index_meta.index_key()).await?; let mut compact_stat = ShardCompactorStat { input_frag_count: frag_infos.len(), ..Default::default() }; let frag_writer_config = FragWriterConfig::default(); let shard_range = ShardingKeyRange::half_open( test_get_start_handle().raw().to_vec(), test_get_end_handle().raw().to_vec(), )?; let compactor = ShardCompactor::new( index_meta.clone(), output_frag_path.to_string(), frag_infos, shard_range, CompactType::Full, frag_writer_config.index_num_threads, frag_writer_config.index_mem_budget.0 as usize, ); let (seg_ids, handle_stat) = compactor.compact(&mut compact_stat).await?; assert_eq!(handle_stat.handle_count, 0); assert!(handle_stat.min_handle.is_empty()); assert!(handle_stat.mid_handle.is_empty()); assert!(handle_stat.max_handle.is_empty()); assert!(handle_stat.delete_count == 0); let handle_set_path = utils::frag_row_entry_path(&output_frag_path); let entries = RowEntry::read_all(&handle_set_path)?; assert!(entries.is_empty()); assert!(seg_ids.is_empty()); let handle_verions = get_handle_versions(&output_frag_path); assert!(handle_verions.is_empty()); Ok(()) } #[tokio::test] async fn test_delete_all_delta_compaction() -> Result<()> { init_logger_for_ut(); let temp_dir = TempDir::new()?; let local_data_path = temp_dir.path().to_str().context("temp dir path to_str")?; let (cdc_log_sender, cdc_log_receiver) = mpsc::channel(1); let (frag_writer_sender, _) = mpsc::channel(1); drop(cdc_log_sender); let index_meta = TEST_INDEX_META_OVERLAP.clone(); let mut cdc_log_frag_writer = CDCLogFragWriter::new( &FragWriterConfig::default(), TEST_SHARD_META.clone(), index_meta.clone(), test_new_frag_path_generator(local_data_path), cdc_log_receiver, frag_writer_sender, ) .await?; let inserts = 1000; let updates = 0; let deletes = 1000; let (raw_docs, expected_raw_docs) = test_gen_raw_docs(inserts, updates, deletes); assert!(expected_raw_docs.is_empty()); let expected_delete_versions: BTreeMap, u64> = raw_docs .iter() .fold(BTreeMap::new(), |mut versions, raw_doc| { versions .entry(raw_doc.handle.raw().to_vec()) .and_modify(|version| *version = (*version).max(raw_doc.version)) .or_insert(raw_doc.version); versions }); let handles: BTreeSet> = raw_docs.iter().map(|r| r.handle.raw().to_vec()).collect(); assert_eq!(handles.len(), inserts); let frag_count = 5; let docs_per_frag = raw_docs.len() / frag_count; let mut frag_reader_metas = Vec::with_capacity(frag_count); let mut current_docs = 0; for raw_doc in raw_docs { cdc_log_frag_writer.add_raw_doc(raw_doc)?; current_docs += 1; if current_docs >= docs_per_frag { let mut frag_writer = cdc_log_frag_writer.switch_frag_writer().await?; let frag_res = frag_writer.finalize()?; frag_reader_metas.push(FragReaderMeta::new( &frag_writer.frag_path(), frag_res.handle_stat.min_handle, frag_res.handle_stat.max_handle, )?); current_docs = 0; } } assert_eq!(frag_reader_metas.len(), frag_count); let temp_dir2 = TempDir::new()?; let new_path = temp_dir2.path().to_str().context("temp dir 2 path")?; let frag_infos = copy_frags(local_data_path, &frag_reader_metas, new_path); let frag_path_gen = test_new_frag_path_generator(new_path); let output_frag_path = frag_path_gen.new_path(index_meta.index_key()).await?; let mut compact_stat = ShardCompactorStat { input_frag_count: frag_infos.len(), ..Default::default() }; let frag_writer_config = FragWriterConfig::default(); let shard_range = ShardingKeyRange::half_open( test_get_start_handle().raw().to_vec(), test_get_end_handle().raw().to_vec(), )?; let compactor = ShardCompactor::new( index_meta.clone(), output_frag_path.to_string(), frag_infos, shard_range, CompactType::Delta, frag_writer_config.index_num_threads, frag_writer_config.index_mem_budget.0 as usize, ); let (seg_ids, handle_stat) = compactor.compact(&mut compact_stat).await?; let expected_handle_set = handles .iter() .cloned() .map(ShardingKey::from_raw) .collect::>(); assert_handle_stat_matches_expected(&handle_stat, &expected_handle_set); assert!(handle_stat.delete_count == handles.len()); let handle_set_path = utils::frag_row_entry_path(&output_frag_path); let entries = RowEntry::read_all(&handle_set_path)?; assert_eq!(entries.len(), handles.len()); for (i, h) in handles.iter().enumerate() { assert_eq!(entries[i].key, *h); assert!(entries[i].del_mark); assert_eq!(entries[i].version, expected_delete_versions[h]); } assert!(seg_ids.is_empty()); let handle_verions = get_handle_versions(&output_frag_path); assert!(handle_verions.is_empty()); Ok(()) } }