try to use select-nth instead of full sort in segment level agg top-k selection

columnar/benches/common.rs

@@ -54,6 +54,6 @@ pub fn generate_columnar_with_name(card: Card, num_docs: u32, column_name: &str)
     }
 
     let mut wrt: Vec<u8> = Vec::new();
-    columnar_writer.serialize(num_docs, None, &mut wrt).unwrap();
+    columnar_writer.serialize(num_docs, &mut wrt).unwrap();
     ColumnarReader::open(wrt).unwrap()
 }

common/src/bitset.rs

@@ -281,16 +281,12 @@ impl BitSet {
     }
 
     /// Inserts an element in the `BitSet`
-    ///
-    /// Returns true if the set changed.
     #[inline]
-    pub fn insert(&mut self, el: u32) -> bool {
+    pub fn insert(&mut self, el: u32) {
         // we do not check saturated els.
         let higher = el / 64u32;
         let lower = el % 64u32;
-        let changed = self.tinysets[higher as usize].insert_mut(lower);
-        self.len += u64::from(changed);
-        changed
+        self.len += u64::from(self.tinysets[higher as usize].insert_mut(lower));
     }
 
     /// Inserts an element in the `BitSet`

src/aggregation/agg_data.rs

@@ -931,9 +931,7 @@ fn build_allowed_term_ids_for_str(
         // add matches
         allowed = Some(BitSet::with_max_value(allowed_capacity));
         let allowed = allowed.as_mut().unwrap();
-        for_each_matching_term_ord(str_col, include, |ord| {
-            let _ = allowed.insert(ord);
-        })?;
+        for_each_matching_term_ord(str_col, include, |ord| allowed.insert(ord))?;
     };
 
     if let Some(exclude) = exclude {

src/aggregation/bucket/term_agg/mod.rs

@@ -981,19 +981,27 @@ where
     ) -> crate::Result<IntermediateBucketResult> {
         let mut entries: Vec<(u64, Bucket)> = term_buckets.into_vec();
 
+        let segment_size = term_req.req.segment_size as usize;
+
+        // select_nth_unstable_by_key(segment_size, ...) places the (k+1)-th element at
+        // entries[segment_size] and guarantees entries[0..segment_size] are the top-k,
+        // unordered. We need this to properly compute term_doc_count_before_cutoff.
         match &term_req.req.order.target {
             OrderTarget::Key => {
                 // We rely on the fact, that term ordinals match the order of the strings
                 // TODO: We could have a special collector, that keeps only TOP n results at any
                 // time.
-                if term_req.req.order.order == Order::Desc {
-                    entries.sort_unstable_by_key(|bucket| std::cmp::Reverse(bucket.0));
-                } else {
-                    entries.sort_unstable_by_key(|bucket| bucket.0);
+                if entries.len() > segment_size {
+                    if term_req.req.order.order == Order::Desc {
+                        entries
+                            .select_nth_unstable_by_key(segment_size, |b| std::cmp::Reverse(b.0));
+                    } else {
+                        entries.select_nth_unstable_by_key(segment_size, |b| b.0);
+                    }
                 }
             }
             OrderTarget::SubAggregation(sub_agg_path) => {
-                // Peek segment-level metric values, sort, then fall through to
+                // Peek segment-level metric values, select top-k, then fall through to
                 // `cut_off_buckets`. Like Elasticsearch, we always cut off when ordering
                 // by a sub-agg: top-K results are approximate and may differ from the
                 // global ordering, especially for non-monotonic metrics like avg/min.
@@ -1003,7 +1011,7 @@ where
                     ))
                 })?;
                 let (agg_name, agg_prop) = get_agg_name_and_property(sub_agg_path);
-                // Fetch values up-front; otherwise sort would re-compute per comparison
+                // Fetch values up-front; otherwise sort would re-compute per call
                 let mut keyed: Vec<(f64, (u64, Bucket))> = entries
                     .into_iter()
                     .map(|bucket| {
@@ -1013,28 +1021,34 @@ where
                         (metric_value, bucket)
                     })
                     .collect();
-                if term_req.req.order.order == Order::Desc {
-                    keyed.sort_unstable_by(|a, b| {
-                        b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal)
-                    });
-                } else {
-                    keyed.sort_unstable_by(|a, b| {
-                        a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal)
-                    });
+                if keyed.len() > segment_size {
+                    if term_req.req.order.order == Order::Desc {
+                        keyed.select_nth_unstable_by(segment_size, |a, b| {
+                            b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal)
+                        });
+                    } else {
+                        keyed.select_nth_unstable_by(segment_size, |a, b| {
+                            a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal)
+                        });
+                    }
                 }
                 entries = keyed.into_iter().map(|(_, e)| e).collect();
             }
             OrderTarget::Count => {
-                if term_req.req.order.order == Order::Desc {
-                    entries.sort_unstable_by_key(|bucket| std::cmp::Reverse(bucket.1.count));
-                } else {
-                    entries.sort_unstable_by_key(|bucket| bucket.1.count);
+                if entries.len() > segment_size {
+                    if term_req.req.order.order == Order::Desc {
+                        entries.select_nth_unstable_by_key(segment_size, |b| {
+                            std::cmp::Reverse(b.1.count)
+                        });
+                    } else {
+                        entries.select_nth_unstable_by_key(segment_size, |b| b.1.count);
+                    }
                 }
             }
         }
 
         let (term_doc_count_before_cutoff, sum_other_doc_count) =
