1mod chunk_maps;
4mod compact;
5mod compact_vectors;
6mod copy;
7pub(crate) use copy::append_and_delete_temp;
8pub(super) use copy::copy_local_range_or_bytes;
9mod dense;
10mod fast_fields;
11mod postings;
12pub(crate) use postings::PostingMergeStats;
13mod sparse;
14mod store;
15mod terms;
16pub(crate) use terms::MergedTerms;
17
18pub(crate) use dense::AnnWriteMode;
19
20use std::sync::Arc;
21use std::sync::atomic::{AtomicBool, Ordering};
22
23use rustc_hash::FxHashMap;
24
25use super::OffsetWriter;
26use super::reader::SegmentReader;
27use super::types::{FieldStats, SegmentFiles, SegmentId, SegmentMeta};
28use crate::Result;
29use crate::directories::{Directory, DirectoryWriter};
30use crate::dsl::{FieldType, Schema};
31use crate::index::{ReorderConcurrencyGate, ReorderPriority};
32use crate::structures::SparseFormat;
33
34fn doc_offsets(segments: &[SegmentReader]) -> Result<Vec<u32>> {
38 let mut offsets = Vec::with_capacity(segments.len());
39 let mut acc = 0u32;
40 for seg in segments {
41 offsets.push(acc);
42 acc = acc.checked_add(seg.num_docs()).ok_or_else(|| {
43 crate::Error::Internal(format!(
44 "Total document count across segments exceeds u32::MAX ({})",
45 u32::MAX
46 ))
47 })?;
48 }
49 Ok(offsets)
50}
51
52#[derive(Clone, Copy, Debug, Default)]
57struct MergeCapacity(u64);
58
59impl MergeCapacity {
60 #[inline]
61 fn add(&mut self, count: u64) -> Option<u64> {
62 self.0 = self.0.saturating_add(count);
63 (self.0 > u64::from(u32::MAX)).then_some(self.0)
64 }
65}
66
67fn field_capacity_error(
68 field_id: u32,
69 field_name: &str,
70 value_kind: &str,
71 count: u64,
72) -> crate::Error {
73 crate::Error::Schema(format!(
74 "merge would produce {count} {value_kind} for field {field_id} ('{field_name}'), \
75 exceeding the segment format limit {}; lower max_segment_docs for this \
76 multi-valued field",
77 u32::MAX,
78 ))
79}
80
81#[derive(Debug, Clone, Default)]
83pub struct MergeStats {
84 pub terms_processed: usize,
86 pub term_dict_bytes: usize,
88 pub postings_bytes: usize,
90 pub store_bytes: usize,
92 pub vectors_bytes: usize,
94 pub sparse_bytes: usize,
96 pub bp_converged: bool,
102 pub fast_bytes: usize,
104 pub posting_blocks_without_bounds: usize,
110}
111
112impl std::fmt::Display for MergeStats {
113 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
114 write!(
115 f,
116 "terms={}, term_dict={}, postings={}, store={}, dense_vectors={}, sparse_vectors={}, fast_fields={}, posting_blocks_without_bounds={}",
117 self.terms_processed,
118 crate::format_bytes(self.term_dict_bytes as u64),
119 crate::format_bytes(self.postings_bytes as u64),
120 crate::format_bytes(self.store_bytes as u64),
121 crate::format_bytes(self.vectors_bytes as u64),
122 crate::format_bytes(self.sparse_bytes as u64),
123 crate::format_bytes(self.fast_bytes as u64),
124 self.posting_blocks_without_bounds,
125 )
126 }
127}
128
129pub use super::types::TrainedVectorStructures;
131
132pub(crate) fn block_in_place_if_multithread<R>(f: impl FnOnce() -> R) -> R {
136 if tokio::runtime::Handle::try_current()
137 .map(|h| h.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread)
138 .unwrap_or(false)
139 {
140 tokio::task::block_in_place(f)
141 } else {
142 f()
143 }
144}
145
146pub struct SegmentMerger {
148 schema: Arc<Schema>,
149 optimization: crate::structures::IndexOptimization,
151 posting_codec: crate::structures::PostingCodec,
155 term_dict_block_size: crate::structures::SSTableBlockSize,
156 reorder_fields: bool,
159 background_pool: Option<Arc<rayon::ThreadPool>>,
163 granularity: crate::segment::reorder::BpGranularity,
167 bp_budget: crate::segment::BpBudget,
172 cancellation: Option<Arc<AtomicBool>>,
175 bp_memory_budget: usize,
177 reorder_permits: Option<Arc<ReorderConcurrencyGate>>,
180 reorder_priority: ReorderPriority,
183}
184
185impl SegmentMerger {
186 pub fn new(schema: Arc<Schema>) -> Self {
187 Self {
188 schema,
189 optimization: crate::structures::IndexOptimization::default(),
190 posting_codec: crate::structures::PostingCodec::default(),
191 term_dict_block_size: crate::structures::SSTableBlockSize::default(),
192 reorder_fields: false,
193 background_pool: None,
194 granularity: crate::segment::reorder::BpGranularity::Auto,
195 bp_budget: crate::segment::BpBudget::full(),
196 cancellation: None,
197 bp_memory_budget: crate::segment::reorder::DEFAULT_MEMORY_BUDGET,
198 reorder_permits: None,
199 reorder_priority: ReorderPriority::AutomaticMerge,
200 }
201 }
202
203 pub fn with_term_dict_block_size(mut self, size: crate::structures::SSTableBlockSize) -> Self {
205 self.term_dict_block_size = size;
206 self
207 }
208
209 pub fn with_posting_config(
211 mut self,
212 optimization: crate::structures::IndexOptimization,
213 posting_codec: crate::structures::PostingCodec,
214 ) -> Self {
215 self.optimization = optimization;
216 self.posting_codec = posting_codec;
217 self
218 }
219
220 pub fn with_reorder_fields(mut self, reorder: bool) -> Self {
223 self.reorder_fields = reorder;
224 self
225 }
226
227 pub fn with_background_pool(mut self, pool: Option<Arc<rayon::ThreadPool>>) -> Self {
229 self.background_pool = pool;
230 self
231 }
232
233 pub fn with_granularity(mut self, granularity: crate::segment::reorder::BpGranularity) -> Self {
235 self.granularity = granularity;
236 self
237 }
238
239 pub fn with_bp_budget(mut self, budget: crate::segment::BpBudget) -> Self {
241 self.bp_budget = budget;
242 self
243 }
244
245 pub(crate) fn with_cancellation(mut self, cancellation: Arc<AtomicBool>) -> Self {
246 self.cancellation = Some(cancellation);
247 self
248 }
249
250 pub fn with_bp_memory_budget(mut self, bytes: usize) -> Self {
252 self.bp_memory_budget = bytes;
253 self
254 }
255
256 pub fn with_reorder_permits(mut self, permits: Arc<ReorderConcurrencyGate>) -> Self {
258 self.reorder_permits = Some(permits);
259 self
260 }
261
262 pub(crate) fn with_reorder_priority(mut self, priority: ReorderPriority) -> Self {
263 self.reorder_priority = priority;
264 self
265 }
266
267 async fn acquire_reorder_permit(&self) -> Result<Option<crate::index::ReorderPermit>> {
268 self.ensure_not_cancelled()?;
269 match &self.reorder_permits {
270 Some(gate) => Ok(Some(gate.acquire(self.reorder_priority).await.map_err(
271 |_| crate::Error::Internal("background reorder scheduler is closed".into()),
272 )?)),
273 None => Ok(None),
274 }
275 }
276
277 pub(super) fn ensure_not_cancelled(&self) -> Result<()> {
278 if self
279 .cancellation
280 .as_ref()
281 .is_some_and(|cancelled| cancelled.load(Ordering::Acquire))
282 {
283 Err(crate::Error::IndexClosed)
284 } else {
285 Ok(())
286 }
287 }
288
289 fn validate_merge_capacities(&self, segments: &[SegmentReader]) -> Result<()> {
293 let mut maxscore_skip_entries = MergeCapacity::default();
295
296 for (field, entry) in self.schema.fields() {
297 match entry.field_type {
298 FieldType::DenseVector | FieldType::BinaryDenseVector => {
299 let mut vectors = MergeCapacity::default();
300 for segment in segments {
301 let Some(flat) = segment.flat_vectors().get(&field.0) else {
302 continue;
303 };
304 if let Some(total) = vectors.add(flat.num_vectors as u64) {
305 let value_kind = if entry.field_type == FieldType::BinaryDenseVector {
306 "binary vectors"
307 } else {
308 "dense vectors"
309 };
310 return Err(field_capacity_error(
311 field.0,
312 &entry.name,
313 value_kind,
314 total,
315 ));
316 }
317 }
318 }
319 FieldType::SparseVector => {
320 let format = entry
321 .sparse_vector_config
322 .as_ref()
323 .map(|config| config.format)
324 .unwrap_or_default();
325 match format {
326 SparseFormat::Seismic => {
327 let mut vectors = MergeCapacity::default();
328 for segment in segments {
329 if let Some(index) = segment.seismic_index(field)
330 && let Some(total) =
331 vectors.add(u64::from(index.total_vectors()))
332 {
333 return Err(field_capacity_error(
334 field.0,
335 &entry.name,
336 "Seismic vectors",
337 total,
338 ));
339 }
340 }
341 }
342
343 SparseFormat::Bmp => {
344 let mut vectors = MergeCapacity::default();
345 let mut blocks = MergeCapacity::default();
346 let mut real_slots = MergeCapacity::default();
347 let mut virtual_slots = MergeCapacity::default();
348
349 for segment in segments {
350 let Some(index) = segment.bmp_indexes().get(&field.0) else {
351 continue;
352 };
353 for (capacity, count, value_kind) in [
354 (&mut vectors, u64::from(index.total_vectors), "BMP vectors"),
355 (&mut blocks, u64::from(index.num_blocks), "BMP blocks"),
356 (
357 &mut real_slots,
358 u64::from(index.num_real_docs()),
359 "BMP real vector slots",
360 ),
361 (
362 &mut virtual_slots,
363 u64::from(index.num_virtual_docs),
364 "BMP padded virtual slots",
365 ),
366 ] {
367 if let Some(total) = capacity.add(count) {
368 return Err(field_capacity_error(
369 field.0,
370 &entry.name,
371 value_kind,
372 total,
373 ));
374 }
375 }
376 }
377 }
378 SparseFormat::MaxScore => {
379 let mut vectors = MergeCapacity::default();
380 let mut dimensions: FxHashMap<u32, (MergeCapacity, MergeCapacity)> =
381 FxHashMap::default();
382
383 for segment in segments {
384 let Some(index) = segment.sparse_indexes().get(&field.0) else {
385 continue;
386 };
387 if let Some(total) = vectors.add(u64::from(index.total_vectors)) {
388 return Err(field_capacity_error(
389 field.0,
390 &entry.name,
391 "MaxScore vectors",
392 total,
393 ));
394 }
395
396 for (dimension, doc_count, block_count) in index.dimension_counts()
397 {
398 let (docs, blocks) = dimensions.entry(dimension).or_default();
399 if let Some(total) = docs.add(u64::from(doc_count)) {
400 return Err(field_capacity_error(
401 field.0,
402 &entry.name,
403 &format!("MaxScore postings for dimension {dimension}"),
404 total,
405 ));
406 }
407 if let Some(total) = blocks.add(u64::from(block_count)) {
408 return Err(field_capacity_error(
409 field.0,
410 &entry.name,
411 &format!("MaxScore blocks for dimension {dimension}"),
412 total,
413 ));
414 }
415 if let Some(total) =
416 maxscore_skip_entries.add(u64::from(block_count))
417 {
418 return Err(crate::Error::Schema(format!(
419 "merge would produce {total} MaxScore skip entries \
420 across sparse fields, exceeding the segment format \
421 limit {}; lower max_segment_docs for multi-valued \
422 sparse fields",
423 u32::MAX,
424 )));
425 }
426 }
427 }
428 }
429 }
430 }
431 _ => {}
432 }
433 }
434 Ok(())
435 }
436
437 pub async fn merge<D: Directory + DirectoryWriter>(
451 &self,
452 dir: &D,
453 segments: &[SegmentReader],
454 new_segment_id: SegmentId,
455 trained: Option<&TrainedVectorStructures>,
456 ) -> Result<(SegmentMeta, MergeStats)> {
457 self.ensure_not_cancelled()?;
458 if segments
459 .iter()
460 .any(|segment| segment.deletion_meta().is_some())
461 {
462 return Err(crate::Error::Schema(
463 "encoded-copy merge cannot discard deletion masks; use IndexWriter::force_merge or SegmentMerger::compact".into(),
464 ));
465 }
466 let total_docs: u32 = segments
470 .iter()
471 .try_fold(0u32, |acc, segment| acc.checked_add(segment.num_docs()))
472 .ok_or_else(|| {
473 crate::Error::Internal(format!(
474 "Total document count exceeds u32::MAX ({})",
475 u32::MAX
476 ))
477 })?;
478
479 self.validate_merge_capacities(segments)?;
480
481 let mut stats = MergeStats::default();
482 let files = SegmentFiles::new(new_segment_id.0);
483
484 let merge_start = std::time::Instant::now();
498
499 let reorder_text = self.reorder_fields
501 && self.schema.fields().any(|(_, entry)| {
502 entry.indexed && entry.reorder && entry.field_type == FieldType::Text
503 });
504 let postings_fut = async {
505 let _permit = if reorder_text {
506 self.acquire_reorder_permit().await?
507 } else {
508 None
509 };
510 let plans = if reorder_text {
511 crate::segment::text_reorder::plan_text_reorders_from_sources(
512 segments,
513 &self.schema,
514 self.bp_memory_budget,
515 self.bp_budget,
516 self.cancellation.as_deref(),
517 self.background_pool.clone(),
518 true,
519 )
520 .await?
521 } else {
522 Vec::new()
523 };
524 let text_converged = plans.iter().all(|plan| plan.converged);
525 let mut postings_writer =
526 OffsetWriter::new(dir.streaming_writer_cold(&files.postings).await?);
527 let mut positions_writer =
528 OffsetWriter::new(dir.streaming_writer_cold(&files.positions).await?);
529 let mut term_dict_writer =
530 OffsetWriter::new(dir.streaming_writer_cold(&files.term_dict).await?);
531
532 let posting_stats = self
533 .merge_postings(
534 segments,
535 &mut term_dict_writer,
536 &mut postings_writer,
537 &mut positions_writer,
538 &plans,
539 )
540 .await?;
541 let terms_processed = posting_stats.terms_processed;
542
543 let postings_bytes = postings_writer.offset() as usize;
544 let term_dict_bytes = term_dict_writer.offset() as usize;
545 let positions_bytes = positions_writer.offset();
546
547 postings_writer.finish()?;
548 term_dict_writer.finish()?;
549 if positions_bytes > 0 {
550 positions_writer.finish()?;
551 } else {
552 drop(positions_writer);
553 let _ = dir.delete(&files.positions).await;
554 }
555 log::info!(
556 "[merge] index={} postings done: {} terms, term_dict={}, postings={}, positions={}",
557 self.schema.index_label(),
558 terms_processed,
559 crate::format_bytes(term_dict_bytes as u64),
560 crate::format_bytes(postings_bytes as u64),
561 crate::format_bytes(positions_bytes),
562 );
563 if reorder_text {
566 self.merge_chunk_maps(dir, segments, &files, &plans).await?;
567 }
568 Ok::<(usize, usize, usize, usize, bool), crate::Error>((
569 terms_processed,
570 term_dict_bytes,
571 postings_bytes,
572 posting_stats.blocks_without_bounds,
573 text_converged,
574 ))
575 };
576
577 let store_fut = async {
578 let mut store_writer =
579 OffsetWriter::new(dir.streaming_writer_cold(&files.store).await?);
580 let store_num_docs = self.merge_store(segments, &mut store_writer).await?;
581 let bytes = store_writer.offset() as usize;
582 store_writer.finish()?;
583 Ok::<(usize, u32), crate::Error>((bytes, store_num_docs))
584 };
585
586 let fast_fut = async { self.merge_fast_fields(dir, segments, &files).await };
587
588 let chunks_fut = async {
589 if reorder_text {
590 Ok(0)
591 } else {
592 self.merge_chunk_maps(dir, segments, &files, &[]).await
593 }
594 };
595 let (postings_result, store_result, fast_bytes, _) =
596 tokio::try_join!(postings_fut, store_fut, fast_fut, chunks_fut)?;
597 self.ensure_not_cancelled()?;
598
599 log::info!(
600 "[merge] index={} stage 1 done in {:.1}s (postings + store + fast)",
601 self.schema.index_label(),
602 merge_start.elapsed().as_secs_f64()
603 );
604
605 let sparse_fut = async { self.merge_sparse_vectors(dir, segments, &files).await };
609
610 let dense_fut = async {
611 self.merge_dense_vectors(dir, segments, &files, trained, AnnWriteMode::Copy)
612 .await
613 };
614
615 let ((sparse_bytes, bp_converged), vectors_bytes) = if self.reorder_fields {
619 let sparse = sparse_fut.await?;
620 let dense = dense_fut.await?;
621 (sparse, dense)
622 } else {
623 tokio::try_join!(sparse_fut, dense_fut)?
624 };
625 self.ensure_not_cancelled()?;
626 let (store_bytes, store_num_docs) = store_result;
627 stats.terms_processed = postings_result.0;
628 stats.term_dict_bytes = postings_result.1;
629 stats.postings_bytes = postings_result.2;
630 stats.posting_blocks_without_bounds = postings_result.3;
631 stats.store_bytes = store_bytes;
632 stats.vectors_bytes = vectors_bytes;
633 stats.sparse_bytes = sparse_bytes;
634 stats.bp_converged = bp_converged && postings_result.4;
635 stats.fast_bytes = fast_bytes;
636 log::info!(
637 "[merge] index={} all phases done in {:.1}s: {}",
638 self.schema.index_label(),
639 merge_start.elapsed().as_secs_f64(),
640 stats
641 );
642
643 self.merge_row_stats(dir, segments, &files).await?;
644
645 self.ensure_not_cancelled()?;
647 let mut merged_field_stats: FxHashMap<u32, FieldStats> = FxHashMap::default();
648 for segment in segments {
649 for (&field_id, field_stats) in &segment.meta().field_stats {
650 let entry = merged_field_stats.entry(field_id).or_default();
651 entry.total_tokens = entry
652 .total_tokens
653 .checked_add(field_stats.total_tokens)
654 .ok_or_else(|| {
655 crate::Error::Corruption(format!(
656 "field {} total-token count overflow while merging",
657 field_id
658 ))
659 })?;
660 entry.doc_count = entry
661 .doc_count
662 .checked_add(field_stats.doc_count)
663 .ok_or_else(|| {
664 crate::Error::Corruption(format!(
665 "field {} document count overflow while merging",
666 field_id
667 ))
668 })?;
669 }
670 }
671
672 if store_num_docs != total_docs {
676 log::error!(
677 "[merge] index={} STORE/META MISMATCH: store has {} docs but metadata expects {}. \
678 Per-segment: {:?}",
679 self.schema.index_label(),
680 store_num_docs,
681 total_docs,
682 segments
683 .iter()
684 .map(|s| (
685 format!("{:016x}", s.meta().id),
686 s.num_docs(),
687 s.store().num_docs()
688 ))
689 .collect::<Vec<_>>()
690 );
691 return Err(crate::Error::Io(std::io::Error::new(
692 std::io::ErrorKind::InvalidData,
693 format!(
694 "Store/meta doc count mismatch: store={}, meta={}",
695 store_num_docs, total_docs
696 ),
697 )));
698 }
699
700 let meta = SegmentMeta {
701 id: new_segment_id.0,
702 num_docs: total_docs,
703 field_stats: merged_field_stats,
704 };
705
706 self.ensure_not_cancelled()?;
707
708 dir.write_durable(&files.meta, &meta.serialize()?).await?;
712
713 log::info!(
718 "[merge] index={} complete: {} docs, {}",
719 self.schema.index_label(),
720 total_docs,
721 stats
722 );
723
724 Ok((meta, stats))
725 }
726}
727
728pub async fn delete_segment<D: Directory + DirectoryWriter>(
730 dir: &D,
731 segment_id: SegmentId,
732) -> Result<()> {
733 let files = SegmentFiles::new(segment_id.0);
734 let paths = files.lifecycle_paths();
735 let results = futures::future::join_all(paths.iter().map(|path| dir.delete(path))).await;
736
737 for result in results {
741 if let Err(error) = result
742 && error.kind() != std::io::ErrorKind::NotFound
743 {
744 return Err(crate::Error::Io(error));
745 }
746 }
747 Ok(())
748}
749
750#[cfg(test)]
751mod capacity_tests {
752 use super::{MergeCapacity, field_capacity_error};
753
754 #[test]
755 fn merge_capacity_accepts_the_exact_u32_boundary() {
756 let mut capacity = MergeCapacity::default();
757 assert_eq!(capacity.add(u64::from(u32::MAX) - 7), None);
758 assert_eq!(capacity.add(7), None);
759 }
760
761 #[test]
762 fn merge_capacity_rejects_the_first_value_beyond_u32() {
763 let mut capacity = MergeCapacity::default();
764 assert_eq!(capacity.add(u64::from(u32::MAX)), None);
765 assert_eq!(capacity.add(1), Some(u64::from(u32::MAX) + 1));
766 }
767
768 #[test]
769 fn merge_capacity_failure_is_not_source_corruption() {
770 let error = field_capacity_error(
771 7,
772 "body_embedding",
773 "dense vectors",
774 u64::from(u32::MAX) + 1,
775 );
776 assert!(matches!(error, crate::Error::Schema(_)));
777 assert!(error.to_string().contains("lower max_segment_docs"));
778 }
779}