From a6fd070e3d1d758df695fcf111ae57e032c0dd7f Mon Sep 17 00:00:00 2001 From: Luca Cominardi Date: Mon, 29 Jun 2026 10:52:41 +0200 Subject: [PATCH] fix: validate manual doc id mappings at construction Co-authored-by: Cursor --- columnar/benches/common.rs | 2 +- src/core/tests.rs | 6 ++++-- src/indexer/doc_id_mapping.rs | 29 +++++++++++++++++++++-------- src/indexer/segment_writer.rs | 12 +----------- 4 files changed, 27 insertions(+), 22 deletions(-) diff --git a/columnar/benches/common.rs b/columnar/benches/common.rs index b087b0ac8..1b9fb57fe 100644 --- a/columnar/benches/common.rs +++ b/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 = Vec::new(); - columnar_writer.serialize(num_docs, &mut wrt).unwrap(); + columnar_writer.serialize(num_docs, None, &mut wrt).unwrap(); ColumnarReader::open(wrt).unwrap() } diff --git a/src/core/tests.rs b/src/core/tests.rs index b5d56891a..54cf08006 100644 --- a/src/core/tests.rs +++ b/src/core/tests.rs @@ -306,8 +306,10 @@ fn test_single_segment_index_writer_with_doc_id_mapping() -> crate::Result<()> { let text_field = schema_builder.add_text_field("text", TEXT | STORED); let schema = schema_builder.build(); let directory = RamDirectory::default(); - let mut settings = IndexSettings::default(); - settings.manual_doc_id_mapping = true; + let settings = IndexSettings { + manual_doc_id_mapping: true, + ..Default::default() + }; let mut single_segment_index_writer = Index::builder() .schema(schema) .settings(settings) diff --git a/src/indexer/doc_id_mapping.rs b/src/indexer/doc_id_mapping.rs index 336523726..feca0383e 100644 --- a/src/indexer/doc_id_mapping.rs +++ b/src/indexer/doc_id_mapping.rs @@ -1,7 +1,6 @@ //! 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::ReadOnlyBitSet; +use common::{BitSet, ReadOnlyBitSet}; use super::SegmentWriter; use crate::schema::{Field, Schema}; @@ -71,11 +70,25 @@ pub struct DocIdMapping { } impl DocIdMapping { + /// Creates a `DocIdMapping` from a mapping of new doc ids to old doc ids, with permutation validation. + /// I.e. every new doc id must appear exactly once in the mapping, and no new doc id may be greater than the length of the mapping. + pub fn new_permutation(new_doc_id_to_old: Vec) -> crate::Result { + // 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 seen_doc_ids = BitSet::with_max_value(max_doc); + for new_doc_id in new_doc_id_to_old.iter().copied() { + if new_doc_id >= max_doc || !seen_doc_ids.insert(new_doc_id) { + return Err(TantivyError::InvalidArgument( + "Mapping must be a permutation of the segment doc ids".to_string(), + )); + } + } + + Ok(Self::from_new_id_to_old_id(new_doc_id_to_old)) + } + /// Creates a `DocIdMapping` from a mapping of new doc ids to old doc ids. - /// - /// The caller MUST ensure that `new_doc_id_to_old` is a permutation of the - /// segment's old doc ids, with every old doc id appearing exactly once. - pub fn from_new_id_to_old_id(new_doc_id_to_old: Vec) -> Self { + pub(crate) fn from_new_id_to_old_id(new_doc_id_to_old: Vec) -> Self { let max_doc = new_doc_id_to_old.len(); let old_max_doc = new_doc_id_to_old .iter() @@ -99,8 +112,8 @@ impl DocIdMapping { } /// iterate over old doc_ids in order of the new doc_ids - pub(crate) fn iter_old_doc_ids(&self) -> impl Iterator + Clone + '_ { - self.new_doc_id_to_old.iter().cloned() + pub(crate) fn iter_old_doc_ids(&self) -> impl Iterator + '_ { + self.new_doc_id_to_old.iter().copied() } /// returns the new doc_ids in order of the old doc_ids diff --git a/src/indexer/segment_writer.rs b/src/indexer/segment_writer.rs index a2ab7de44..41572174f 100644 --- a/src/indexer/segment_writer.rs +++ b/src/indexer/segment_writer.rs @@ -1,5 +1,5 @@ use columnar::MonotonicallyMappableToU64; -use common::{BitSet, JsonPathWriter}; +use common::JsonPathWriter; use itertools::Itertools; use tokenizer_api::BoxTokenStream; @@ -179,16 +179,6 @@ impl SegmentWriter { ))); } - // Check that the mapping is a permutation of the segment doc ids. - let mut seen_doc_ids = BitSet::with_max_value(self.max_doc); - for old_doc_id in mapping.iter_old_doc_ids() { - if old_doc_id >= self.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(), - )); - } - } - self.finalize_inner(Some(mapping)) }