-            cut_off_buckets(&mut entries, term_req.req.segment_size as usize);
+            cut_off_buckets(&mut entries, segment_size);
 
         let mut dict: FxHashMap<IntermediateKey, IntermediateTermBucketEntry> = Default::default();
         dict.reserve(entries.len());

src/core/tests.rs

@@ -1,14 +1,14 @@
 use crate::collector::Count;
 use crate::directory::{RamDirectory, WatchCallback};
 use crate::index::SegmentId;
-use crate::indexer::{DocIdMapping, LogMergePolicy, NoMergePolicy};
+use crate::indexer::{LogMergePolicy, NoMergePolicy};
 use crate::postings::Postings;
 use crate::query::TermQuery;
-use crate::schema::{Field, IndexRecordOption, Schema, Value, INDEXED, STORED, STRING, TEXT};
+use crate::schema::{Field, IndexRecordOption, Schema, INDEXED, STRING, TEXT};
 use crate::tokenizer::TokenizerManager;
 use crate::{
-    Directory, DocAddress, DocSet, Index, IndexBuilder, IndexReader, IndexSettings, IndexWriter,
-    ReloadPolicy, TantivyDocument, Term,
+    Directory, DocSet, Index, IndexBuilder, IndexReader, IndexSettings, IndexWriter, ReloadPolicy,
+    TantivyDocument, Term,
 };
 
 #[test]
