From f7355e60cdb4f076ce6ea46cfd1b9e9ec733c16c Mon Sep 17 00:00:00 2001
From: Luca Cominardi <luca.cominardi@datadoghq.com>
Date: Fri, 26 Jun 2026 16:47:14 +0200
Subject: [PATCH] feat: add single segment doc id mapping finalization

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 src/core/tests.rs                          | 48 ++++++++++++++++++++--
 src/indexer/single_segment_index_writer.rs | 30 +++++++++++---
 2 files changed, 69 insertions(+), 9 deletions(-)

diff --git a/src/core/tests.rs b/src/core/tests.rs
index 62baedf1d..eacb9c9da 100644
--- a/src/core/tests.rs
+++ b/src/core/tests.rs
@@ -1,14 +1,14 @@
 use crate::collector::Count;
 use crate::directory::{RamDirectory, WatchCallback};
 use crate::index::SegmentId;
-use crate::indexer::{LogMergePolicy, NoMergePolicy};
+use crate::indexer::{DocIdMapping, LogMergePolicy, NoMergePolicy};
 use crate::postings::Postings;
 use crate::query::TermQuery;
-use crate::schema::{Field, IndexRecordOption, Schema, INDEXED, STRING, TEXT};
+use crate::schema::{Field, IndexRecordOption, Schema, Value, INDEXED, STORED, STRING, TEXT};
 use crate::tokenizer::TokenizerManager;
 use crate::{
-    Directory, DocSet, Index, IndexBuilder, IndexReader, IndexSettings, IndexWriter, ReloadPolicy,
-    TantivyDocument, Term,
+    Directory, DocAddress, DocSet, Index, IndexBuilder, IndexReader, IndexSettings, IndexWriter,
+    ReloadPolicy, TantivyDocument, Term,
 };
 
 #[test]
@@ -300,6 +300,46 @@ fn test_single_segment_index_writer() -> crate::Result<()> {
     Ok(())
 }
 
+#[test]
+fn test_single_segment_index_writer_with_doc_id_mapping() -> crate::Result<()> {
+    let mut schema_builder = Schema::builder();
+    let text_field = schema_builder.add_text_field("text", TEXT | STORED);
+    let schema = schema_builder.build();
+    let directory = RamDirectory::default();
+    let mut single_segment_index_writer = Index::builder()
+        .schema(schema)
+        .single_segment_index_writer(directory, 15_000_000)?;
+
+    single_segment_index_writer.add_document(doc!(text_field=>"alpha beta"))?;
+    single_segment_index_writer.add_document(doc!())?;
+    single_segment_index_writer.add_document(doc!(text_field=>"gamma"))?;
+
+    let mapping = DocIdMapping::from_new_id_to_old_id(vec![2, 1, 0]);
+    let index = single_segment_index_writer.finalize_with_doc_id_mapping(&mapping)?;
+
+    let searcher = index.reader()?.searcher();
+    let segment_reader = searcher.segment_reader(0);
+    let fieldnorm_reader = segment_reader.get_fieldnorms_reader(text_field)?;
+
+    assert_eq!(fieldnorm_reader.fieldnorm(0), 1);
+    assert_eq!(fieldnorm_reader.fieldnorm(1), 0);
+    assert_eq!(fieldnorm_reader.fieldnorm(2), 2);
+
+    let doc_0 = searcher.doc::<TantivyDocument>(DocAddress::new(0, 0))?;
+    assert_eq!(
+        doc_0.get_first(text_field).and_then(|val| val.as_str()),
+        Some("gamma")
+    );
+    let doc_1 = searcher.doc::<TantivyDocument>(DocAddress::new(0, 1))?;
+    assert!(doc_1.get_first(text_field).is_none());
+    let doc_2 = searcher.doc::<TantivyDocument>(DocAddress::new(0, 2))?;
+    assert_eq!(
+        doc_2.get_first(text_field).and_then(|val| val.as_str()),
+        Some("alpha beta")
+    );
+    Ok(())
+}
+
 #[test]
 fn test_merging_segment_update_docfreq() {
     let mut schema_builder = Schema::builder();
diff --git a/src/indexer/single_segment_index_writer.rs b/src/indexer/single_segment_index_writer.rs
index 673accae3..a30f1f1f5 100644
--- a/src/indexer/single_segment_index_writer.rs
+++ b/src/indexer/single_segment_index_writer.rs
@@ -2,7 +2,7 @@ use std::marker::PhantomData;
 
 use crate::indexer::operation::AddOperation;
 use crate::indexer::segment_updater::save_metas;
-use crate::indexer::SegmentWriter;
+use crate::indexer::{DocIdMapping, SegmentWriter};
 use crate::schema::document::Document;
 use crate::{Directory, Index, IndexMeta, Opstamp, Segment, TantivyDocument};
 
@@ -38,9 +38,29 @@ impl<D: Document> SingleSegmentIndexWriter<D> {
     }
 
     pub fn finalize(self) -> crate::Result<Index> {
-        let max_doc = self.segment_writer.max_doc();
-        self.segment_writer.finalize()?;
-        let segment: Segment = self.segment.with_max_doc(max_doc);
+        let Self {
+            segment,
+            segment_writer,
+            ..
+        } = self;
+        let max_doc = segment_writer.max_doc();
+        segment_writer.finalize()?;
+        Self::finalize_inner(segment, max_doc)
+    }
+
+    pub fn finalize_with_doc_id_mapping(self, mapping: &DocIdMapping) -> crate::Result<Index> {
+        let Self {
+            segment,
+            segment_writer,
+            ..
+        } = self;
+        let max_doc = segment_writer.max_doc();
+        segment_writer.finalize_with_doc_id_mapping(mapping)?;
+        Self::finalize_inner(segment, max_doc)
+    }
+
+    fn finalize_inner(segment: Segment, max_doc: u32) -> crate::Result<Index> {
+        let segment: Segment = segment.with_max_doc(max_doc);
         let index = segment.index();
         let index_meta = IndexMeta {
             index_settings: index.settings().clone(),
@@ -51,6 +71,6 @@ impl<D: Document> SingleSegmentIndexWriter<D> {
         };
         save_metas(&index_meta, index.directory())?;
         index.directory().sync_directory()?;
-        Ok(segment.index().clone())
+        Ok(index.clone())
     }
 }