@@ -300,51 +300,6 @@ 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 settings = IndexSettings {
-        manual_doc_id_mapping: true,
-        ..Default::default()
-    };
-    let mut single_segment_index_writer = Index::builder()
-        .schema(schema)
-        .settings(settings)
-        .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::new_permutation(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();

src/index/index_meta.rs

@@ -250,11 +250,6 @@ pub struct IndexSettings {
     /// provided in `IndexSortByField`
     #[serde(skip_serializing_if = "Option::is_none")]
     pub sort_by_field: Option<IndexSortByField>,
-    /// If true, enables caller-provided doc id mappings at segment finalization time.
-    /// Always skip serializing this field since it's only used at segment finalization time.
-    #[doc(hidden)]
-    #[serde(skip)]
-    pub manual_doc_id_mapping: bool,
     /// The `Compressor` used to compress the doc store.
     #[serde(default)]
     pub docstore_compression: Compressor,
@@ -278,7 +273,6 @@ impl Default for IndexSettings {
     fn default() -> Self {
         Self {
             sort_by_field: None,
-            manual_doc_id_mapping: false,
             docstore_compression: Compressor::default(),
             docstore_blocksize: default_docstore_blocksize(),
             docstore_compress_dedicated_thread: true,
@@ -466,7 +460,6 @@ mod tests {
                     field: "text".to_string(),
                     order: Order::Asc,
                 }),
-                manual_doc_id_mapping: false,
                 docstore_compression: crate::store::Compressor::Zstd(ZstdCompressor {
                     compression_level: Some(4),
                 }),
@@ -536,7 +529,6 @@ mod tests {
             index_settings,
             IndexSettings {
                 sort_by_field: None,
-                manual_doc_id_mapping: false,
                 docstore_compression: Compressor::default(),
                 docstore_compress_dedicated_thread: true,
                 docstore_blocksize: 16_384
@@ -555,18 +547,6 @@ mod tests {
                 serde_json::from_value(index_settings_json).unwrap();
             assert_eq!(index_settings_deser, index_settings);
         }
-        {
-            index_settings.manual_doc_id_mapping = true;
-            let index_settings_json = serde_json::to_value(&index_settings).unwrap();
-            assert_eq!(
-                index_settings_json,
-                serde_json::json!({
-                    "docstore_compression": "lz4",
-                    "docstore_blocksize": 16384
-                })
-            );
-            index_settings.manual_doc_id_mapping = false;
-        }
         {
             index_settings.docstore_compress_dedicated_thread = false;
             let index_settings_json = serde_json::to_value(&index_settings).unwrap();

src/indexer/doc_id_mapping.rs

@@ -1,6 +1,7 @@
 //! This module is used when sorting the index by a property, e.g.
 //! to get mappings from old doc_id to new doc_id and vice versa, after sorting
-use common::{BitSet, ReadOnlyBitSet};
+
+use common::ReadOnlyBitSet;
 
 use super::SegmentWriter;
 use crate::schema::{Field, Schema};
@@ -70,33 +71,7 @@ pub struct DocIdMapping {
 }
 
 impl DocIdMapping {
-    /// Creates a `DocIdMapping` from a mapping of new doc ids to old doc ids, with permutation validation.
-    /// The mapping is validated by checking that every old doc id appears exactly once in the mapping.
-    /// I.e., doc ids must be consecutive from `0` to `new_doc_id_to_old.len() - 1`, inclusive.
-    pub fn new_permutation(new_doc_id_to_old: Vec<DocId>) -> crate::Result<Self> {
-        // Check that the mapping is a permutation of the segment doc ids.
-        let max_doc = new_doc_id_to_old.len() as DocId;
-        let mut old_doc_id_to_new = vec![0; max_doc as usize];
-
-        let mut seen_doc_ids = BitSet::with_max_value(max_doc);
-        for (i, old_doc_id) in new_doc_id_to_old.iter().copied().enumerate() {
-            if old_doc_id >= max_doc || !seen_doc_ids.insert(old_doc_id) {
-                return Err(TantivyError::InvalidArgument(
-                    "Mapping must be a permutation of the segment doc ids".to_string(),
-                ));
-            }
-            old_doc_id_to_new[new_doc_id_to_old[i] as usize] = i as DocId;
-        }
-
-        let doc_id_mapping = DocIdMapping {
-            new_doc_id_to_old,
-            old_doc_id_to_new,
-        };
-        Ok(doc_id_mapping)
-    }
-
-    /// Creates a `DocIdMapping` from a mapping of new doc ids to old doc ids.
-    pub(crate) fn from_new_id_to_old_id(new_doc_id_to_old: Vec<DocId>) -> Self {
+    pub fn from_new_id_to_old_id(new_doc_id_to_old: Vec<DocId>) -> Self {
         let max_doc = new_doc_id_to_old.len();
         let old_max_doc = new_doc_id_to_old
             .iter()
@@ -114,41 +89,35 @@ impl DocIdMapping {
         }
     }
 
-    /// Returns the new doc_id for the old doc_id
-    pub(crate) fn get_new_doc_id(&self, doc_id: DocId) -> DocId {
+    /// returns the new doc_id for the old doc_id
+    pub fn get_new_doc_id(&self, doc_id: DocId) -> DocId {
         self.old_doc_id_to_new[doc_id as usize]
     }
-
-    /// Iiterate over old doc_ids in order of the new doc_ids
-    pub(crate) fn iter_old_doc_ids(&self) -> impl Iterator<Item = DocId> + Clone + '_ {
-        self.new_doc_id_to_old.iter().copied()
+    /// returns the old doc_id for the new doc_id
+    pub fn get_old_doc_id(&self, doc_id: DocId) -> DocId {
+        self.new_doc_id_to_old[doc_id as usize]
+    }
+    /// iterate over old doc_ids in order of the new doc_ids
+    pub fn iter_old_doc_ids(&self) -> impl Iterator<Item = DocId> + Clone + '_ {
+        self.new_doc_id_to_old.iter().cloned()
     }
 
-    /// Returns the new doc_ids in order of the old doc_ids
-    pub(crate) fn old_to_new_ids(&self) -> &[DocId] {
+    pub fn old_to_new_ids(&self) -> &[DocId] {
         &self.old_doc_id_to_new[..]
     }
 
     /// Remaps a given array to the new doc ids.
-    pub(crate) fn remap<T: Copy>(&self, els: &[T]) -> Vec<T> {
+    pub fn remap<T: Copy>(&self, els: &[T]) -> Vec<T> {
         self.new_doc_id_to_old
             .iter()
             .map(|old_doc| els[*old_doc as usize])
             .collect()
     }
-
-    /// Returns the number of documents in the mapping.
-    pub(crate) fn len(&self) -> usize {
-        // new_doc_id_to_old and old_doc_id_to_new have the same length by construction.
+    pub fn num_new_doc_ids(&self) -> usize {
         self.new_doc_id_to_old.len()
     }
-}
-
-#[cfg(test)]
-impl DocIdMapping {
-    /// returns the old doc_id for the new doc_id
-    fn get_old_doc_id(&self, doc_id: DocId) -> DocId {
-        self.new_doc_id_to_old[doc_id as usize]
+    pub fn num_old_doc_ids(&self) -> usize {
+        self.old_doc_id_to_new.len()
     }
 }
 
@@ -189,7 +158,7 @@ mod tests_indexsorting {
     use crate::indexer::doc_id_mapping::DocIdMapping;
     use crate::indexer::NoMergePolicy;
     use crate::query::QueryParser;
-    use crate::{schema::*, TantivyError};
+    use crate::schema::*;
     use crate::{DocAddress, Index, IndexBuilder, IndexSettings, IndexSortByField, Order};
 
     fn create_test_index(
@@ -581,18 +550,6 @@ mod tests_indexsorting {
         assert_eq!(doc_mapping.get_new_doc_id(5), 2);
     }
 
-    #[test]
-    fn test_doc_mapping_new_permutation_rejects_out_of_range() {
-        let result = DocIdMapping::new_permutation(vec![5, 0]);
-        assert!(matches!(result, Err(TantivyError::InvalidArgument(_)),));
-    }
-
-    #[test]
-    fn test_doc_mapping_new_permutation_rejects_duplicates() {
-        let result = DocIdMapping::new_permutation(vec![0, 1, 0]);
-        assert!(matches!(result, Err(TantivyError::InvalidArgument(_)),));
-    }
-
     #[test]
     fn test_doc_mapping_remap() {
         let doc_mapping = DocIdMapping::from_new_id_to_old_id(vec![2, 8, 3]);

src/indexer/mod.rs

@@ -33,7 +33,6 @@ mod stamper;
 use crossbeam_channel as channel;
 use smallvec::SmallVec;
 
-pub use self::doc_id_mapping::DocIdMapping;
 pub use self::index_writer::{advance_deletes, IndexWriter, IndexWriterOptions};
 pub use self::log_merge_policy::LogMergePolicy;
 pub use self::merge_operation::MergeOperation;

src/indexer/segment_serializer.rs

@@ -4,7 +4,7 @@ use crate::directory::WritePtr;
 use crate::fieldnorm::FieldNormsSerializer;
 use crate::index::{Segment, SegmentComponent};
 use crate::postings::InvertedIndexSerializer;
-use crate::store::{Compressor, StoreWriter};
+use crate::store::StoreWriter;
 
 /// Segment serializer is in charge of laying out on disk
 /// the data accumulated and sorted by the `SegmentWriter`.
@@ -25,18 +25,17 @@ impl SegmentSerializer {
         // If the segment is going to be sorted, we stream the docs first to a temporary file.
         // In the merge case this is not necessary because we can kmerge the already sorted
         // segments
+        let remapping_required = segment.index().settings().sort_by_field.is_some() && !is_in_merge;
         let settings = segment.index().settings().clone();
-        let remapping_required =
-            (settings.sort_by_field.is_some() || settings.manual_doc_id_mapping) && !is_in_merge;
         let store_writer = if remapping_required {
             let store_write = segment.open_write(SegmentComponent::TempStore)?;
             StoreWriter::new(
                 store_write,
-                Compressor::None,
+                crate::store::Compressor::None,
                 // We want fast random access on the docs, so we choose a small block size.
                 // If this is zero, the skip index will contain too many checkpoints and
                 // therefore will be relatively slow.
-                16_000,
+                16000,
                 settings.docstore_compress_dedicated_thread,
             )?
         } else {

src/indexer/segment_writer.rs

@@ -136,8 +136,10 @@ impl SegmentWriter {
 
     /// Lay on disk the current content of the `SegmentWriter`
     ///
-    /// Finalize consumes the `SegmentWriter`, so that it cannot be used afterwards.
-    pub fn finalize(self) -> crate::Result<Vec<u64>> {
+    /// Finalize consumes the `SegmentWriter`, so that it cannot
+    /// be used afterwards.
+    pub fn finalize(mut self) -> crate::Result<Vec<u64>> {
+        self.fieldnorms_writer.fill_up_to_max_doc(self.max_doc);
         let mapping: Option<DocIdMapping> = self
             .segment_serializer
             .segment()
@@ -147,41 +149,6 @@ impl SegmentWriter {
             .clone()
             .map(|sort_by_field| get_doc_id_mapping_from_field(sort_by_field, &self))
             .transpose()?;
-        self.finalize_inner(mapping.as_ref())
-    }
-
-    /// Lay on disk the current content of the `SegmentWriter` using the provided doc id mapping.
-    ///
-    /// Finalize consumes the `SegmentWriter`, so that it cannot be used afterwards.
-    pub fn finalize_with_doc_id_mapping(self, mapping: &DocIdMapping) -> crate::Result<Vec<u64>> {
-        // Ensure the segment writer was created in remap mode so the docstore can be reordered.
-        if !self
-            .segment_serializer
-            .segment()
-            .index()
-            .settings()
-            .manual_doc_id_mapping
-        {
-            return Err(TantivyError::InvalidArgument(
-                "IndexSettings::manual_doc_id_mapping must be set to true".to_string(),
-            ));
-        }
-
-        // Check that the mapping eventually covers all documents in the segment.
-        if mapping.len() != self.max_doc as usize {
-            return Err(TantivyError::InvalidArgument(format!(
-                "Mapping must cover all documents in this segment. Expected {} documents, got {}",
-                self.max_doc,
-                mapping.len()
-            )));
-        }
-
-        self.finalize_inner(Some(mapping))
-    }
-
-    fn finalize_inner(mut self, mapping: Option<&DocIdMapping>) -> crate::Result<Vec<u64>> {
-        // Pad before remapping; the mapping indexes fieldnorms by old doc id.
-        self.fieldnorms_writer.fill_up_to_max_doc(self.max_doc);
         remap_and_write(
             self.schema,
             &self.per_field_postings_writers,
@@ -189,9 +156,9 @@ impl SegmentWriter {
             self.fast_field_writers,
             &self.fieldnorms_writer,
             self.segment_serializer,
-            mapping,
+            mapping.as_ref(),
         )?;
-        let doc_opstamps = remap_doc_opstamps(self.doc_opstamps, mapping);
+        let doc_opstamps = remap_doc_opstamps(self.doc_opstamps, mapping.as_ref());
         Ok(doc_opstamps)
     }
 
@@ -518,7 +485,6 @@ mod tests {
     use crate::collector::{Count, TopDocs};
     use crate::directory::RamDirectory;
     use crate::fastfield::FastValue;
-    use crate::indexer::doc_id_mapping::DocIdMapping;
     use crate::postings::{Postings, TermInfo};
     use crate::query::{PhraseQuery, QueryParser};
     use crate::schema::{
@@ -531,7 +497,7 @@ mod tests {
     use crate::tokenizer::{PreTokenizedString, Token};
     use crate::{
         DateTime, Directory, DocAddress, DocSet, Index, IndexWriter, SegmentReader,
-        TantivyDocument, TantivyError, Term, TERMINATED,
+        TantivyDocument, Term, TERMINATED,
     };
 
     #[test]
@@ -1170,87 +1136,4 @@ mod tests {
             "Schema error: 'Error getting tokenizer for field: title'"
         );
     }
-
-    /// Builds a `SegmentWriter` with a fast `u64` field and a text field that only some
-    /// documents populate, so the text field is missing fieldnorms on some docs.
-    ///
-    /// The `texts` slice provides, for each document, an optional text value. The order
-    /// number is always recorded in the `order` fast field so callers can recover the
-    /// original document via that value.
-    fn build_segment_writer_with_doc_id_mapping(
-        texts: &[Option<&str>],
-    ) -> (Index, crate::Segment, super::SegmentWriter) {
-        let mut schema_builder = Schema::builder();
-        schema_builder.add_u64_field("order", FAST | STORED);
-        schema_builder.add_text_field("text", TEXT);
-        let schema = schema_builder.build();
-        let mut index = Index::create_in_ram(schema);
-        index.settings_mut().manual_doc_id_mapping = true;
-        let segment = index.new_segment();
-        let order = index.schema().get_field("order").unwrap();
-        let text = index.schema().get_field("text").unwrap();
-        let mut segment_writer =
-            super::SegmentWriter::for_segment(15_000_000, segment.clone()).unwrap();
-        for (opstamp, text_opt) in texts.iter().enumerate() {
-            let mut doc = TantivyDocument::default();
-            doc.add_u64(order, opstamp as u64);
-            if let Some(text_value) = text_opt {
-                doc.add_text(text, *text_value);
-            }
-            segment_writer
-                .add_document(crate::indexer::AddOperation {
-                    opstamp: opstamp as u64,
-                    document: doc,
-                })
-                .unwrap();
-        }
-        (index, segment, segment_writer)
-    }
-
-    #[test]
-    fn test_finalize_with_doc_id_mapping_rejects_wrong_length() {
-        let (_index, _segment, segment_writer) =
-            build_segment_writer_with_doc_id_mapping(&[Some("a"), Some("b"), Some("c")]);
-        // Mapping only covers 2 of the 3 documents.
-        let mapping = DocIdMapping::new_permutation(vec![1, 0]).unwrap();
-        let err = segment_writer
-            .finalize_with_doc_id_mapping(&mapping)
-            .unwrap_err();
-        assert!(
-            matches!(err, TantivyError::InvalidArgument(_)),
-            "unexpected error: {err:?}"
-        );
-    }
-
-    #[test]
-    fn test_finalize_with_doc_id_mapping_remaps_missing_fieldnorms() -> crate::Result<()> {
-        // doc 0: "alpha beta"  (2 tokens)
-        // doc 1: <no text>     (missing fieldnorm -> 0)
-        // doc 2: "gamma"       (1 token)
-        // doc 3: <no text>     (missing fieldnorm -> 0)
-        let (index, segment, segment_writer) = build_segment_writer_with_doc_id_mapping(&[
-            Some("alpha beta"),
-            None,
-            Some("gamma"),
-            None,
-        ]);
-        let max_doc = segment_writer.max_doc();
-
-        // Reverse the documents. New doc id i maps to old doc id (3 - i).
-        let mapping = DocIdMapping::new_permutation(vec![3, 2, 1, 0])?;
-        segment_writer.finalize_with_doc_id_mapping(&mapping)?;
-
-        let segment = segment.with_max_doc(max_doc);
-        let segment_reader = SegmentReader::open(&segment)?;
-        let text = index.schema().get_field("text").unwrap();
-        let fieldnorm_reader = segment_reader.get_fieldnorms_reader(text)?;
-
-        // After remapping, fieldnorms follow the reversed order:
-        // new 0 <- old 3 (0), new 1 <- old 2 (1), new 2 <- old 1 (0), new 3 <- old 0 (2)
-        assert_eq!(fieldnorm_reader.fieldnorm(0), 0);
-        assert_eq!(fieldnorm_reader.fieldnorm(1), 1);
-        assert_eq!(fieldnorm_reader.fieldnorm(2), 0);
-        assert_eq!(fieldnorm_reader.fieldnorm(3), 2);
-        Ok(())
-    }
 }

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::{DocIdMapping, SegmentWriter};
+use crate::indexer::SegmentWriter;
 use crate::schema::document::Document;
 use crate::{Directory, Index, IndexMeta, Opstamp, Segment, TantivyDocument};
 
@@ -38,29 +38,9 @@ impl<D: Document> SingleSegmentIndexWriter<D> {
     }
 
     pub fn finalize(self) -> crate::Result<Index> {
-        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 max_doc = self.segment_writer.max_doc();
+        self.segment_writer.finalize()?;
+        let segment: Segment = self.segment.with_max_doc(max_doc);
         let index = segment.index();
         let index_meta = IndexMeta {
             index_settings: index.settings().clone(),
@@ -71,6 +51,6 @@ impl<D: Document> SingleSegmentIndexWriter<D> {
         };
         save_metas(&index_meta, index.directory())?;
         index.directory().sync_directory()?;
-        Ok(index.clone())
+        Ok(segment.index().clone())
     }
 }