diff --git a/columnar/src/column_values/u128_based/compact_space/mod.rs b/columnar/src/column_values/u128_based/compact_space/mod.rs index 9fa36e0201..3d58574639 100644 --- a/columnar/src/column_values/u128_based/compact_space/mod.rs +++ b/columnar/src/column_values/u128_based/compact_space/mod.rs @@ -529,6 +529,7 @@ impl CompactSpaceDecompressor { #[cfg(test)] mod tests { + use common::file_slice::FileSlice; use itertools::Itertools; use super::*; diff --git a/src/directory/mmap_directory/mod.rs b/src/directory/mmap_directory/mod.rs index 60ef82b304..633e3ee758 100644 --- a/src/directory/mmap_directory/mod.rs +++ b/src/directory/mmap_directory/mod.rs @@ -675,9 +675,10 @@ mod tests { reader.reload().unwrap(); let num_segments = reader.searcher().segment_readers().len(); assert!(num_segments <= 4); - let num_components_except_deletes_and_tempstore = - crate::index::SegmentComponent::iterator().len() - 2; - let max_num_mmapped = num_components_except_deletes_and_tempstore * num_segments; + // Built-in segment components excluding the delete file (no deletes here): + // postings, positions, fast fields, field norms, terms, store, temp store. + let max_components_per_segment = 7; + let max_num_mmapped = max_components_per_segment * num_segments; assert_eventually(|| { let num_mmapped = mmap_directory.get_cache_info().mmapped.len(); if num_mmapped > max_num_mmapped { diff --git a/src/error.rs b/src/error.rs index 39fee20d78..8b010cf67e 100644 --- a/src/error.rs +++ b/src/error.rs @@ -110,6 +110,26 @@ pub enum TantivyError { #[error("Deserialize error: {0}")] /// An error occurred while attempting to deserialize a document. DeserializeError(DeserializeError), + /// A segment plugin required by the index is not registered. + #[error( + "Missing required segment plugin(s) for extension(s): {0}. Re-register via \ + Index::register_plugin before writing, merging, or garbage collecting." + )] + MissingPlugin(String), + /// Multiple registered segment plugins claim ownership of the same extension(s). + #[error( + "Conflicting segment plugins claim the same extension(s): {0}. Each extension must be \ + owned by exactly one registered plugin." + )] + ConflictingPlugins(String), + /// A segment plugin was registered on a non-empty index whose existing segments do not + /// contain its extension(s). + #[error( + "Segment plugin(s) for extension(s) {0} were registered on a non-empty index whose \ + existing segments do not contain them. The plugin set is fixed once an index has \ + segments; create a new index to change it." + )] + UnexpectedPlugin(String), } impl From for TantivyError { diff --git a/src/fastfield/mod.rs b/src/fastfield/mod.rs index a5060f8be8..d981d6508d 100644 --- a/src/fastfield/mod.rs +++ b/src/fastfield/mod.rs @@ -25,6 +25,8 @@ use columnar::MonotonicallyMappableToU64; pub use self::alive_bitset::{intersect_alive_bitsets, write_alive_bitset, AliveBitSet}; pub use self::error::{FastFieldNotAvailableError, Result}; pub use self::facet_reader::FacetReader; +pub mod plugin; +pub use self::plugin::{FastFieldsPlugin, FastFieldsPluginWriter}; pub use self::readers::FastFieldReaders; pub use self::writer::FastFieldsWriter; use crate::schema::Type; diff --git a/src/fastfield/plugin.rs b/src/fastfield/plugin.rs new file mode 100644 index 0000000000..091f1987f7 --- /dev/null +++ b/src/fastfield/plugin.rs @@ -0,0 +1,182 @@ +//! FastFields as a [`SegmentPlugin`] implementation. +//! +//! This wraps the existing `FastFieldsWriter` and `FastFieldReaders` types behind +//! the plugin interface so that fast fields participate in the unified plugin lifecycle. + +use std::any::Any; +use std::collections::BTreeMap; + +use columnar::{ + ColumnType, ColumnarReader, MergeRowOrder, RowAddr, ShuffleMergeOrder, StackMergeOrder, +}; +use common::TerminatingWrite; +use measure_time::debug_time; + +use crate::directory::{Directory, WritePtr}; +use crate::fastfield::{FastFieldReaders, FastFieldsWriter}; +use crate::index::{SegmentComponent, SegmentReader}; +use crate::indexer::doc_id_mapping::{DocIdMapping, MappingType, SegmentDocIdMapping}; +use crate::plugin::{PluginMergeContext, PluginWriter, PluginWriterContext, SegmentPlugin}; +use crate::schema::document::Document; +use crate::schema::{value_type_to_column_type, Schema}; +use crate::space_usage::{ComponentSpaceUsage, FAST_FIELDS}; +use crate::Segment; + +pub struct FastFieldsPlugin; + +impl SegmentPlugin for FastFieldsPlugin { + fn extensions(&self) -> &[&str] { + &["fast"] + } + + fn create_writer(&self, _ctx: &PluginWriterContext) -> crate::Result> { + unimplemented!("FastFieldsPlugin is a built-in; use FastFieldsPluginWriter::new") + } + + fn merge(&self, ctx: PluginMergeContext) -> crate::Result<()> { + debug_time!("write-fast-fields"); + let path = ctx + .target_segment + .relative_path(SegmentComponent::FastFields); + let mut fast_field_wrt: WritePtr = + ctx.target_segment.index().directory().open_write(&path)?; + + let required_columns = extract_fast_field_required_columns(ctx.schema); + let columnars: Vec<&ColumnarReader> = ctx + .readers + .iter() + .map(|reader| reader.fast_fields().columnar()) + .collect(); + + // Clone the doc_id_mapping since convert_to_merge_order consumes it by value + let doc_id_mapping = ctx.doc_id_mapping.clone(); + let merge_row_order = convert_to_merge_order(&columnars[..], doc_id_mapping); + + columnar::merge_columnar( + &columnars[..], + &required_columns, + merge_row_order, + &mut fast_field_wrt, + )?; + + fast_field_wrt.terminate()?; + Ok(()) + } + + fn space_usage( + &self, + segment_reader: &SegmentReader, + ) -> crate::Result> { + let file = segment_reader.open_read(SegmentComponent::FastFields)?; + let readers = FastFieldReaders::open(file, segment_reader.schema().clone()) + .map_err(|e| crate::TantivyError::InternalError(e.to_string()))?; + let usage = readers.space_usage()?; + Ok(BTreeMap::from([( + FAST_FIELDS.to_string(), + ComponentSpaceUsage::PerField(usage), + )])) + } +} + +pub struct FastFieldsPluginWriter { + pub writer: FastFieldsWriter, + fast_field_write: WritePtr, +} + +impl FastFieldsPluginWriter { + pub(crate) fn new(ctx: &PluginWriterContext) -> crate::Result { + let index = ctx.segment.index(); + let tokenizer_manager = index.fast_field_tokenizer().clone(); + let writer = FastFieldsWriter::from_schema_and_tokenizer_manager( + &ctx.segment.schema(), + tokenizer_manager, + )?; + + let path = ctx.segment.relative_path(SegmentComponent::FastFields); + let fast_field_write = index.directory().open_write(&path)?; + + Ok(FastFieldsPluginWriter { + writer, + fast_field_write, + }) + } + + /// Generic, zero-copy document ingestion. Called directly by the `SegmentWriter` for + /// this built-in, so a custom `Document` is indexed in place without materializing a + /// `TantivyDocument`. + pub(crate) fn index_document(&mut self, doc: &D) -> crate::Result<()> { + self.writer.add_document(doc) + } + + pub fn writer_mut(&mut self) -> &mut FastFieldsWriter { + &mut self.writer + } + + pub fn writer(&self) -> &FastFieldsWriter { + &self.writer + } +} + +impl PluginWriter for FastFieldsPluginWriter { + fn serialize( + mut self: Box, + _segment: &Segment, + doc_id_map: Option<&DocIdMapping>, + ) -> crate::Result<()> { + self.writer + .serialize(&mut self.fast_field_write, doc_id_map) + .map_err(|e| crate::TantivyError::InternalError(e.to_string()))?; + self.fast_field_write.terminate()?; + Ok(()) + } + + fn mem_usage(&self) -> usize { + self.writer.mem_usage() + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } +} + +// --- Helper functions moved from merger.rs --- + +fn convert_to_merge_order( + columnars: &[&ColumnarReader], + doc_id_mapping: SegmentDocIdMapping, +) -> MergeRowOrder { + match doc_id_mapping.mapping_type() { + MappingType::Stacked => MergeRowOrder::Stack(StackMergeOrder::stack(columnars)), + MappingType::StackedWithDeletes | MappingType::Shuffled => { + let new_row_id_to_old_row_id: Vec = doc_id_mapping + .new_doc_id_to_old_doc_addr + .into_iter() + .map(|doc_addr| RowAddr { + segment_ord: doc_addr.segment_ord, + row_id: doc_addr.doc_id, + }) + .collect(); + MergeRowOrder::Shuffled(ShuffleMergeOrder { + new_row_id_to_old_row_id, + alive_bitsets: doc_id_mapping.alive_bitsets, + }) + } + } +} + +fn extract_fast_field_required_columns(schema: &Schema) -> Vec<(String, ColumnType)> { + schema + .fields() + .map(|(_, field_entry)| field_entry) + .filter(|field_entry| field_entry.is_fast()) + .filter_map(|field_entry| { + let column_name = field_entry.name().to_string(); + let column_type = value_type_to_column_type(field_entry.field_type().value_type())?; + Some((column_name, column_type)) + }) + .collect() +} diff --git a/src/fastfield/writer.rs b/src/fastfield/writer.rs index 2f4196078e..4eab0c1908 100644 --- a/src/fastfield/writer.rs +++ b/src/fastfield/writer.rs @@ -234,7 +234,7 @@ impl FastFieldsWriter { /// Serializes all of the `FastFieldWriter`s by pushing them in /// order to the fast field serializer. pub fn serialize( - mut self, + &mut self, wrt: &mut dyn io::Write, doc_id_map_opt: Option<&DocIdMapping>, ) -> io::Result<()> { diff --git a/src/index/index.rs b/src/index/index.rs index 2cb466a4ee..cc933b8f3f 100644 --- a/src/index/index.rs +++ b/src/index/index.rs @@ -1,8 +1,9 @@ -use std::collections::HashSet; +use std::collections::{BTreeSet, HashSet}; use std::fmt; #[cfg(feature = "mmap")] use std::path::Path; use std::path::PathBuf; +use std::sync::Arc; use std::thread::available_parallelism; use super::segment::Segment; @@ -14,15 +15,20 @@ use crate::directory::error::OpenReadError; use crate::directory::MmapDirectory; use crate::directory::{Directory, ManagedDirectory, RamDirectory, INDEX_WRITER_LOCK}; use crate::error::{DataCorruption, TantivyError}; -use crate::index::{IndexMeta, SegmentId, SegmentMeta, SegmentMetaInventory}; +use crate::fastfield::FastFieldsPlugin; +use crate::index::{ + IndexMeta, InvertedIndexPlugin, SegmentComponent, SegmentId, SegmentMeta, SegmentMetaInventory, +}; use crate::indexer::index_writer::{ IndexWriterOptions, MAX_NUM_THREAD, MEMORY_BUDGET_NUM_BYTES_MIN, }; use crate::indexer::segment_updater::save_metas; use crate::indexer::{IndexWriter, SingleSegmentIndexWriter}; +use crate::plugin::SegmentPlugin; use crate::reader::{IndexReader, IndexReaderBuilder}; use crate::schema::document::Document; use crate::schema::{Field, FieldType, Schema, Type}; +use crate::store::StorePlugin; use crate::tokenizer::{TextAnalyzer, TokenizerManager}; use crate::SegmentReader; @@ -48,6 +54,49 @@ fn load_metas( .map_err(From::from) } +/// The built-in plugins, in write order. +pub(crate) fn builtin_plugins() -> Vec> { + vec![ + Arc::new(InvertedIndexPlugin), + Arc::new(FastFieldsPlugin), + Arc::new(StorePlugin), + ] +} + +/// The files owned by the given segments: one per plugin extension (the built-in +/// extensions plus the index's required custom ones), plus each segment's delete bitset +/// and — while it is still being written — its temp store. +/// +/// Callers pass the segment set they care about (e.g. all tracked segments for garbage +/// collection, or just the committed ones) and the index's persisted custom extensions. +pub(crate) fn list_segment_files( + segment_metas: &[SegmentMeta], + persisted_custom_extensions: &[String], +) -> HashSet { + let extensions: Vec = builtin_plugins() + .iter() + .flat_map(|plugin| plugin.extensions().iter().map(|ext| ext.to_string())) + .chain(persisted_custom_extensions.iter().cloned()) + .collect(); + let mut files = HashSet::new(); + for segment_meta in segment_metas { + for ext in &extensions { + files.insert(segment_meta.relative_path(SegmentComponent::Custom(ext.clone()))); + } + files.insert(segment_meta.relative_path(SegmentComponent::Delete)); + if segment_meta.include_temp_store() { + files.insert(segment_meta.relative_path(SegmentComponent::TempStore)); + } + } + files +} + +/// File extensions for built-in segment components that no plugin owns: the temp +/// store (`.store.temp`) and the delete bitset (`..del`). +/// Both spellings of each are reserved — the bare [`SegmentComponent`] name and the +/// on-disk extension — so a custom plugin cannot claim either and contend for the file. +const RESERVED_NON_PLUGIN_EXTENSIONS: &[&str] = &["temp", "store.temp", "del"]; + /// Save the index meta file. /// This operation is atomic : /// Either @@ -58,11 +107,18 @@ fn load_metas( fn save_new_metas( schema: Schema, index_settings: IndexSettings, + plugins: &[Arc], directory: &dyn Directory, ) -> crate::Result<()> { + let persisted_custom_extensions: Vec = plugins + .iter() + .flat_map(|plugin| plugin.extensions().iter().copied()) + .map(str::to_string) + .collect(); save_metas( &IndexMeta { index_settings, + persisted_custom_extensions, segments: Vec::new(), schema, opstamp: 0u64, @@ -106,6 +162,7 @@ pub struct IndexBuilder { index_settings: IndexSettings, tokenizer_manager: TokenizerManager, fast_field_tokenizer_manager: TokenizerManager, + custom_plugins: Vec>, } impl Default for IndexBuilder { fn default() -> Self { @@ -120,6 +177,7 @@ impl IndexBuilder { index_settings: IndexSettings::default(), tokenizer_manager: TokenizerManager::default(), fast_field_tokenizer_manager: TokenizerManager::default(), + custom_plugins: Vec::new(), } } @@ -149,6 +207,12 @@ impl IndexBuilder { self } + #[must_use] + pub fn register_plugin(mut self, plugin: Arc) -> Self { + self.custom_plugins.push(plugin); + self + } + /// Creates a new index using the [`RamDirectory`]. /// /// The index will be allocated in anonymous memory. @@ -223,6 +287,7 @@ impl IndexBuilder { let mut index = Index::open(dir)?; index.set_tokenizers(self.tokenizer_manager.clone()); if index.schema() == self.get_expect_schema()? { + index.custom_plugins.extend(self.custom_plugins); Ok(index) } else { Err(TantivyError::SchemaError( @@ -275,13 +340,14 @@ impl IndexBuilder { /// Creates a new index given an implementation of the trait `Directory`. /// /// If a directory previously existed, it will be erased. - fn create>>(self, dir: T) -> crate::Result { + pub(crate) fn create>>(self, dir: T) -> crate::Result { self.validate()?; let dir = dir.into(); let directory = ManagedDirectory::wrap(dir)?; save_new_metas( self.get_expect_schema()?, self.index_settings.clone(), + &self.custom_plugins, &directory, )?; let mut metas = IndexMeta::with_schema(self.get_expect_schema()?); @@ -289,6 +355,7 @@ impl IndexBuilder { let mut index = Index::open_from_metas(directory, &metas, SegmentMetaInventory::default()); index.set_tokenizers(self.tokenizer_manager); index.set_fast_field_tokenizers(self.fast_field_tokenizer_manager); + index.custom_plugins.extend(self.custom_plugins); Ok(index) } } @@ -303,6 +370,7 @@ pub struct Index { tokenizers: TokenizerManager, fast_field_tokenizers: TokenizerManager, inventory: SegmentMetaInventory, + custom_plugins: Vec>, } impl Index { @@ -422,6 +490,7 @@ impl Index { fast_field_tokenizers: TokenizerManager::default(), executor: Executor::single_thread(), inventory, + custom_plugins: Vec::new(), } } @@ -430,6 +499,21 @@ impl Index { self.tokenizers = tokenizers; } + /// Re-attach a custom segment plugin after the index has been opened. + /// + /// Plugin registration is not persisted, so `Index::open` restores only the + /// built-in plugins. Custom plugins must be re-registered before any write, + /// merge, or garbage collection, the same way custom tokenizers are + /// re-attached via [`Index::set_tokenizers`]. + /// + /// The index records, at creation, the custom extensions it requires, so these + /// operations fail closed with [`TantivyError::MissingPlugin`] when an owning + /// plugin is not registered — rather than dropping its data from a merge or + /// deleting its files during garbage collection. + pub fn register_plugin(&mut self, plugin: Arc) { + self.custom_plugins.push(plugin); + } + /// Accessor for the tokenizer manager. pub fn tokenizers(&self) -> &TokenizerManager { &self.tokenizers @@ -586,7 +670,7 @@ impl Index { ) })?; - IndexWriter::new(self, options, directory_lock) + IndexWriter::new(self.clone(), options, directory_lock) } /// Open a new index writer. Attempts to acquire a lockfile. @@ -653,6 +737,85 @@ impl Index { self.writer_with_num_threads(num_threads, memory_budget_in_bytes) } + /// The custom (non-built-in) plugins registered on this index. + pub fn custom_plugins(&self) -> &[Arc] { + &self.custom_plugins + } + + /// The sorted set of extensions claimed by the registered custom (non-built-in) plugins. + pub(crate) fn custom_extensions(&self) -> BTreeSet<&str> { + self.custom_plugins + .iter() + .flat_map(|plugin| plugin.extensions().iter().copied()) + .collect() + } + + /// All plugins that participate in writing and merging, in write order: the built-ins + /// (field norms, postings, fast fields, store) followed by the registered custom + /// plugins. Built-ins are not stored; they are produced fresh here from `builtin_plugins`. + pub fn all_plugins(&self) -> impl Iterator> + '_ { + builtin_plugins() + .into_iter() + .chain(self.custom_plugins.iter().cloned()) + } + + /// Validates that the registered plugin set matches the index's recorded plugin set: + /// no two plugins claim the same extension, none claim a reserved one, and the + /// registered custom plugins exactly match the extensions the index was created with. + /// The plugin set is part of the index's structure, like the schema, and is fixed at + /// creation. + pub(crate) fn validate_plugins(&self) -> crate::Result<()> { + // No two plugins may claim the same extension, and none may claim a reserved one. + let all_plugins: Vec<_> = self.all_plugins().collect(); + let mut claimed_extensions: HashSet<&str> = + RESERVED_NON_PLUGIN_EXTENSIONS.iter().copied().collect(); + let mut conflicting: std::collections::BTreeSet = std::collections::BTreeSet::new(); + for plugin in &all_plugins { + for &ext in plugin.extensions() { + if !claimed_extensions.insert(ext) { + conflicting.insert(ext.to_string()); + } + } + } + if !conflicting.is_empty() { + return Err(crate::TantivyError::ConflictingPlugins( + conflicting.into_iter().collect::>().join(", "), + )); + } + + // The index persists, at creation, the custom plugin extensions it requires; that + // index-wide record is the single source of truth for the plugin set. The currently + // registered custom plugins must match it exactly. + let metas = self.load_metas()?; + let persisted_custom_extensions: BTreeSet<&str> = metas + .persisted_custom_extensions + .iter() + .map(String::as_str) + .collect(); + let registered_custom_extensions = self.custom_extensions(); + + // Persisted but not registered: writing or merging would silently drop that plugin's data. + let missing: Vec<&str> = persisted_custom_extensions + .difference(®istered_custom_extensions) + .copied() + .collect(); + if !missing.is_empty() { + return Err(crate::TantivyError::MissingPlugin(missing.join(", "))); + } + + // Registered but not persisted: the plugin set is fixed at creation, so a plugin not + // persisted there cannot be added to an existing index. + let unexpected: Vec<&str> = registered_custom_extensions + .difference(&persisted_custom_extensions) + .copied() + .collect(); + if !unexpected.is_empty() { + return Err(crate::TantivyError::UnexpectedPlugin(unexpected.join(", "))); + } + + Ok(()) + } + /// Accessor to the index settings pub fn settings(&self) -> &IndexSettings { &self.settings @@ -720,11 +883,9 @@ impl Index { /// Returns the set of corrupted files pub fn validate_checksum(&self) -> crate::Result> { let managed_files = self.directory.list_managed_files(); - let active_segments_files: HashSet = self - .searchable_segment_metas()? - .iter() - .flat_map(|segment_meta| segment_meta.list_files()) - .collect(); + let metas = self.load_metas()?; + let active_segments_files = + list_segment_files(&metas.segments, &metas.persisted_custom_extensions); let active_existing_files: HashSet<&PathBuf> = active_segments_files.intersection(&managed_files).collect(); diff --git a/src/index/index_meta.rs b/src/index/index_meta.rs index 28a27c72e8..b10d05a5c8 100644 --- a/src/index/index_meta.rs +++ b/src/index/index_meta.rs @@ -1,4 +1,3 @@ -use std::collections::HashSet; use std::fmt; use std::path::PathBuf; use std::sync::atomic::AtomicBool; @@ -103,28 +102,10 @@ impl SegmentMeta { .unwrap_or(0u32) } - /// Returns the list of files that - /// are required for the segment meta. - /// Note: Some of the returned files may not exist depending on the state of the segment. - /// - /// This is useful as the way tantivy removes files - /// is by removing all files that have been created by tantivy - /// and are not used by any segment anymore. - pub fn list_files(&self) -> HashSet { - if self - .tracked + pub(crate) fn include_temp_store(&self) -> bool { + self.tracked .include_temp_doc_store .load(std::sync::atomic::Ordering::Relaxed) - { - SegmentComponent::iterator() - .map(|component| self.relative_path(*component)) - .collect::>() - } else { - SegmentComponent::iterator() - .filter(|comp| *comp != &SegmentComponent::TempStore) - .map(|component| self.relative_path(*component)) - .collect::>() - } } /// Returns the relative path of a component of our segment. @@ -142,6 +123,7 @@ impl SegmentMeta { SegmentComponent::FastFields => ".fast".to_string(), SegmentComponent::FieldNorms => ".fieldnorm".to_string(), SegmentComponent::Delete => format!(".{}.del", self.delete_opstamp().unwrap_or(0)), + SegmentComponent::Custom(ext) => format!(".{ext}"), }); PathBuf::from(path) } @@ -324,6 +306,16 @@ pub struct IndexMeta { /// `IndexSettings` to configure index options. #[serde(default)] pub index_settings: IndexSettings, + /// Custom (non-built-in) plugin extensions this index was created with. + /// + /// Fixed at index creation. This is the single source of truth for which + /// custom plugins the index requires: garbage collection keeps their files, + /// and the writer/merger fails closed if an owning plugin is not registered. + /// An index created before this field existed deserializes to empty, i.e. + /// built-in plugins only. + #[serde(default)] + #[serde(skip_serializing_if = "Vec::is_empty")] + pub persisted_custom_extensions: Vec, /// List of `SegmentMeta` information associated with each finalized segment of the index. pub segments: Vec, /// Index `Schema` @@ -344,6 +336,8 @@ struct UntrackedIndexMeta { pub segments: Vec, #[serde(default)] pub index_settings: IndexSettings, + #[serde(default)] + pub persisted_custom_extensions: Vec, pub schema: Schema, pub opstamp: Opstamp, #[serde(skip_serializing_if = "Option::is_none")] @@ -354,6 +348,7 @@ impl UntrackedIndexMeta { pub fn track(self, inventory: &SegmentMetaInventory) -> IndexMeta { IndexMeta { index_settings: self.index_settings, + persisted_custom_extensions: self.persisted_custom_extensions, segments: self .segments .into_iter() @@ -375,6 +370,7 @@ impl IndexMeta { pub fn with_schema(schema: Schema) -> IndexMeta { IndexMeta { index_settings: IndexSettings::default(), + persisted_custom_extensions: Vec::new(), segments: vec![], schema, opstamp: 0u64, @@ -429,6 +425,7 @@ mod tests { }), ..Default::default() }, + persisted_custom_extensions: Vec::new(), segments: Vec::new(), schema, opstamp: 0u64, @@ -466,6 +463,7 @@ mod tests { docstore_blocksize: 1_000_000, docstore_compress_dedicated_thread: true, }, + persisted_custom_extensions: Vec::new(), segments: Vec::new(), schema, opstamp: 0u64, diff --git a/src/index/inverted_index_plugin.rs b/src/index/inverted_index_plugin.rs new file mode 100644 index 0000000000..dacbced127 --- /dev/null +++ b/src/index/inverted_index_plugin.rs @@ -0,0 +1,686 @@ +//! The inverted index as a [`SegmentPlugin`] implementation. +//! +//! Field norms, the term dictionary, postings, and positions form a single subsystem: +//! field norms are produced by the same tokenization pass that feeds the postings, and +//! postings serialization reads the field norms back. This plugin therefore owns all four +//! files (`fieldnorm`, `term`, `idx`, `pos`) and drives them together, both while indexing +//! (`add_document`) and while merging. + +use std::any::Any; +use std::collections::BTreeMap; +use std::sync::Arc; + +use columnar::MonotonicallyMappableToU64; +use common::JsonPathWriter; +use itertools::Itertools; +use measure_time::debug_time; +use tokenizer_api::BoxTokenStream; + +use crate::directory::{CompositeFile, Directory}; +use crate::docset::{DocSet, TERMINATED}; +use crate::error::DataCorruption; +use crate::fieldnorm::{FieldNormReader, FieldNormReaders, FieldNormsSerializer, FieldNormsWriter}; +use crate::index::{Segment, SegmentComponent, SegmentReader}; +use crate::indexer::doc_id_mapping::{DocIdMapping, SegmentDocIdMapping}; +use crate::indexer::indexing_term::IndexingTerm; +use crate::json_utils::{index_json_value, IndexingPositionsPerPath}; +use crate::plugin::{PluginMergeContext, PluginWriter, PluginWriterContext, SegmentPlugin}; +use crate::postings::{ + compute_initial_table_size, serialize_postings, IndexingContext, IndexingPosition, + InvertedIndexSerializer, PerFieldPostingsWriter, Postings, PostingsWriter, SegmentPostings, +}; +use crate::schema::document::{Document, Value}; +use crate::schema::{Field, FieldType, Schema, DATE_TIME_PRECISION_INDEXED}; +use crate::space_usage::{ComponentSpaceUsage, FIELDNORMS, POSITIONS, POSTINGS, TERMDICT}; +use crate::termdict::{TermMerger, TermOrdinal}; +use crate::tokenizer::{FacetTokenizer, PreTokenizedStream, TextAnalyzer, Tokenizer}; +use crate::{DocId, InvertedIndexReader, TantivyError}; + +pub struct InvertedIndexPlugin; + +impl SegmentPlugin for InvertedIndexPlugin { + fn extensions(&self) -> &[&str] { + &["fieldnorm", "term", "idx", "pos"] + } + + fn create_writer(&self, _ctx: &PluginWriterContext) -> crate::Result> { + unimplemented!("InvertedIndexPlugin is a built-in; use InvertedIndexPluginWriter::new") + } + + fn merge(&self, ctx: PluginMergeContext) -> crate::Result<()> { + // Field norms first: postings merge reads them back from the target segment. + merge_fieldnorms(&ctx)?; + + debug_time!("write-postings"); + debug!("write-postings"); + let target_segment = ctx.target_segment; + let mut serializer = InvertedIndexSerializer::open(target_segment)?; + let fieldnorm_data = target_segment.open_read(SegmentComponent::FieldNorms)?; + let fieldnorm_readers = FieldNormReaders::open(fieldnorm_data)?; + write_postings_merge( + ctx.readers, + ctx.schema, + &mut serializer, + fieldnorm_readers, + ctx.doc_id_mapping, + )?; + serializer.close()?; + Ok(()) + } + + fn space_usage( + &self, + segment_reader: &SegmentReader, + ) -> crate::Result> { + let schema = segment_reader.schema(); + + let fieldnorms = + FieldNormReaders::open(segment_reader.open_read(SegmentComponent::FieldNorms)?)? + .space_usage(schema); + let termdict = CompositeFile::open(&segment_reader.open_read(SegmentComponent::Terms)?)? + .space_usage(schema); + let postings = CompositeFile::open(&segment_reader.open_read(SegmentComponent::Postings)?)? + .space_usage(schema); + let positions = match segment_reader.open_read(SegmentComponent::Positions) { + Ok(file) => CompositeFile::open(&file)?.space_usage(schema), + Err(_) => CompositeFile::empty().space_usage(schema), + }; + Ok(BTreeMap::from([ + ( + FIELDNORMS.to_string(), + ComponentSpaceUsage::PerField(fieldnorms), + ), + ( + TERMDICT.to_string(), + ComponentSpaceUsage::PerField(termdict), + ), + ( + POSTINGS.to_string(), + ComponentSpaceUsage::PerField(postings), + ), + ( + POSITIONS.to_string(), + ComponentSpaceUsage::PerField(positions), + ), + ])) + } +} + +// --- Plugin writer --- + +pub struct InvertedIndexPluginWriter { + schema: Schema, + per_field_postings_writers: PerFieldPostingsWriter, + per_field_text_analyzers: Vec, + fieldnorms_writer: FieldNormsWriter, + term_buffer: IndexingTerm, + json_path_writer: JsonPathWriter, + json_positions_per_path: IndexingPositionsPerPath, + ctx: IndexingContext, + postings_serializer: InvertedIndexSerializer, + fieldnorm_serializer: FieldNormsSerializer, + max_doc: DocId, +} + +impl InvertedIndexPluginWriter { + pub(crate) fn new(ctx: &PluginWriterContext) -> crate::Result { + let segment = ctx.segment; + let schema = segment.schema(); + let tokenizer_manager = segment.index().tokenizers().clone(); + + let per_field_text_analyzers = schema + .fields() + .map(|(_, field_entry)| { + let text_options = match field_entry.field_type() { + FieldType::Str(ref text_options) => text_options.get_indexing_options(), + FieldType::JsonObject(ref json_object_options) => { + json_object_options.get_text_indexing_options() + } + _ => None, + }; + let tokenizer_name = text_options + .map(|text_index_option| text_index_option.tokenizer()) + .unwrap_or("default"); + tokenizer_manager.get(tokenizer_name).ok_or_else(|| { + TantivyError::SchemaError(format!( + "Error getting tokenizer for field: {}", + field_entry.name() + )) + }) + }) + .collect::, _>>()?; + + let fieldnorm_path = segment.relative_path(SegmentComponent::FieldNorms); + let fieldnorm_write = segment.index().directory().open_write(&fieldnorm_path)?; + + let table_size = compute_initial_table_size(ctx.memory_budget_in_bytes)?; + Ok(InvertedIndexPluginWriter { + per_field_postings_writers: PerFieldPostingsWriter::for_schema(&schema), + per_field_text_analyzers, + fieldnorms_writer: FieldNormsWriter::for_schema(&schema), + term_buffer: IndexingTerm::with_capacity(16), + json_path_writer: JsonPathWriter::default(), + json_positions_per_path: IndexingPositionsPerPath::default(), + ctx: IndexingContext::new(table_size), + postings_serializer: InvertedIndexSerializer::open(segment)?, + fieldnorm_serializer: FieldNormsSerializer::from_write(fieldnorm_write)?, + schema, + max_doc: 0, + }) + } + + /// Generic, zero-copy document ingestion. The `SegmentWriter` calls this directly on the + /// built-in inverted-index writer, so a custom `Document` is tokenized and indexed in + /// place — no `TantivyDocument` is materialized. + pub(crate) fn index_document( + &mut self, + doc_id: DocId, + doc: &D, + ) -> crate::Result<()> { + let vals_grouped_by_field = doc + .iter_fields_and_values() + .sorted_by_key(|(field, _)| *field) + .chunk_by(|(field, _)| *field); + + for (field, field_values) in &vals_grouped_by_field { + let values = field_values.map(|el| el.1); + + let field_entry = self.schema.get_field_entry(field); + let make_schema_error = || { + TantivyError::SchemaError(format!( + "Expected a {:?} for field {:?}", + field_entry.field_type().value_type(), + field_entry.name() + )) + }; + if !field_entry.is_indexed() { + continue; + } + + let (term_buffer, ctx) = (&mut self.term_buffer, &mut self.ctx); + let postings_writer: &mut dyn PostingsWriter = + self.per_field_postings_writers.get_for_field_mut(field); + term_buffer.clear_with_field(field); + + match field_entry.field_type() { + FieldType::Facet(_) => { + let mut facet_tokenizer = FacetTokenizer::default(); // this can be global + for value in values { + let value = value.as_value(); + + let facet_str = value.as_facet().ok_or_else(make_schema_error)?; + let mut facet_tokenizer = facet_tokenizer.token_stream(facet_str); + let mut indexing_position = IndexingPosition::default(); + postings_writer.index_text( + doc_id, + &mut facet_tokenizer, + term_buffer, + ctx, + &mut indexing_position, + ); + } + } + FieldType::Str(_) => { + let mut indexing_position = IndexingPosition::default(); + for value in values { + let value = value.as_value(); + + let mut token_stream = if let Some(text) = value.as_str() { + let text_analyzer = + &mut self.per_field_text_analyzers[field.field_id() as usize]; + text_analyzer.token_stream(text) + } else if let Some(tok_str) = value.into_pre_tokenized_text() { + BoxTokenStream::new(PreTokenizedStream::from(*tok_str.clone())) + } else { + continue; + }; + + assert!(term_buffer.is_empty()); + postings_writer.index_text( + doc_id, + &mut *token_stream, + term_buffer, + ctx, + &mut indexing_position, + ); + } + if field_entry.has_fieldnorms() { + self.fieldnorms_writer + .record(doc_id, field, indexing_position.num_tokens); + } + } + FieldType::U64(_) => { + let mut num_vals = 0; + for value in values { + let value = value.as_value(); + + num_vals += 1; + let u64_val = value.as_u64().ok_or_else(make_schema_error)?; + term_buffer.set_u64(u64_val); + postings_writer.subscribe(doc_id, 0u32, term_buffer, ctx); + } + if field_entry.has_fieldnorms() { + self.fieldnorms_writer.record(doc_id, field, num_vals); + } + } + FieldType::Date(_) => { + let mut num_vals = 0; + for value in values { + let value = value.as_value(); + + num_vals += 1; + let date_val = value.as_datetime().ok_or_else(make_schema_error)?; + term_buffer + .set_u64(date_val.truncate(DATE_TIME_PRECISION_INDEXED).to_u64()); + postings_writer.subscribe(doc_id, 0u32, term_buffer, ctx); + } + if field_entry.has_fieldnorms() { + self.fieldnorms_writer.record(doc_id, field, num_vals); + } + } + FieldType::I64(_) => { + let mut num_vals = 0; + for value in values { + let value = value.as_value(); + + num_vals += 1; + let i64_val = value.as_i64().ok_or_else(make_schema_error)?; + term_buffer.set_i64(i64_val); + postings_writer.subscribe(doc_id, 0u32, term_buffer, ctx); + } + if field_entry.has_fieldnorms() { + self.fieldnorms_writer.record(doc_id, field, num_vals); + } + } + FieldType::F64(_) => { + let mut num_vals = 0; + for value in values { + let value = value.as_value(); + num_vals += 1; + let f64_val = value.as_f64().ok_or_else(make_schema_error)?; + term_buffer.set_f64(f64_val); + postings_writer.subscribe(doc_id, 0u32, term_buffer, ctx); + } + if field_entry.has_fieldnorms() { + self.fieldnorms_writer.record(doc_id, field, num_vals); + } + } + FieldType::Bool(_) => { + let mut num_vals = 0; + for value in values { + let value = value.as_value(); + num_vals += 1; + let bool_val = value.as_bool().ok_or_else(make_schema_error)?; + term_buffer.set_bool(bool_val); + postings_writer.subscribe(doc_id, 0u32, term_buffer, ctx); + } + if field_entry.has_fieldnorms() { + self.fieldnorms_writer.record(doc_id, field, num_vals); + } + } + FieldType::Bytes(_) => { + let mut num_vals = 0; + for value in values { + let value = value.as_value(); + num_vals += 1; + let bytes = value.as_bytes().ok_or_else(make_schema_error)?; + term_buffer.set_bytes(bytes); + postings_writer.subscribe(doc_id, 0u32, term_buffer, ctx); + } + if field_entry.has_fieldnorms() { + self.fieldnorms_writer.record(doc_id, field, num_vals); + } + } + FieldType::JsonObject(json_options) => { + let text_analyzer = + &mut self.per_field_text_analyzers[field.field_id() as usize]; + + self.json_positions_per_path.clear(); + self.json_path_writer + .set_expand_dots(json_options.is_expand_dots_enabled()); + for json_value in values { + self.json_path_writer.clear(); + + index_json_value( + doc_id, + json_value, + text_analyzer, + term_buffer, + &mut self.json_path_writer, + postings_writer, + ctx, + &mut self.json_positions_per_path, + ); + } + } + FieldType::IpAddr(_) => { + let mut num_vals = 0; + for value in values { + let value = value.as_value(); + + num_vals += 1; + let ip_addr = value.as_ip_addr().ok_or_else(make_schema_error)?; + term_buffer.set_ip_addr(ip_addr); + postings_writer.subscribe(doc_id, 0u32, term_buffer, ctx); + } + if field_entry.has_fieldnorms() { + self.fieldnorms_writer.record(doc_id, field, num_vals); + } + } + } + } + self.max_doc = doc_id + 1; + Ok(()) + } +} + +impl PluginWriter for InvertedIndexPluginWriter { + fn serialize( + mut self: Box, + segment: &Segment, + doc_id_map: Option<&DocIdMapping>, + ) -> crate::Result<()> { + // Field norms first: postings serialization reads them back below. + self.fieldnorms_writer.fill_up_to_max_doc(self.max_doc); + self.fieldnorms_writer + .serialize(self.fieldnorm_serializer, doc_id_map) + .map_err(|e| crate::TantivyError::InternalError(e.to_string()))?; + + let fieldnorm_data = segment.open_read(SegmentComponent::FieldNorms)?; + let fieldnorm_readers = FieldNormReaders::open(fieldnorm_data)?; + serialize_postings( + self.ctx, + self.schema, + &self.per_field_postings_writers, + fieldnorm_readers, + doc_id_map, + &mut self.postings_serializer, + )?; + self.postings_serializer.close()?; + Ok(()) + } + + fn mem_usage(&self) -> usize { + self.ctx.mem_usage() + self.fieldnorms_writer.mem_usage() + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } +} + +// --- Field norm merge --- + +fn merge_fieldnorms(ctx: &PluginMergeContext) -> crate::Result<()> { + let path = ctx + .target_segment + .relative_path(SegmentComponent::FieldNorms); + let write = ctx.target_segment.index().directory().open_write(&path)?; + let mut serializer = FieldNormsSerializer::from_write(write)?; + + let schema = ctx.schema; + let fields = FieldNormsWriter::fields_with_fieldnorm(schema); + let max_doc: usize = ctx + .readers + .iter() + .map(|reader| reader.num_docs() as usize) + .sum(); + let mut fieldnorms_data = Vec::with_capacity(max_doc); + + for field in fields { + fieldnorms_data.clear(); + let fieldnorms_readers: Vec = ctx + .readers + .iter() + .map(|reader| reader.get_fieldnorms_reader(field)) + .collect::>()?; + for old_doc_addr in ctx.doc_id_mapping.iter_old_doc_addrs() { + let reader = &fieldnorms_readers[old_doc_addr.segment_ord as usize]; + let fieldnorm_id = reader.fieldnorm_id(old_doc_addr.doc_id); + fieldnorms_data.push(fieldnorm_id); + } + serializer.serialize_field(field, &fieldnorms_data)?; + } + serializer.close()?; + Ok(()) +} + +// --- Postings merge helpers (moved from IndexMerger) --- + +struct DeltaComputer { + buffer: Vec, +} + +impl DeltaComputer { + fn new() -> DeltaComputer { + DeltaComputer { + buffer: vec![0u32; 512], + } + } + + fn compute_delta(&mut self, positions: &[u32]) -> &[u32] { + if positions.len() > self.buffer.len() { + self.buffer.resize(positions.len(), 0u32); + } + let mut last_pos = 0u32; + for (cur_pos, dest) in positions.iter().cloned().zip(self.buffer.iter_mut()) { + *dest = cur_pos - last_pos; + last_pos = cur_pos; + } + &self.buffer[..positions.len()] + } +} + +fn estimate_total_num_tokens_in_single_segment( + reader: &SegmentReader, + field: Field, +) -> crate::Result { + if !reader.has_deletes() { + return Ok(reader.inverted_index(field)?.total_num_tokens()); + } + if let Some(fieldnorm_reader) = reader.fieldnorms_readers().get_field(field)? { + let mut count: [usize; 256] = [0; 256]; + for doc in reader.doc_ids_alive() { + let fieldnorm_id = fieldnorm_reader.fieldnorm_id(doc); + count[fieldnorm_id as usize] += 1; + } + let total_num_tokens = count + .iter() + .cloned() + .enumerate() + .map(|(fieldnorm_ord, count)| { + count as u64 * u64::from(FieldNormReader::id_to_fieldnorm(fieldnorm_ord as u8)) + }) + .sum::(); + return Ok(total_num_tokens); + } + let segment_num_tokens = reader.inverted_index(field)?.total_num_tokens(); + if reader.max_doc() == 0 { + return Ok(0u64); + } + let ratio = reader.num_docs() as f64 / reader.max_doc() as f64; + Ok((segment_num_tokens as f64 * ratio) as u64) +} + +fn estimate_total_num_tokens(readers: &[SegmentReader], field: Field) -> crate::Result { + let mut total_num_tokens: u64 = 0; + for reader in readers { + total_num_tokens += estimate_total_num_tokens_in_single_segment(reader, field)?; + } + Ok(total_num_tokens) +} + +fn write_postings_for_field( + readers: &[SegmentReader], + schema: &Schema, + indexed_field: Field, + serializer: &mut InvertedIndexSerializer, + fieldnorm_reader: Option, + doc_id_mapping: &SegmentDocIdMapping, +) -> crate::Result<()> { + debug_time!("write-postings-for-field"); + let mut positions_buffer: Vec = Vec::with_capacity(1_000); + let mut delta_computer = DeltaComputer::new(); + + let mut max_term_ords: Vec = Vec::new(); + + let field_readers: Vec> = readers + .iter() + .map(|reader| reader.inverted_index(indexed_field)) + .collect::>>()?; + + let mut field_term_streams = Vec::new(); + for field_reader in &field_readers { + let terms = field_reader.terms(); + field_term_streams.push(terms.stream()?); + max_term_ords.push(terms.num_terms() as u64); + } + + let mut merged_terms = TermMerger::new(field_term_streams); + + let mut merged_doc_id_map: Vec>> = readers + .iter() + .map(|reader| { + let mut segment_local_map = vec![]; + segment_local_map.resize(reader.max_doc() as usize, None); + segment_local_map + }) + .collect(); + for (new_doc_id, old_doc_addr) in doc_id_mapping.iter_old_doc_addrs().enumerate() { + let segment_map = &mut merged_doc_id_map[old_doc_addr.segment_ord as usize]; + segment_map[old_doc_addr.doc_id as usize] = Some(new_doc_id as DocId); + } + + let total_num_tokens: u64 = estimate_total_num_tokens(readers, indexed_field)?; + + let mut field_serializer = + serializer.new_field(indexed_field, total_num_tokens, fieldnorm_reader)?; + + let field_entry = schema.get_field_entry(indexed_field); + + let segment_postings_option = field_entry.field_type().get_index_record_option().expect( + "Encountered a field that is not supposed to be + indexed. Have you modified the schema?", + ); + + let mut segment_postings_containing_the_term: Vec<(usize, SegmentPostings)> = vec![]; + let mut doc_id_and_positions = vec![]; + + while merged_terms.advance() { + segment_postings_containing_the_term.clear(); + let term_bytes: &[u8] = merged_terms.key(); + + let mut total_doc_freq = 0; + + for (segment_ord, term_info) in merged_terms.current_segment_ords_and_term_infos() { + let segment_reader = &readers[segment_ord]; + let inverted_index: &InvertedIndexReader = &field_readers[segment_ord]; + let segment_postings = + inverted_index.read_postings_from_terminfo(&term_info, segment_postings_option)?; + let alive_bitset_opt = segment_reader.alive_bitset(); + let doc_freq = if let Some(alive_bitset) = alive_bitset_opt { + segment_postings.doc_freq_given_deletes(alive_bitset) + } else { + segment_postings.doc_freq() + }; + if doc_freq > 0u32 { + total_doc_freq += doc_freq; + segment_postings_containing_the_term.push((segment_ord, segment_postings)); + } + } + + if total_doc_freq == 0u32 { + continue; + } + + assert!(!segment_postings_containing_the_term.is_empty()); + + let has_term_freq = { + let has_term_freq = !segment_postings_containing_the_term[0] + .1 + .block_cursor + .freqs() + .is_empty(); + for (_, postings) in &segment_postings_containing_the_term[1..] { + if has_term_freq == postings.block_cursor.freqs().is_empty() { + return Err(DataCorruption::comment_only( + "Term freqs are inconsistent across segments", + ) + .into()); + } + } + has_term_freq + }; + + field_serializer.new_term(term_bytes, total_doc_freq, has_term_freq)?; + + for (segment_ord, mut segment_postings) in segment_postings_containing_the_term.drain(..) { + let old_to_new_doc_id = &merged_doc_id_map[segment_ord]; + + let mut doc = segment_postings.doc(); + while doc != TERMINATED { + if let Some(remapped_doc_id) = old_to_new_doc_id[doc as usize] { + let term_freq = if has_term_freq { + segment_postings.positions(&mut positions_buffer); + segment_postings.term_freq() + } else { + positions_buffer.clear(); + 0u32 + }; + + if !doc_id_mapping.is_trivial() { + doc_id_and_positions.push(( + remapped_doc_id, + term_freq, + positions_buffer.to_vec(), + )); + } else { + let delta_positions = delta_computer.compute_delta(&positions_buffer); + field_serializer.write_doc(remapped_doc_id, term_freq, delta_positions); + } + } + + doc = segment_postings.advance(); + } + } + if !doc_id_mapping.is_trivial() { + doc_id_and_positions.sort_unstable_by_key(|&(doc_id, _, _)| doc_id); + + for (doc_id, term_freq, positions) in &doc_id_and_positions { + let delta_positions = delta_computer.compute_delta(positions); + field_serializer.write_doc(*doc_id, *term_freq, delta_positions); + } + doc_id_and_positions.clear(); + } + field_serializer.close_term()?; + } + field_serializer.close()?; + Ok(()) +} + +fn write_postings_merge( + readers: &[SegmentReader], + schema: &Schema, + serializer: &mut InvertedIndexSerializer, + fieldnorm_readers: FieldNormReaders, + doc_id_mapping: &SegmentDocIdMapping, +) -> crate::Result<()> { + for (field, field_entry) in schema.fields() { + let fieldnorm_reader = fieldnorm_readers.get_field(field)?; + if field_entry.is_indexed() { + write_postings_for_field( + readers, + schema, + field, + serializer, + fieldnorm_reader, + doc_id_mapping, + )?; + } + } + Ok(()) +} diff --git a/src/index/mod.rs b/src/index/mod.rs index fcd9a14751..2dfa58842a 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -4,15 +4,18 @@ mod index; mod index_meta; +mod inverted_index_plugin; mod inverted_index_reader; mod segment; mod segment_component; mod segment_id; mod segment_reader; +pub(crate) use self::index::{builtin_plugins, list_segment_files}; pub use self::index::{Index, IndexBuilder}; pub(crate) use self::index_meta::SegmentMetaInventory; pub use self::index_meta::{IndexMeta, IndexSettings, IndexSortByField, Order, SegmentMeta}; +pub use self::inverted_index_plugin::{InvertedIndexPlugin, InvertedIndexPluginWriter}; pub use self::inverted_index_reader::InvertedIndexReader; pub use self::segment::Segment; pub use self::segment_component::SegmentComponent; diff --git a/src/index/segment.rs b/src/index/segment.rs index fcd32a1fff..98cac3837f 100644 --- a/src/index/segment.rs +++ b/src/index/segment.rs @@ -82,9 +82,9 @@ impl Segment { } /// Open one of the component file for *regular* write. - pub fn open_write(&mut self, component: SegmentComponent) -> Result { + pub fn open_write(&self, component: SegmentComponent) -> Result { let path = self.relative_path(component); - let write = self.index.directory_mut().open_write(&path)?; + let write = self.index.directory().open_write(&path)?; Ok(write) } } diff --git a/src/index/segment_component.rs b/src/index/segment_component.rs index 42ac1d178c..04989d4733 100644 --- a/src/index/segment_component.rs +++ b/src/index/segment_component.rs @@ -1,11 +1,11 @@ -use std::slice; +use std::fmt::{Display, Formatter}; /// Enum describing each component of a tantivy segment. /// /// Each component is stored in its own file, /// using the pattern `segment_uuid`.`component_extension`, /// except the delete component that takes an `segment_uuid`.`delete_opstamp`.`component_extension` -#[derive(Copy, Clone, Eq, PartialEq)] +#[derive(Debug, Clone, Hash, Eq, PartialEq)] pub enum SegmentComponent { /// Postings (or inverted list). Sorted lists of document ids, associated with terms Postings, @@ -28,21 +28,41 @@ pub enum SegmentComponent { /// Bitset describing which document of the segment is alive. /// (It was representing deleted docs but changed to represent alive docs from v0.17) Delete, + /// A custom component defined by a [`SegmentPlugin`](crate::SegmentPlugin). + /// The string is the file extension for this component. + Custom(String), } -impl SegmentComponent { - /// Iterates through the components. - pub fn iterator() -> slice::Iter<'static, SegmentComponent> { - static SEGMENT_COMPONENTS: [SegmentComponent; 8] = [ - SegmentComponent::Postings, - SegmentComponent::Positions, - SegmentComponent::FastFields, - SegmentComponent::FieldNorms, - SegmentComponent::Terms, - SegmentComponent::Store, - SegmentComponent::TempStore, - SegmentComponent::Delete, - ]; - SEGMENT_COMPONENTS.iter() +impl TryFrom<&str> for SegmentComponent { + type Error = String; + + fn try_from(value: &str) -> Result { + match value { + "idx" => Ok(SegmentComponent::Postings), + "pos" => Ok(SegmentComponent::Positions), + "term" => Ok(SegmentComponent::Terms), + "store" => Ok(SegmentComponent::Store), + "temp" => Ok(SegmentComponent::TempStore), + "fast" => Ok(SegmentComponent::FastFields), + "fieldnorm" => Ok(SegmentComponent::FieldNorms), + "del" => Ok(SegmentComponent::Delete), + other => Ok(SegmentComponent::Custom(other.to_string())), + } + } +} + +impl Display for SegmentComponent { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + match self { + SegmentComponent::Postings => write!(f, "idx"), + SegmentComponent::Positions => write!(f, "pos"), + SegmentComponent::FastFields => write!(f, "fast"), + SegmentComponent::FieldNorms => write!(f, "fieldnorm"), + SegmentComponent::Terms => write!(f, "term"), + SegmentComponent::Store => write!(f, "store"), + SegmentComponent::TempStore => write!(f, "temp"), + SegmentComponent::Delete => write!(f, "del"), + SegmentComponent::Custom(ext) => write!(f, "{ext}"), + } } } diff --git a/src/index/segment_reader.rs b/src/index/segment_reader.rs index cfb177064c..5b6ea0dfcf 100644 --- a/src/index/segment_reader.rs +++ b/src/index/segment_reader.rs @@ -1,4 +1,5 @@ -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; +use std::path::PathBuf; use std::sync::{Arc, RwLock}; use std::{fmt, io}; @@ -7,14 +8,14 @@ use fnv::FnvHashMap; use itertools::Itertools; use crate::directory::error::OpenReadError; -use crate::directory::{CompositeFile, FileSlice}; +use crate::directory::{CompositeFile, Directory, FileSlice}; use crate::error::DataCorruption; use crate::fastfield::{intersect_alive_bitsets, AliveBitSet, FacetReader, FastFieldReaders}; use crate::fieldnorm::{FieldNormReader, FieldNormReaders}; -use crate::index::{InvertedIndexReader, Segment, SegmentComponent, SegmentId}; +use crate::index::{Index, InvertedIndexReader, Segment, SegmentComponent, SegmentId}; use crate::json_utils::json_path_sep_to_dot; use crate::schema::{Field, IndexRecordOption, Schema, Type}; -use crate::space_usage::SegmentSpaceUsage; +use crate::space_usage::{ComponentSpaceUsage, SegmentSpaceUsage}; use crate::store::StoreReader; use crate::termdict::TermDictionary; use crate::{DocId, Opstamp}; @@ -31,6 +32,7 @@ use crate::{DocId, Opstamp}; /// as close to all of the memory data is mmapped. #[derive(Clone)] pub struct SegmentReader { + index: Index, inv_idx_reader_cache: Arc>>>, segment_id: SegmentId, @@ -190,6 +192,7 @@ impl SegmentReader { .unwrap_or(max_doc); Ok(SegmentReader { + index: segment.index().clone(), inv_idx_reader_cache: Default::default(), num_docs, max_doc, @@ -451,21 +454,54 @@ impl SegmentReader { } /// Summarize total space usage of this segment. + /// + /// The per-component usage is assembled from each registered plugin's + /// [`SegmentPlugin::space_usage`](crate::plugin::SegmentPlugin::space_usage), plus + /// the non-plugin `deletes` entry (the alive bitset). pub fn space_usage(&self) -> io::Result { - Ok(SegmentSpaceUsage::new( + let mut components: BTreeMap = BTreeMap::new(); + for plugin in self.index.all_plugins() { + let plugin_usage = plugin + .space_usage(self) + .map_err(|err| io::Error::other(err.to_string()))?; + components.extend(plugin_usage); + } + let deletes = self + .alive_bitset_opt + .as_ref() + .map(AliveBitSet::space_usage) + .unwrap_or_default(); + components.insert( + crate::space_usage::DELETES.to_string(), + ComponentSpaceUsage::Basic(deletes), + ); + Ok(SegmentSpaceUsage::from_components( self.num_docs(), - self.termdict_composite.space_usage(self.schema()), - self.postings_composite.space_usage(self.schema()), - self.positions_composite.space_usage(self.schema()), - self.fast_fields_readers.space_usage()?, - self.fieldnorm_readers.space_usage(self.schema()), - self.get_store_reader(0)?.space_usage(), - self.alive_bitset_opt - .as_ref() - .map(AliveBitSet::space_usage) - .unwrap_or_default(), + components, )) } + + fn relative_path(&self, component: SegmentComponent) -> PathBuf { + let mut path = self.segment_id().uuid_string(); + path.push_str(&match component { + SegmentComponent::Postings => ".idx".to_string(), + SegmentComponent::Positions => ".pos".to_string(), + SegmentComponent::Terms => ".term".to_string(), + SegmentComponent::Store => ".store".to_string(), + SegmentComponent::TempStore => ".store.temp".to_string(), + SegmentComponent::FastFields => ".fast".to_string(), + SegmentComponent::FieldNorms => ".fieldnorm".to_string(), + SegmentComponent::Delete => format!(".{}.del", self.delete_opstamp().unwrap_or(0)), + SegmentComponent::Custom(ext) => format!(".{ext}"), + }); + PathBuf::from(path) + } + + /// Opens one of the component files for reading. + pub fn open_read(&self, component: SegmentComponent) -> Result { + let path = self.relative_path(component); + self.index.directory().open_read(&path) + } } #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)] diff --git a/src/indexer/doc_id_mapping.rs b/src/indexer/doc_id_mapping.rs index 59ef2a128f..12fe7c07ab 100644 --- a/src/indexer/doc_id_mapping.rs +++ b/src/indexer/doc_id_mapping.rs @@ -16,7 +16,7 @@ pub enum MappingType { /// Struct to provide mapping from new doc_id to old doc_id and segment. #[derive(Clone)] -pub(crate) struct SegmentDocIdMapping { +pub struct SegmentDocIdMapping { pub(crate) new_doc_id_to_old_doc_addr: Vec, pub(crate) alive_bitsets: Vec>, mapping_type: MappingType, @@ -139,9 +139,9 @@ pub(crate) fn get_doc_id_mapping_from_field( sort_by_field: IndexSortByField, segment_writer: &SegmentWriter, ) -> crate::Result { - let schema = segment_writer.segment_serializer.segment().schema(); + let schema = segment_writer.segment.schema(); expect_field_id_for_sort_field(&schema, &sort_by_field)?; // for now expect - let new_doc_id_to_old = segment_writer.fast_field_writers.sort_order( + let new_doc_id_to_old = segment_writer.fast_fields.writer().sort_order( sort_by_field.field.as_str(), segment_writer.max_doc(), sort_by_field.order.is_desc(), diff --git a/src/indexer/index_writer.rs b/src/indexer/index_writer.rs index 721e816edd..9cf9984d47 100644 --- a/src/indexer/index_writer.rs +++ b/src/indexer/index_writer.rs @@ -278,10 +278,11 @@ impl IndexWriter { /// If the memory arena per thread is too small or too big, returns /// `TantivyError::InvalidArgument` pub(crate) fn new( - index: &Index, + index: Index, options: IndexWriterOptions, directory_lock: DirectoryLock, ) -> crate::Result { + index.validate_plugins()?; if options.memory_budget_per_thread < MEMORY_BUDGET_NUM_BYTES_MIN { let err_msg = format!( "The memory arena in bytes per thread needs to be at least \ @@ -574,7 +575,8 @@ impl IndexWriter { .take() .expect("The IndexWriter does not have any lock. This is a bug, please report."); - let new_index_writer = IndexWriter::new(&self.index, self.options.clone(), directory_lock)?; + let new_index_writer = + IndexWriter::new(self.index.clone(), self.options.clone(), directory_lock)?; // the current `self` is dropped right away because of this call. // diff --git a/src/indexer/merger.rs b/src/indexer/merger.rs index f373b1d877..4d0d9d2deb 100644 --- a/src/indexer/merger.rs +++ b/src/indexer/merger.rs @@ -1,28 +1,16 @@ use std::sync::Arc; -use columnar::{ - compute_merged_term_ord_mapping, BytesColumn, Column, ColumnType, ColumnarReader, - MergeRowOrder, RowAddr, ShuffleMergeOrder, StackMergeOrder, -}; +use columnar::{compute_merged_term_ord_mapping, BytesColumn, Column}; use common::ReadOnlyBitSet; use itertools::Itertools; -use measure_time::debug_time; -use crate::directory::WritePtr; -use crate::docset::{DocSet, TERMINATED}; -use crate::error::DataCorruption; use crate::fastfield::{AliveBitSet, FastFieldNotAvailableError}; -use crate::fieldnorm::{FieldNormReader, FieldNormReaders, FieldNormsSerializer, FieldNormsWriter}; -use crate::index::{Segment, SegmentComponent, SegmentReader}; +use crate::index::{builtin_plugins, Segment, SegmentReader}; use crate::indexer::doc_id_mapping::{MappingType, SegmentDocIdMapping}; -use crate::indexer::SegmentSerializer; -use crate::postings::{InvertedIndexSerializer, Postings, SegmentPostings}; -use crate::schema::{value_type_to_column_type, Field, FieldType, Schema, Type}; -use crate::store::StoreWriter; -use crate::termdict::{TermMerger, TermOrdinal}; -use crate::{ - DocAddress, DocId, IndexSettings, IndexSortByField, InvertedIndexReader, Order, SegmentOrdinal, -}; +use crate::plugin::{PluginMergeContext, SegmentPlugin}; +use crate::schema::{Schema, Type}; +use crate::termdict::TermOrdinal; +use crate::{DocAddress, DocId, IndexSettings, IndexSortByField, Order, SegmentOrdinal}; /// Per-segment accessor for Str/Bytes sort fields during index merging. /// @@ -63,123 +51,12 @@ enum ReaderSortFieldAccessors { /// We do not allow segments with more than pub const MAX_DOC_LIMIT: u32 = 1 << 31; -fn estimate_total_num_tokens_in_single_segment( - reader: &SegmentReader, - field: Field, -) -> crate::Result { - // There are no deletes. We can simply use the exact value saved into the posting list. - // Note that this value is not necessarily exact as it could have been the result of a merge - // between segments themselves containing deletes. - if !reader.has_deletes() { - return Ok(reader.inverted_index(field)?.total_num_tokens()); - } - - // When there are deletes, we use an approximation either - // by using the fieldnorm. - if let Some(fieldnorm_reader) = reader.fieldnorms_readers().get_field(field)? { - let mut count: [usize; 256] = [0; 256]; - for doc in reader.doc_ids_alive() { - let fieldnorm_id = fieldnorm_reader.fieldnorm_id(doc); - count[fieldnorm_id as usize] += 1; - } - let total_num_tokens = count - .iter() - .cloned() - .enumerate() - .map(|(fieldnorm_ord, count)| { - count as u64 * u64::from(FieldNormReader::id_to_fieldnorm(fieldnorm_ord as u8)) - }) - .sum::(); - return Ok(total_num_tokens); - } - - // There are no fieldnorms available. - // Here we just do a pro-rata with the overall number of tokens an the ratio of - // documents alive. - let segment_num_tokens = reader.inverted_index(field)?.total_num_tokens(); - if reader.max_doc() == 0 { - // That supposedly never happens, but let's be a bit defensive here. - return Ok(0u64); - } - let ratio = reader.num_docs() as f64 / reader.max_doc() as f64; - Ok((segment_num_tokens as f64 * ratio) as u64) -} - -fn estimate_total_num_tokens(readers: &[SegmentReader], field: Field) -> crate::Result { - let mut total_num_tokens: u64 = 0; - for reader in readers { - total_num_tokens += estimate_total_num_tokens_in_single_segment(reader, field)?; - } - Ok(total_num_tokens) -} - pub struct IndexMerger { index_settings: IndexSettings, schema: Schema, pub(crate) readers: Vec, max_doc: u32, -} - -struct DeltaComputer { - buffer: Vec, -} - -impl DeltaComputer { - fn new() -> DeltaComputer { - DeltaComputer { - buffer: vec![0u32; 512], - } - } - - fn compute_delta(&mut self, positions: &[u32]) -> &[u32] { - if positions.len() > self.buffer.len() { - self.buffer.resize(positions.len(), 0u32); - } - let mut last_pos = 0u32; - for (cur_pos, dest) in positions.iter().cloned().zip(self.buffer.iter_mut()) { - *dest = cur_pos - last_pos; - last_pos = cur_pos; - } - &self.buffer[..positions.len()] - } -} - -fn convert_to_merge_order( - columnars: &[&ColumnarReader], - doc_id_mapping: SegmentDocIdMapping, -) -> MergeRowOrder { - match doc_id_mapping.mapping_type() { - MappingType::Stacked => MergeRowOrder::Stack(StackMergeOrder::stack(columnars)), - MappingType::StackedWithDeletes | MappingType::Shuffled => { - // RUST/LLVM is amazing. The following conversion is actually a no-op: - // no allocation, no copy. - let new_row_id_to_old_row_id: Vec = doc_id_mapping - .new_doc_id_to_old_doc_addr - .into_iter() - .map(|doc_addr| RowAddr { - segment_ord: doc_addr.segment_ord, - row_id: doc_addr.doc_id, - }) - .collect(); - MergeRowOrder::Shuffled(ShuffleMergeOrder { - new_row_id_to_old_row_id, - alive_bitsets: doc_id_mapping.alive_bitsets, - }) - } - } -} - -fn extract_fast_field_required_columns(schema: &Schema) -> Vec<(String, ColumnType)> { - schema - .fields() - .map(|(_, field_entry)| field_entry) - .filter(|field_entry| field_entry.is_fast()) - .filter_map(|field_entry| { - let column_name = field_entry.name().to_string(); - let column_type = value_type_to_column_type(field_entry.field_type().value_type())?; - Some((column_name, column_type)) - }) - .collect() + custom_plugins: Vec>, } impl IndexMerger { @@ -269,14 +146,40 @@ impl IndexMerger { ); return Err(crate::TantivyError::InvalidArgument(err_msg)); } + // The merged output carries the first segment's custom plugin set, so every source + // segment must register the same set — otherwise a later segment's custom component + // would be silently dropped, or the first index's plugin would read a component the + // later segment lacks. Fail closed on any mismatch. + let custom_plugins = if let Some(first_segment) = segments.first() { + let expected = first_segment.index().custom_extensions(); + for segment in &segments[1..] { + if segment.index().custom_extensions() != expected { + return Err(crate::TantivyError::InvalidArgument( + "Attempt to merge segments whose source indices register different plugin \ + sets" + .to_string(), + )); + } + } + first_segment.index().custom_plugins().to_vec() + } else { + Vec::new() + }; + Ok(IndexMerger { index_settings, schema, readers, max_doc, + custom_plugins, }) } + /// The custom (non-built-in) plugins of the segments being merged. + pub(crate) fn custom_plugins(&self) -> &[Arc] { + &self.custom_plugins + } + fn sort_by_field_type(&self, sort_by_field: &IndexSortByField) -> crate::Result { let schema_field = self.schema.get_field(&sort_by_field.field)?; let field_entry = self.schema.get_field_entry(schema_field); @@ -314,53 +217,6 @@ impl IndexMerger { .collect()) } - fn write_fieldnorms( - &self, - mut fieldnorms_serializer: FieldNormsSerializer, - doc_id_mapping: &SegmentDocIdMapping, - ) -> crate::Result<()> { - let fields = FieldNormsWriter::fields_with_fieldnorm(&self.schema); - let mut fieldnorms_data = Vec::with_capacity(self.max_doc as usize); - for field in fields { - fieldnorms_data.clear(); - let fieldnorms_readers: Vec = self - .readers - .iter() - .map(|reader| reader.get_fieldnorms_reader(field)) - .collect::>()?; - for old_doc_addr in doc_id_mapping.iter_old_doc_addrs() { - let fieldnorms_reader = &fieldnorms_readers[old_doc_addr.segment_ord as usize]; - let fieldnorm_id = fieldnorms_reader.fieldnorm_id(old_doc_addr.doc_id); - fieldnorms_data.push(fieldnorm_id); - } - fieldnorms_serializer.serialize_field(field, &fieldnorms_data[..])?; - } - fieldnorms_serializer.close()?; - Ok(()) - } - - fn write_fast_fields( - &self, - fast_field_wrt: &mut WritePtr, - doc_id_mapping: SegmentDocIdMapping, - ) -> crate::Result<()> { - debug_time!("write-fast-fields"); - let required_columns = extract_fast_field_required_columns(&self.schema); - let columnars: Vec<&ColumnarReader> = self - .readers - .iter() - .map(|reader| reader.fast_fields().columnar()) - .collect(); - let merge_row_order = convert_to_merge_order(&columnars[..], doc_id_mapping); - columnar::merge_columnar( - &columnars[..], - &required_columns, - merge_row_order, - fast_field_wrt, - )?; - Ok(()) - } - /// Checks if segments can use the fast disjunct-stack path (byte concatenation) /// instead of a full k-way merge. /// @@ -645,309 +501,14 @@ impl IndexMerger { )) } - fn write_postings_for_field( - &self, - indexed_field: Field, - _field_type: &FieldType, - serializer: &mut InvertedIndexSerializer, - fieldnorm_reader: Option, - doc_id_mapping: &SegmentDocIdMapping, - ) -> crate::Result<()> { - debug_time!("write-postings-for-field"); - let mut positions_buffer: Vec = Vec::with_capacity(1_000); - let mut delta_computer = DeltaComputer::new(); - - let mut max_term_ords: Vec = Vec::new(); - - let field_readers: Vec> = self - .readers - .iter() - .map(|reader| reader.inverted_index(indexed_field)) - .collect::>>()?; - - let mut field_term_streams = Vec::new(); - for field_reader in &field_readers { - let terms = field_reader.terms(); - field_term_streams.push(terms.stream()?); - max_term_ords.push(terms.num_terms() as u64); - } - - let mut merged_terms = TermMerger::new(field_term_streams); - - // map from segment doc ids to the resulting merged segment doc id. - - let mut merged_doc_id_map: Vec>> = self - .readers - .iter() - .map(|reader| { - let mut segment_local_map = vec![]; - segment_local_map.resize(reader.max_doc() as usize, None); - segment_local_map - }) - .collect(); - for (new_doc_id, old_doc_addr) in doc_id_mapping.iter_old_doc_addrs().enumerate() { - let segment_map = &mut merged_doc_id_map[old_doc_addr.segment_ord as usize]; - segment_map[old_doc_addr.doc_id as usize] = Some(new_doc_id as DocId); - } - - // Note that the total number of tokens is not exact. - // It is only used as a parameter in the BM25 formula. - let total_num_tokens: u64 = estimate_total_num_tokens(&self.readers, indexed_field)?; - - // Create the total list of doc ids - // by stacking the doc ids from the different segment. - // - // In the new segments, the doc id from the different - // segment are stacked so that : - // - Segment 0's doc ids become doc id [0, seg.max_doc] - // - Segment 1's doc ids become [seg0.max_doc, seg0.max_doc + seg.max_doc] - // - Segment 2's doc ids become [seg0.max_doc + seg1.max_doc, seg0.max_doc + seg1.max_doc + - // seg2.max_doc] - // - // This stacking applies only when the index is not sorted, in that case the - // doc_ids are kmerged by their sort property - let mut field_serializer = - serializer.new_field(indexed_field, total_num_tokens, fieldnorm_reader)?; - - let field_entry = self.schema.get_field_entry(indexed_field); - - // ... set segment postings option the new field. - let segment_postings_option = field_entry.field_type().get_index_record_option().expect( - "Encountered a field that is not supposed to be - indexed. Have you modified the schema?", - ); - - let mut segment_postings_containing_the_term: Vec<(usize, SegmentPostings)> = vec![]; - let mut doc_id_and_positions = vec![]; - - while merged_terms.advance() { - segment_postings_containing_the_term.clear(); - let term_bytes: &[u8] = merged_terms.key(); - - let mut total_doc_freq = 0; - - // Let's compute the list of non-empty posting lists - for (segment_ord, term_info) in merged_terms.current_segment_ords_and_term_infos() { - let segment_reader = &self.readers[segment_ord]; - let inverted_index: &InvertedIndexReader = &field_readers[segment_ord]; - let segment_postings = inverted_index - .read_postings_from_terminfo(&term_info, segment_postings_option)?; - let alive_bitset_opt = segment_reader.alive_bitset(); - let doc_freq = if let Some(alive_bitset) = alive_bitset_opt { - segment_postings.doc_freq_given_deletes(alive_bitset) - } else { - segment_postings.doc_freq() - }; - if doc_freq > 0u32 { - total_doc_freq += doc_freq; - segment_postings_containing_the_term.push((segment_ord, segment_postings)); - } - } - - // At this point, `segment_postings` contains the posting list - // of all of the segments containing the given term (and that are non-empty) - // - // These segments are non-empty and advance has already been called. - if total_doc_freq == 0u32 { - // All docs that used to contain the term have been deleted. The `term` will be - // entirely removed. - continue; - } - - // This should never happen as we early exited for total_doc_freq == 0. - assert!(!segment_postings_containing_the_term.is_empty()); - - let has_term_freq = { - let has_term_freq = !segment_postings_containing_the_term[0] - .1 - .block_cursor - .freqs() - .is_empty(); - for (_, postings) in &segment_postings_containing_the_term[1..] { - // This may look at a strange way to test whether we have term freq or not. - // With JSON object, the schema is not sufficient to know whether a term - // has its term frequency encoded or not: - // strings may have term frequencies, while number terms never have one. - // - // Ideally, we should have burnt one bit of two in the `TermInfo`. - // However, we preferred not changing the codec too much and detect this - // instead by - // - looking at the size of the skip data for bitpacked blocks - // - observing the absence of remaining data after reading the docs for vint - // blocks. - // - // Overall the reliable way to know if we have actual frequencies loaded or not - // is to check whether the actual decoded array is empty or not. - if has_term_freq == postings.block_cursor.freqs().is_empty() { - return Err(DataCorruption::comment_only( - "Term freqs are inconsistent across segments", - ) - .into()); - } - } - has_term_freq - }; - - field_serializer.new_term(term_bytes, total_doc_freq, has_term_freq)?; - - // We can now serialize this postings, by pushing each document to the - // postings serializer. - for (segment_ord, mut segment_postings) in - segment_postings_containing_the_term.drain(..) - { - let old_to_new_doc_id = &merged_doc_id_map[segment_ord]; - - let mut doc = segment_postings.doc(); - while doc != TERMINATED { - // deleted doc are skipped as they do not have a `remapped_doc_id`. - if let Some(remapped_doc_id) = old_to_new_doc_id[doc as usize] { - // we make sure to only write the term if - // there is at least one document. - let term_freq = if has_term_freq { - segment_postings.positions(&mut positions_buffer); - segment_postings.term_freq() - } else { - // The positions_buffer may contain positions from the previous term - // Existence of positions depend on the value type in JSON fields. - // https://github.com/quickwit-oss/tantivy/issues/2283 - positions_buffer.clear(); - 0u32 - }; - - // if doc_id_mapping exists, the doc_ids are reordered, they are - // not just stacked. The field serializer expects monotonically increasing - // doc_ids, so we collect and sort them first, before writing. - // - // I think this is not strictly necessary, it would be possible to - // avoid the loading into a vec via some form of kmerge, but then the merge - // logic would deviate much more from the stacking case (unsorted index) - if !doc_id_mapping.is_trivial() { - doc_id_and_positions.push(( - remapped_doc_id, - term_freq, - positions_buffer.to_vec(), - )); - } else { - let delta_positions = delta_computer.compute_delta(&positions_buffer); - field_serializer.write_doc(remapped_doc_id, term_freq, delta_positions); - } - } - - doc = segment_postings.advance(); - } - } - if !doc_id_mapping.is_trivial() { - doc_id_and_positions.sort_unstable_by_key(|&(doc_id, _, _)| doc_id); - - for (doc_id, term_freq, positions) in &doc_id_and_positions { - let delta_positions = delta_computer.compute_delta(positions); - field_serializer.write_doc(*doc_id, *term_freq, delta_positions); - } - doc_id_and_positions.clear(); - } - // closing the term. - field_serializer.close_term()?; - } - field_serializer.close()?; - Ok(()) - } - - fn write_postings( - &self, - serializer: &mut InvertedIndexSerializer, - fieldnorm_readers: FieldNormReaders, - doc_id_mapping: &SegmentDocIdMapping, - ) -> crate::Result<()> { - for (field, field_entry) in self.schema.fields() { - let fieldnorm_reader = fieldnorm_readers.get_field(field)?; - if field_entry.is_indexed() { - self.write_postings_for_field( - field, - field_entry.field_type(), - serializer, - fieldnorm_reader, - doc_id_mapping, - )?; - } - } - Ok(()) - } - - fn write_storable_fields( - &self, - store_writer: &mut StoreWriter, - doc_id_mapping: &SegmentDocIdMapping, - ) -> crate::Result<()> { - debug_time!("write-storable-fields"); - debug!("write-storable-field"); - - if !doc_id_mapping.is_trivial() { - debug!("non-trivial-doc-id-mapping"); - - let store_readers: Vec<_> = self - .readers - .iter() - .map(|reader| reader.get_store_reader(50)) - .collect::>()?; - - let mut document_iterators: Vec<_> = store_readers - .iter() - .enumerate() - .map(|(i, store)| store.iter_raw(self.readers[i].alive_bitset())) - .collect(); - - for old_doc_addr in doc_id_mapping.iter_old_doc_addrs() { - let doc_bytes_it = &mut document_iterators[old_doc_addr.segment_ord as usize]; - if let Some(doc_bytes_res) = doc_bytes_it.next() { - let doc_bytes = doc_bytes_res?; - store_writer.store_bytes(&doc_bytes)?; - } else { - return Err(DataCorruption::comment_only(format!( - "unexpected missing document in docstore on merge, doc address \ - {old_doc_addr:?}", - )) - .into()); - } - } - } else { - debug!("trivial-doc-id-mapping"); - for reader in &self.readers { - let store_reader = reader.get_store_reader(1)?; - if reader.has_deletes() - // If there is not enough data in the store, we avoid stacking in order to - // avoid creating many small blocks in the doc store. Once we have 5 full blocks, - // we start stacking. In the worst case 2/7 of the blocks would be very small. - // [segment 1 - {1 doc}][segment 2 - {fullblock * 5}{1doc}] - // => 5 * full blocks, 2 * 1 document blocks - // - // In a more realistic scenario the segments are of the same size, so 1/6 of - // the doc stores would be on average half full, given total randomness (which - // is not the case here, but not sure how it behaves exactly). - // - // https://github.com/quickwit-oss/tantivy/issues/1053 - // - // take 7 in order to not walk over all checkpoints. - || store_reader.block_checkpoints().take(7).count() < 6 - || store_reader.decompressor() != store_writer.compressor().into() - { - for doc_bytes_res in store_reader.iter_raw(reader.alive_bitset()) { - let doc_bytes = doc_bytes_res?; - store_writer.store_bytes(&doc_bytes)?; - } - } else { - store_writer.stack(store_reader)?; - } - } - } - Ok(()) - } - - /// Writes the merged segment by pushing information - /// to the `SegmentSerializer`. + /// Writes the merged segment into `target_segment`. + /// + /// Each plugin's [`SegmentPlugin::merge`] opens and closes its own serializers, + /// so the merge path drives no incremental plugin writers. /// /// # Returns /// The number of documents in the resulting segment. - pub fn write(&self, mut serializer: SegmentSerializer) -> crate::Result { + pub fn write(&self, target_segment: &Segment) -> crate::Result { let doc_id_mapping = if let Some(sort_by_field) = self.index_settings.sort_by_field.as_ref() { if self.is_disjunct_and_sorted_on_sort_property(sort_by_field)? { @@ -958,28 +519,23 @@ impl IndexMerger { } else { self.get_doc_id_from_concatenated_data()? }; - debug!("write-fieldnorms"); - if let Some(fieldnorms_serializer) = serializer.extract_fieldnorms_serializer() { - self.write_fieldnorms(fieldnorms_serializer, &doc_id_mapping)?; + + // Merge plugins in write order (matching `Index::all_plugins`): + // inverted index, then fast_fields, store, and custom plugins. + for plugin in builtin_plugins() + .into_iter() + .chain(self.custom_plugins.iter().cloned()) + { + debug!("merge-plugin: {:?}", plugin.extensions()); + plugin.merge(PluginMergeContext { + readers: &self.readers, + doc_id_mapping: &doc_id_mapping, + target_segment, + schema: &self.schema, + settings: &self.index_settings, + })?; } - debug!("write-postings"); - let fieldnorm_data = serializer - .segment() - .open_read(SegmentComponent::FieldNorms)?; - let fieldnorm_readers = FieldNormReaders::open(fieldnorm_data)?; - self.write_postings( - serializer.get_postings_serializer(), - fieldnorm_readers, - &doc_id_mapping, - )?; - - debug!("write-storagefields"); - self.write_storable_fields(serializer.get_store_writer(), &doc_id_mapping)?; - debug!("write-fastfields"); - self.write_fast_fields(serializer.get_fast_field_write(), doc_id_mapping)?; - - debug!("close-serializer"); - serializer.close()?; + Ok(self.max_doc) } } diff --git a/src/indexer/mod.rs b/src/indexer/mod.rs index 44d89b3a8e..ee87409604 100644 --- a/src/indexer/mod.rs +++ b/src/indexer/mod.rs @@ -8,7 +8,7 @@ pub(crate) mod delete_queue; pub(crate) mod path_to_unordered_id; -pub(crate) mod doc_id_mapping; +pub mod doc_id_mapping; mod doc_opstamp_mapping; mod flat_map_with_buffer; pub(crate) mod index_writer; @@ -24,7 +24,6 @@ pub(crate) mod prepared_commit; mod segment_entry; mod segment_manager; mod segment_register; -pub(crate) mod segment_serializer; pub(crate) mod segment_updater; pub(crate) mod segment_writer; pub(crate) mod single_segment_index_writer; @@ -40,7 +39,6 @@ pub use self::merge_policy::{MergeCandidate, MergePolicy, NoMergePolicy}; pub use self::operation::{AddOperation, DeleteOperation, UserOperation}; pub use self::prepared_commit::PreparedCommit; pub use self::segment_entry::SegmentEntry; -pub(crate) use self::segment_serializer::SegmentSerializer; pub use self::segment_updater::{merge_filtered_segments, merge_indices}; pub use self::segment_writer::SegmentWriter; pub use self::single_segment_index_writer::SingleSegmentIndexWriter; diff --git a/src/indexer/segment_serializer.rs b/src/indexer/segment_serializer.rs deleted file mode 100644 index cdb7a8ef50..0000000000 --- a/src/indexer/segment_serializer.rs +++ /dev/null @@ -1,111 +0,0 @@ -use common::TerminatingWrite; - -use crate::directory::WritePtr; -use crate::fieldnorm::FieldNormsSerializer; -use crate::index::{Segment, SegmentComponent}; -use crate::postings::InvertedIndexSerializer; -use crate::store::StoreWriter; - -/// Segment serializer is in charge of laying out on disk -/// the data accumulated and sorted by the `SegmentWriter`. -pub struct SegmentSerializer { - segment: Segment, - pub(crate) store_writer: StoreWriter, - fast_field_write: WritePtr, - fieldnorms_serializer: Option, - postings_serializer: InvertedIndexSerializer, -} - -impl SegmentSerializer { - /// Creates a new `SegmentSerializer`. - pub fn for_segment( - mut segment: Segment, - is_in_merge: bool, - ) -> crate::Result { - // 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 store_writer = if remapping_required { - let store_write = segment.open_write(SegmentComponent::TempStore)?; - StoreWriter::new( - store_write, - 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. - 16000, - settings.docstore_compress_dedicated_thread, - )? - } else { - let store_write = segment.open_write(SegmentComponent::Store)?; - StoreWriter::new( - store_write, - settings.docstore_compression, - settings.docstore_blocksize, - settings.docstore_compress_dedicated_thread, - )? - }; - - let fast_field_write = segment.open_write(SegmentComponent::FastFields)?; - - let fieldnorms_write = segment.open_write(SegmentComponent::FieldNorms)?; - let fieldnorms_serializer = FieldNormsSerializer::from_write(fieldnorms_write)?; - - let postings_serializer = InvertedIndexSerializer::open(&mut segment)?; - Ok(SegmentSerializer { - segment, - store_writer, - fast_field_write, - fieldnorms_serializer: Some(fieldnorms_serializer), - postings_serializer, - }) - } - - /// The memory used (inclusive childs) - pub fn mem_usage(&self) -> usize { - self.store_writer.mem_usage() - } - - pub fn segment(&self) -> &Segment { - &self.segment - } - - pub fn segment_mut(&mut self) -> &mut Segment { - &mut self.segment - } - - /// Accessor to the `PostingsSerializer`. - pub fn get_postings_serializer(&mut self) -> &mut InvertedIndexSerializer { - &mut self.postings_serializer - } - - /// Accessor to the `FastFieldSerializer`. - pub fn get_fast_field_write(&mut self) -> &mut WritePtr { - &mut self.fast_field_write - } - - /// Extract the field norm serializer. - /// - /// Note the fieldnorms serializer can only be extracted once. - pub fn extract_fieldnorms_serializer(&mut self) -> Option { - self.fieldnorms_serializer.take() - } - - /// Accessor to the `StoreWriter`. - pub fn get_store_writer(&mut self) -> &mut StoreWriter { - &mut self.store_writer - } - - /// Finalize the segment serialization. - pub fn close(mut self) -> crate::Result<()> { - if let Some(fieldnorms_serializer) = self.extract_fieldnorms_serializer() { - fieldnorms_serializer.close()?; - } - self.fast_field_write.terminate()?; - self.postings_serializer.close()?; - self.store_writer.close()?; - Ok(()) - } -} diff --git a/src/indexer/segment_updater.rs b/src/indexer/segment_updater.rs index 147ff06d30..9f132e78f4 100644 --- a/src/indexer/segment_updater.rs +++ b/src/indexer/segment_updater.rs @@ -13,7 +13,9 @@ use super::segment_manager::SegmentManager; use crate::core::META_FILEPATH; use crate::directory::{Directory, DirectoryClone, GarbageCollectionResult}; use crate::fastfield::AliveBitSet; -use crate::index::{Index, IndexMeta, IndexSettings, Segment, SegmentId, SegmentMeta}; +use crate::index::{ + list_segment_files, Index, IndexMeta, IndexSettings, Segment, SegmentId, SegmentMeta, +}; use crate::indexer::delete_queue::DeleteCursor; use crate::indexer::index_writer::advance_deletes; use crate::indexer::merge_operation::MergeOperationInventory; @@ -22,7 +24,6 @@ use crate::indexer::segment_manager::SegmentsStatus; use crate::indexer::stamper::Stamper; use crate::indexer::{ DefaultMergePolicy, MergeCandidate, MergeOperation, MergePolicy, SegmentEntry, - SegmentSerializer, }; use crate::{FutureResult, Opstamp, TantivyError}; @@ -118,9 +119,7 @@ fn merge( IndexMerger::open(index.schema(), index.settings().clone(), &segments[..])?; // ... we just serialize this index merger in our new segment to merge the segments. - let segment_serializer = SegmentSerializer::for_segment(merged_segment.clone(), true)?; - - let num_docs = merger.write(segment_serializer)?; + let num_docs = merger.write(&merged_segment)?; let merged_segment_id = merged_segment.id(); @@ -139,6 +138,10 @@ fn merge( /// This function does NOT check or take the `IndexWriter` is running. It is not /// meant to work if you have an `IndexWriter` running for the origin indices, or /// the destination `Index`. +/// +/// Every custom [`SegmentPlugin`](crate::SegmentPlugin) must be registered on the source +/// indices, and all source indices must register the same plugin set; the merge errors +/// otherwise. The merged output carries that shared plugin set forward. #[doc(hidden)] pub fn merge_indices>>( indices: &[Index], @@ -185,6 +188,10 @@ pub fn merge_indices>>( /// This function does NOT check or take the `IndexWriter` is running. It is not /// meant to work if you have an `IndexWriter` running for the origin indices, or /// the destination `Index`. +/// +/// Every custom [`SegmentPlugin`](crate::SegmentPlugin) must be registered on the segments' +/// source indices, and all of those indices must register the same plugin set; the merge +/// errors otherwise. The merged output carries that shared plugin set forward. #[doc(hidden)] pub fn merge_filtered_segments>>( segments: &[Segment], @@ -212,21 +219,33 @@ pub fn merge_filtered_segments>>( )); } + // The source index's writer was already validated against its recorded plugin + // set (else `MissingPlugin`); the merger reuses those registered plugins. + let merger = IndexMerger::open_with_custom_alive_set( + target_schema.clone(), + target_settings.clone(), + segments, + filter_doc_ids, + )?; + let mut merged_index = Index::create( output_directory, target_schema.clone(), target_settings.clone(), )?; + + // Carry the merger's custom plugins onto the merged index (built-ins come for free) and + // record their extensions as the merged index's required set; without this the first GC + // drops the merged files. + let mut persisted_custom_extensions: Vec = Vec::new(); + for plugin in merger.custom_plugins() { + persisted_custom_extensions.extend(plugin.extensions().iter().map(|ext| ext.to_string())); + merged_index.register_plugin(plugin.clone()); + } + let merged_segment = merged_index.new_segment(); let merged_segment_id = merged_segment.id(); - let merger: IndexMerger = IndexMerger::open_with_custom_alive_set( - merged_index.schema(), - merged_index.settings().clone(), - segments, - filter_doc_ids, - )?; - let segment_serializer = SegmentSerializer::for_segment(merged_segment, true)?; - let num_docs = merger.write(segment_serializer)?; + let num_docs = merger.write(&merged_segment)?; let segment_meta = merged_index.new_segment_meta(merged_segment_id, num_docs); @@ -243,6 +262,7 @@ pub fn merge_filtered_segments>>( let index_meta = IndexMeta { index_settings: target_settings, // index_settings of all segments should be the same + persisted_custom_extensions, segments: vec![segment_meta], schema: target_schema, opstamp: 0u64, @@ -412,6 +432,8 @@ impl SegmentUpdater { .sort_by_key(|segment_meta| std::cmp::Reverse(segment_meta.max_doc())); let index_meta = IndexMeta { index_settings: index.settings().clone(), + // The required plugin set is fixed at index creation; carry it forward. + persisted_custom_extensions: self.load_meta().persisted_custom_extensions.clone(), segments: committed_segment_metas, schema: index.schema(), opstamp, @@ -434,12 +456,13 @@ impl SegmentUpdater { /// This does not include lock files, or files that are obsolete /// but have not yet been deleted by the garbage collector. fn list_files(&self) -> HashSet { - let mut files: HashSet = self - .index - .list_all_segment_metas() - .into_iter() - .flat_map(|segment_meta| segment_meta.list_files()) - .collect(); + // All tracked segments (including in-flight ones), so GC keeps files for segments + // not yet in the committed meta. Custom extensions come from the persisted record + // (not the live registry), so GC is correct even before a plugin is re-registered. + let mut files = list_segment_files( + &self.index.list_all_segment_metas(), + &self.load_meta().persisted_custom_extensions, + ); files.insert(META_FILEPATH.to_path_buf()); files } diff --git a/src/indexer/segment_writer.rs b/src/indexer/segment_writer.rs index 5ffda31448..5065a37574 100644 --- a/src/indexer/segment_writer.rs +++ b/src/indexer/segment_writer.rs @@ -1,46 +1,12 @@ -use columnar::MonotonicallyMappableToU64; -use common::JsonPathWriter; -use itertools::Itertools; -use tokenizer_api::BoxTokenStream; - use super::doc_id_mapping::{get_doc_id_mapping_from_field, DocIdMapping}; use super::operation::AddOperation; -use crate::fastfield::FastFieldsWriter; -use crate::fieldnorm::{FieldNormReaders, FieldNormsWriter}; -use crate::index::{Segment, SegmentComponent}; -use crate::indexer::indexing_term::IndexingTerm; -use crate::indexer::segment_serializer::SegmentSerializer; -use crate::json_utils::{index_json_value, IndexingPositionsPerPath}; -use crate::postings::{ - compute_table_memory_size, serialize_postings, IndexingContext, IndexingPosition, - PerFieldPostingsWriter, PostingsWriter, -}; -use crate::schema::document::{Document, Value}; -use crate::schema::{FieldEntry, FieldType, Schema, DATE_TIME_PRECISION_INDEXED}; -use crate::store::{StoreReader, StoreWriter}; -use crate::tokenizer::{FacetTokenizer, PreTokenizedStream, TextAnalyzer, Tokenizer}; -use crate::{DocId, Opstamp, TantivyError}; - -/// Computes the initial size of the hash table. -/// -/// Returns the recommended initial table size as a power of 2. -/// -/// Note this is a very dumb way to compute log2, but it is easier to proofread that way. -fn compute_initial_table_size(per_thread_memory_budget: usize) -> crate::Result { - let table_memory_upper_bound = per_thread_memory_budget / 3; - (10..20) // We cap it at 2^19 = 512K capacity. - // TODO: There are cases where this limit causes a - // reallocation in the hashmap. Check if this affects performance. - .map(|power| 1 << power) - .take_while(|capacity| compute_table_memory_size(*capacity) < table_memory_upper_bound) - .last() - .ok_or_else(|| { - crate::TantivyError::InvalidArgument(format!( - "per thread memory budget (={per_thread_memory_budget}) is too small. Raise the \ - memory budget or lower the number of threads." - )) - }) -} +use crate::fastfield::FastFieldsPluginWriter; +use crate::index::{InvertedIndexPluginWriter, Segment}; +use crate::plugin::{PluginWriter, PluginWriterContext}; +use crate::schema::document::{Document, ErasedDocument}; +use crate::schema::Schema; +use crate::store::StorePluginWriter; +use crate::{DocId, Opstamp}; fn remap_doc_opstamps( opstamps: Vec, @@ -63,16 +29,12 @@ fn remap_doc_opstamps( /// The segment is laid on disk when the segment gets `finalized`. pub struct SegmentWriter { pub(crate) max_doc: DocId, - pub(crate) ctx: IndexingContext, - pub(crate) per_field_postings_writers: PerFieldPostingsWriter, - pub(crate) segment_serializer: SegmentSerializer, - pub(crate) fast_field_writers: FastFieldsWriter, - pub(crate) fieldnorms_writer: FieldNormsWriter, - pub(crate) json_path_writer: JsonPathWriter, - pub(crate) json_positions_per_path: IndexingPositionsPerPath, + pub(crate) segment: Segment, + pub(crate) inverted_index: InvertedIndexPluginWriter, + pub(crate) fast_fields: FastFieldsPluginWriter, + pub(crate) store_writer: StorePluginWriter, + custom_plugins: Vec>, pub(crate) doc_opstamps: Vec, - per_field_text_analyzers: Vec, - term_buffer: IndexingTerm, schema: Schema, } @@ -88,76 +50,69 @@ impl SegmentWriter { /// - schema pub fn for_segment(memory_budget_in_bytes: usize, segment: Segment) -> crate::Result { let schema = segment.schema(); - let tokenizer_manager = segment.index().tokenizers().clone(); - let tokenizer_manager_fast_field = segment.index().fast_field_tokenizer().clone(); - let table_size = compute_initial_table_size(memory_budget_in_bytes)?; - let segment_serializer = SegmentSerializer::for_segment(segment, false)?; - let per_field_postings_writers = PerFieldPostingsWriter::for_schema(&schema); - let per_field_text_analyzers = schema - .fields() - .map(|(_, field_entry): (_, &FieldEntry)| { - let text_options = match field_entry.field_type() { - FieldType::Str(ref text_options) => text_options.get_indexing_options(), - FieldType::JsonObject(ref json_object_options) => { - json_object_options.get_text_indexing_options() - } - _ => None, - }; - let tokenizer_name = text_options - .map(|text_index_option| text_index_option.tokenizer()) - .unwrap_or("default"); - - tokenizer_manager.get(tokenizer_name).ok_or_else(|| { - TantivyError::SchemaError(format!( - "Error getting tokenizer for field: {}", - field_entry.name() - )) - }) - }) - .collect::, _>>()?; + let ctx = PluginWriterContext { + segment: &segment, + memory_budget_in_bytes, + }; + let inverted_index = InvertedIndexPluginWriter::new(&ctx)?; + let fast_fields = FastFieldsPluginWriter::new(&ctx)?; + let store_writer = StorePluginWriter::new(&ctx)?; + let custom_plugins = segment + .index() + .custom_plugins() + .iter() + .map(|plugin| plugin.create_writer(&ctx)) + .collect::>>()?; Ok(Self { max_doc: 0, - ctx: IndexingContext::new(table_size), - per_field_postings_writers, - fieldnorms_writer: FieldNormsWriter::for_schema(&schema), - json_path_writer: JsonPathWriter::default(), - json_positions_per_path: IndexingPositionsPerPath::default(), - segment_serializer, - fast_field_writers: FastFieldsWriter::from_schema_and_tokenizer_manager( - &schema, - tokenizer_manager_fast_field, - )?, + segment, + inverted_index, + fast_fields, + store_writer, + custom_plugins, doc_opstamps: Vec::with_capacity(1_000), - per_field_text_analyzers, - term_buffer: IndexingTerm::with_capacity(16), schema, }) } + /// Resolves a custom plugin writer by one of its registered extensions. + /// Callers downcast the returned writer to their concrete type via `as_any_mut`. + pub fn custom_plugin_writer_mut(&mut self, extension: &str) -> Option<&mut dyn PluginWriter> { + let idx = self + .segment + .index() + .custom_plugins() + .iter() + .position(|plugin| plugin.extensions().contains(&extension))?; + self.custom_plugins + .get_mut(idx) + .map(|writer| writer.as_mut()) + } + /// Lay on disk the current content of the `SegmentWriter` /// /// Finalize consumes the `SegmentWriter`, so that it cannot /// be used afterwards. - pub fn finalize(mut self) -> crate::Result> { - self.fieldnorms_writer.fill_up_to_max_doc(self.max_doc); + pub fn finalize(self) -> crate::Result> { let mapping: Option = self - .segment_serializer - .segment() + .segment .index() .settings() .sort_by_field .clone() .map(|sort_by_field| get_doc_id_mapping_from_field(sort_by_field, &self)) .transpose()?; - remap_and_write( - self.schema, - &self.per_field_postings_writers, - self.ctx, - self.fast_field_writers, - &self.fieldnorms_writer, - self.segment_serializer, - mapping.as_ref(), - )?; + + // Serialize in write order: inverted index, fast fields, store, then custom plugins. + let segment = &self.segment; + let map = mapping.as_ref(); + Box::new(self.inverted_index).serialize(segment, map)?; + Box::new(self.fast_fields).serialize(segment, map)?; + Box::new(self.store_writer).serialize(segment, map)?; + for writer in self.custom_plugins { + writer.serialize(segment, map)?; + } + let doc_opstamps = remap_doc_opstamps(self.doc_opstamps, mapping.as_ref()); Ok(doc_opstamps) } @@ -165,209 +120,14 @@ impl SegmentWriter { /// Returns an estimation of the current memory usage of the segment writer. /// If the mem usage exceeds the `memory_budget`, the segment be serialized. pub fn mem_usage(&self) -> usize { - self.ctx.mem_usage() - + self.fieldnorms_writer.mem_usage() - + self.fast_field_writers.mem_usage() - + self.segment_serializer.mem_usage() - } - - fn index_document(&mut self, doc: &D) -> crate::Result<()> { - let doc_id = self.max_doc; - - // TODO: Can this be optimised a bit? - let vals_grouped_by_field = doc - .iter_fields_and_values() - .sorted_by_key(|(field, _)| *field) - .chunk_by(|(field, _)| *field); - - for (field, field_values) in &vals_grouped_by_field { - let values = field_values.map(|el| el.1); - - let field_entry = self.schema.get_field_entry(field); - let make_schema_error = || { - TantivyError::SchemaError(format!( - "Expected a {:?} for field {:?}", - field_entry.field_type().value_type(), - field_entry.name() - )) - }; - if !field_entry.is_indexed() { - continue; - } - - let (term_buffer, ctx) = (&mut self.term_buffer, &mut self.ctx); - let postings_writer: &mut dyn PostingsWriter = - self.per_field_postings_writers.get_for_field_mut(field); - term_buffer.clear_with_field(field); - - match field_entry.field_type() { - FieldType::Facet(_) => { - let mut facet_tokenizer = FacetTokenizer::default(); // this can be global - for value in values { - let value = value.as_value(); - - let facet_str = value.as_facet().ok_or_else(make_schema_error)?; - let mut facet_tokenizer = facet_tokenizer.token_stream(facet_str); - let mut indexing_position = IndexingPosition::default(); - postings_writer.index_text( - doc_id, - &mut facet_tokenizer, - term_buffer, - ctx, - &mut indexing_position, - ); - } - } - FieldType::Str(_) => { - let mut indexing_position = IndexingPosition::default(); - for value in values { - let value = value.as_value(); - - let mut token_stream = if let Some(text) = value.as_str() { - let text_analyzer = - &mut self.per_field_text_analyzers[field.field_id() as usize]; - text_analyzer.token_stream(text) - } else if let Some(tok_str) = value.into_pre_tokenized_text() { - BoxTokenStream::new(PreTokenizedStream::from(*tok_str.clone())) - } else { - continue; - }; - - assert!(term_buffer.is_empty()); - postings_writer.index_text( - doc_id, - &mut *token_stream, - term_buffer, - ctx, - &mut indexing_position, - ); - } - if field_entry.has_fieldnorms() { - self.fieldnorms_writer - .record(doc_id, field, indexing_position.num_tokens); - } - } - FieldType::U64(_) => { - let mut num_vals = 0; - for value in values { - let value = value.as_value(); - - num_vals += 1; - let u64_val = value.as_u64().ok_or_else(make_schema_error)?; - term_buffer.set_u64(u64_val); - postings_writer.subscribe(doc_id, 0u32, term_buffer, ctx); - } - if field_entry.has_fieldnorms() { - self.fieldnorms_writer.record(doc_id, field, num_vals); - } - } - FieldType::Date(_) => { - let mut num_vals = 0; - for value in values { - let value = value.as_value(); - - num_vals += 1; - let date_val = value.as_datetime().ok_or_else(make_schema_error)?; - term_buffer - .set_u64(date_val.truncate(DATE_TIME_PRECISION_INDEXED).to_u64()); - postings_writer.subscribe(doc_id, 0u32, term_buffer, ctx); - } - if field_entry.has_fieldnorms() { - self.fieldnorms_writer.record(doc_id, field, num_vals); - } - } - FieldType::I64(_) => { - let mut num_vals = 0; - for value in values { - let value = value.as_value(); - - num_vals += 1; - let i64_val = value.as_i64().ok_or_else(make_schema_error)?; - term_buffer.set_i64(i64_val); - postings_writer.subscribe(doc_id, 0u32, term_buffer, ctx); - } - if field_entry.has_fieldnorms() { - self.fieldnorms_writer.record(doc_id, field, num_vals); - } - } - FieldType::F64(_) => { - let mut num_vals = 0; - for value in values { - let value = value.as_value(); - num_vals += 1; - let f64_val = value.as_f64().ok_or_else(make_schema_error)?; - term_buffer.set_f64(f64_val); - postings_writer.subscribe(doc_id, 0u32, term_buffer, ctx); - } - if field_entry.has_fieldnorms() { - self.fieldnorms_writer.record(doc_id, field, num_vals); - } - } - FieldType::Bool(_) => { - let mut num_vals = 0; - for value in values { - let value = value.as_value(); - num_vals += 1; - let bool_val = value.as_bool().ok_or_else(make_schema_error)?; - term_buffer.set_bool(bool_val); - postings_writer.subscribe(doc_id, 0u32, term_buffer, ctx); - } - if field_entry.has_fieldnorms() { - self.fieldnorms_writer.record(doc_id, field, num_vals); - } - } - FieldType::Bytes(_) => { - let mut num_vals = 0; - for value in values { - let value = value.as_value(); - num_vals += 1; - let bytes = value.as_bytes().ok_or_else(make_schema_error)?; - term_buffer.set_bytes(bytes); - postings_writer.subscribe(doc_id, 0u32, term_buffer, ctx); - } - if field_entry.has_fieldnorms() { - self.fieldnorms_writer.record(doc_id, field, num_vals); - } - } - FieldType::JsonObject(json_options) => { - let text_analyzer = - &mut self.per_field_text_analyzers[field.field_id() as usize]; - - self.json_positions_per_path.clear(); - self.json_path_writer - .set_expand_dots(json_options.is_expand_dots_enabled()); - for json_value in values { - self.json_path_writer.clear(); - - index_json_value( - doc_id, - json_value, - text_analyzer, - term_buffer, - &mut self.json_path_writer, - postings_writer, - ctx, - &mut self.json_positions_per_path, - ); - } - } - FieldType::IpAddr(_) => { - let mut num_vals = 0; - for value in values { - let value = value.as_value(); - - num_vals += 1; - let ip_addr = value.as_ip_addr().ok_or_else(make_schema_error)?; - term_buffer.set_ip_addr(ip_addr); - postings_writer.subscribe(doc_id, 0u32, term_buffer, ctx); - } - if field_entry.has_fieldnorms() { - self.fieldnorms_writer.record(doc_id, field, num_vals); - } - } - } - } - Ok(()) + self.inverted_index.mem_usage() + + self.fast_fields.mem_usage() + + self.store_writer.mem_usage() + + self + .custom_plugins + .iter() + .map(|writer| writer.mem_usage()) + .sum::() } /// Indexes a new document @@ -378,11 +138,23 @@ impl SegmentWriter { add_operation: AddOperation, ) -> crate::Result<()> { let AddOperation { document, opstamp } = add_operation; + let doc_id = self.max_doc; self.doc_opstamps.push(opstamp); - self.fast_field_writers.add_document(&document)?; - self.index_document(&document)?; - let doc_writer = self.segment_serializer.get_store_writer(); - doc_writer.store(&document, &self.schema)?; + + // Built-in writers index the document generically — zero-copy for any `Document` type. + self.inverted_index.index_document(doc_id, &document)?; + self.fast_fields.index_document(&document)?; + self.store_writer.store(&document, &self.schema)?; + + // Custom plugins are trait objects, so they take the document type-erased. Leaves stay + // borrowed; a plugin that knows the concrete type recovers it for free via `as_any`. + if !self.custom_plugins.is_empty() { + let erased: &dyn ErasedDocument = &document; + for writer in &mut self.custom_plugins { + writer.add_document(doc_id, erased, &self.schema)?; + } + } + self.max_doc += 1; Ok(()) } @@ -408,73 +180,6 @@ impl SegmentWriter { } } -/// This method is used as a trick to workaround the borrow checker -/// Writes a view of a segment by pushing information -/// to the `SegmentSerializer`. -/// -/// `doc_id_map` is used to map to the new doc_id order. -fn remap_and_write( - schema: Schema, - per_field_postings_writers: &PerFieldPostingsWriter, - ctx: IndexingContext, - fast_field_writers: FastFieldsWriter, - fieldnorms_writer: &FieldNormsWriter, - mut serializer: SegmentSerializer, - doc_id_map: Option<&DocIdMapping>, -) -> crate::Result<()> { - debug!("remap-and-write"); - if let Some(fieldnorms_serializer) = serializer.extract_fieldnorms_serializer() { - fieldnorms_writer.serialize(fieldnorms_serializer, doc_id_map)?; - } - let fieldnorm_data = serializer - .segment() - .open_read(SegmentComponent::FieldNorms)?; - let fieldnorm_readers = FieldNormReaders::open(fieldnorm_data)?; - serialize_postings( - ctx, - schema, - per_field_postings_writers, - fieldnorm_readers, - doc_id_map, - serializer.get_postings_serializer(), - )?; - debug!("fastfield-serialize"); - fast_field_writers.serialize(serializer.get_fast_field_write(), doc_id_map)?; - - // finalize temp docstore and create version, which reflects the doc_id_map - if let Some(doc_id_map) = doc_id_map { - debug!("resort-docstore"); - let store_write = serializer - .segment_mut() - .open_write(SegmentComponent::Store)?; - let settings = serializer.segment().index().settings(); - let store_writer = StoreWriter::new( - store_write, - settings.docstore_compression, - settings.docstore_blocksize, - settings.docstore_compress_dedicated_thread, - )?; - let old_store_writer = std::mem::replace(&mut serializer.store_writer, store_writer); - old_store_writer.close()?; - let store_read = StoreReader::open( - serializer - .segment() - .open_read(SegmentComponent::TempStore)?, - 1, /* The docstore is configured to have one doc per block, and each doc is - * accessed only once: we don't need caching. */ - )?; - for old_doc_id in doc_id_map.iter_old_doc_ids() { - let doc_bytes = store_read.get_document_bytes(old_doc_id)?; - serializer.get_store_writer().store_bytes(&doc_bytes)?; - } - } - - debug!("serializer-close"); - serializer.close()?; - - Ok(()) -} - #[cfg(test)] mod tests { use std::collections::BTreeMap; @@ -503,7 +208,7 @@ mod tests { #[test] #[cfg(not(feature = "compare_hash_only"))] fn test_hashmap_size() { - use super::compute_initial_table_size; + use crate::postings::compute_initial_table_size; assert_eq!(compute_initial_table_size(100_000).unwrap(), 1 << 12); assert_eq!(compute_initial_table_size(1_000_000).unwrap(), 1 << 15); assert_eq!(compute_initial_table_size(15_000_000).unwrap(), 1 << 19); diff --git a/src/indexer/single_segment_index_writer.rs b/src/indexer/single_segment_index_writer.rs index 673accae32..6b523a8537 100644 --- a/src/indexer/single_segment_index_writer.rs +++ b/src/indexer/single_segment_index_writer.rs @@ -42,8 +42,10 @@ impl SingleSegmentIndexWriter { self.segment_writer.finalize()?; let segment: Segment = self.segment.with_max_doc(max_doc); let index = segment.index(); + let previous_meta = index.load_metas()?; let index_meta = IndexMeta { index_settings: index.settings().clone(), + persisted_custom_extensions: previous_meta.persisted_custom_extensions.clone(), segments: vec![segment.meta().clone()], schema: index.schema(), opstamp: 0, diff --git a/src/lib.rs b/src/lib.rs index 1f96237384..53a4fb10f0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -198,6 +198,7 @@ pub mod index; pub mod positions; pub mod postings; +pub mod plugin; /// Module containing the different query implementations. pub mod query; pub mod schema; @@ -230,6 +231,7 @@ pub use crate::index::{ Segment, SegmentMeta, SegmentReader, }; pub use crate::indexer::{IndexWriter, SingleSegmentIndexWriter}; +pub use crate::plugin::{PluginMergeContext, PluginWriter, PluginWriterContext, SegmentPlugin}; pub use crate::schema::{Document, TantivyDocument, Term}; /// Index format version. diff --git a/src/plugin.rs b/src/plugin.rs new file mode 100644 index 0000000000..783e28f9bd --- /dev/null +++ b/src/plugin.rs @@ -0,0 +1,515 @@ +//! Extensible segment component plugin system. +//! +//! This module defines the [`SegmentPlugin`] trait and supporting types that allow +//! custom data structures to participate in the segment lifecycle (write, read, merge). +//! +//! The built-in components (postings, fast fields, field norms, store) are themselves +//! implemented as plugins; external code attaches new data to segments through the same +//! trait without modifying tantivy internals. + +use std::any::Any; +use std::collections::BTreeMap; + +use common::HasLen; + +use crate::index::{IndexSettings, SegmentComponent, SegmentReader}; +use crate::indexer::doc_id_mapping::SegmentDocIdMapping; +use crate::schema::document::ErasedDocument; +use crate::schema::Schema; +use crate::space_usage::ComponentSpaceUsage; +use crate::{DocId, Segment}; + +/// A pluggable segment component that participates in writing and merging. +/// +/// Each plugin manages one or more files within a segment. The plugin is a factory +/// that creates writers and handles merging. The actual data APIs are +/// component-specific and accessed via downcasting on the concrete types. +/// +/// # Ordering +/// +/// Plugins are written and merged in [`Index::all_plugins`] order: the built-ins first +/// (field norms, postings, fast fields, store — postings reads field norms back from +/// disk, so field norms come first), then custom plugins in the order they were +/// registered. A plugin can read any earlier plugin's output from disk; custom plugins +/// run after all built-ins, so they can read built-in output, and a built-in never +/// depends on a custom plugin. +/// +/// [`Index::all_plugins`]: crate::Index::all_plugins +pub trait SegmentPlugin: Send + Sync + 'static { + /// File extensions this component manages (e.g., `["idx", "pos", "term"]` for postings). + fn extensions(&self) -> &[&str]; + + /// Create a writer for accumulating and serializing data during indexing. + /// + /// Returns a type-erased writer. The `SegmentWriter` will downcast to the concrete + /// type when it needs to call component-specific APIs (e.g., feeding terms to + /// the postings writer). + fn create_writer(&self, ctx: &PluginWriterContext) -> crate::Result>; + + /// Merge data from multiple source segments into a target segment. + fn merge(&self, ctx: PluginMergeContext) -> crate::Result<()>; + + /// Report on-disk space usage of this component, keyed by component name. + /// + /// The returned entries are merged into [`SegmentSpaceUsage`]. The default + /// implementation emits one [`ComponentSpaceUsage::Basic`] entry per file in + /// [`extensions()`](Self::extensions); built-in plugins override this to report + /// richer per-field breakdowns under the keys the named accessors expect. + /// + /// [`SegmentSpaceUsage`]: crate::space_usage::SegmentSpaceUsage + fn space_usage( + &self, + segment_reader: &SegmentReader, + ) -> crate::Result> { + let mut usage = BTreeMap::new(); + for &ext in self.extensions() { + let file = segment_reader.open_read(SegmentComponent::Custom(ext.to_string()))?; + usage.insert( + ext.to_string(), + ComponentSpaceUsage::Basic(file.len().into()), + ); + } + Ok(usage) + } +} + +/// Writer for a single component within a segment. +/// +/// The writer accumulates data during indexing (via [`add_document`](Self::add_document) +/// and component-specific APIs on the concrete type) and serializes it to segment files +/// during finalization. +pub trait PluginWriter: Send + Any { + /// Records a single document during indexing. + /// + /// Called once per document added to the segment, in doc-id order, for every custom + /// plugin writer. The document is passed type-erased as [`&dyn ErasedDocument`] — custom + /// plugins are trait objects, so they cannot take a generic `Document`. A plugin that + /// knows the concrete document type can recover it for free with + /// `doc.as_any().downcast_ref::()`; otherwise it can walk `doc.erased_fields()`. + /// The default is a no-op; override it to accumulate per-document state. + fn add_document( + &mut self, + _doc_id: DocId, + _doc: &dyn ErasedDocument, + _schema: &Schema, + ) -> crate::Result<()> { + Ok(()) + } + + /// Serialize accumulated data to segment files and finalize them. + /// Called once, during `SegmentWriter::finalize()`. + /// + /// Consumes the writer: it is the terminal step of the writer's lifecycle, responsible + /// for both writing the segment files and closing their handles. + fn serialize( + self: Box, + segment: &Segment, + doc_id_map: Option<&crate::indexer::doc_id_mapping::DocIdMapping>, + ) -> crate::Result<()>; + + /// Current memory usage of this writer. + fn mem_usage(&self) -> usize; + + // Downcast support for accessing component-specific APIs. Once the crate MSRV + // reaches Rust 1.86, these can be dropped: trait upcasting lets callers coerce + // `&dyn PluginWriter` to `&dyn Any` and call `downcast_ref`/`downcast_mut` directly. + fn as_any(&self) -> &dyn Any; + fn as_any_mut(&mut self) -> &mut dyn Any; +} + +/// Context provided to [`SegmentPlugin::create_writer`]. +/// +/// The schema, settings, and directory are reachable from `segment` +/// (`segment.schema()`, `segment.index().settings()`, `segment.index().directory()`). +pub struct PluginWriterContext<'a> { + /// The segment being written to. + pub segment: &'a Segment, + /// Per-thread indexing memory budget in bytes. Plugins that keep an in-memory arena + /// (e.g. the inverted index) size it from this. + pub memory_budget_in_bytes: usize, +} + +/// Context provided to [`SegmentPlugin::merge`]. +pub struct PluginMergeContext<'a> { + pub readers: &'a [SegmentReader], + pub doc_id_mapping: &'a SegmentDocIdMapping, + pub target_segment: &'a Segment, + pub schema: &'a Schema, + pub settings: &'a IndexSettings, +} + +#[cfg(test)] +mod tests { + //! Round-trip integration test for the segment plugin system. + //! + //! Defines a custom marker plugin, then verifies it works through the + //! full lifecycle: write → read → merge → read. + + use std::sync::Arc; + + use super::*; + use crate::index::SegmentComponent; + use crate::schema::{Schema, STORED, TEXT}; + use crate::{Index, IndexWriter}; + + const MARKER: u32 = 0xDEADBEEF; + + /// A simple plugin that writes a fixed marker to a custom file. + struct MarkerPlugin; + + impl SegmentPlugin for MarkerPlugin { + fn extensions(&self) -> &[&str] { + &["marker"] + } + + fn create_writer( + &self, + _ctx: &PluginWriterContext, + ) -> crate::Result> { + Ok(Box::new(MarkerWriter)) + } + + fn merge(&self, ctx: PluginMergeContext) -> crate::Result<()> { + let component = SegmentComponent::Custom("marker".to_string()); + let mut write = ctx.target_segment.open_write(component)?; + use std::io::Write; + write.write_all(&MARKER.to_le_bytes())?; + common::TerminatingWrite::terminate(write)?; + Ok(()) + } + } + + struct MarkerWriter; + + impl PluginWriter for MarkerWriter { + fn serialize( + self: Box, + segment: &Segment, + _doc_id_map: Option<&crate::indexer::doc_id_mapping::DocIdMapping>, + ) -> crate::Result<()> { + let component = SegmentComponent::Custom("marker".to_string()); + let mut write = segment.open_write(component)?; + use std::io::Write; + write.write_all(&MARKER.to_le_bytes())?; + common::TerminatingWrite::terminate(write)?; + Ok(()) + } + + fn mem_usage(&self) -> usize { + std::mem::size_of::() + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } + } + + #[test] + fn test_plugin_full_lifecycle() -> crate::Result<()> { + use crate::indexer::NoMergePolicy; + + let mut schema_builder = Schema::builder(); + let text_field = schema_builder.add_text_field("text", TEXT | STORED); + let schema = schema_builder.build(); + + let plugin: Arc = Arc::new(MarkerPlugin); + let index = Index::builder() + .schema(schema) + .register_plugin(plugin) + .create_in_ram()?; + + assert!(index.all_plugins().count() >= 2); + assert!( + index + .all_plugins() + .any(|plugin| plugin.extensions().contains(&"marker")), + "marker plugin should be registered" + ); + assert!( + index + .all_plugins() + .any(|plugin| plugin.extensions().contains(&"fieldnorm")), + "fieldnorms built-in plugin should be registered" + ); + + // write: two commits, no auto-merge, so we get two distinct segments. + let mut writer: IndexWriter = index.writer_with_num_threads(1, 15_000_000)?; + writer.set_merge_policy(Box::new(NoMergePolicy)); + writer.add_document(crate::doc!(text_field => "hello world"))?; + writer.add_document(crate::doc!(text_field => "foo bar"))?; + writer.commit()?; + writer.add_document(crate::doc!(text_field => "baz qux"))?; + writer.commit()?; + + // read: each segment carries the marker written by MarkerWriter::serialize. + let searcher = index.reader()?.searcher(); + assert_eq!(searcher.num_docs(), 3); + let segment_readers = searcher.segment_readers(); + assert_eq!(segment_readers.len(), 2); + for segment_reader in segment_readers { + let data = segment_reader + .open_read(SegmentComponent::Custom("marker".to_string()))? + .read_bytes()?; + assert_eq!( + u32::from_le_bytes([data[0], data[1], data[2], data[3]]), + MARKER + ); + } + + // merge: exercises MarkerPlugin::merge. + writer.merge(&index.searchable_segment_ids()?).wait()?; + + // read: the merged segment carries the marker written by MarkerPlugin::merge. + let searcher = index.reader()?.searcher(); + assert_eq!(searcher.num_docs(), 3); + let segment_readers = searcher.segment_readers(); + assert_eq!(segment_readers.len(), 1); + let data = segment_readers[0] + .open_read(SegmentComponent::Custom("marker".to_string()))? + .read_bytes()?; + assert_eq!( + u32::from_le_bytes([data[0], data[1], data[2], data[3]]), + MARKER + ); + + Ok(()) + } + + #[test] + fn test_plugin_extensions() { + let plugin = MarkerPlugin; + assert_eq!(plugin.extensions(), &["marker"]); + } + + #[test] + fn test_reopen_without_plugin_fails_closed() -> crate::Result<()> { + use crate::directory::RamDirectory; + use crate::TantivyError; + + let mut schema_builder = Schema::builder(); + let text_field = schema_builder.add_text_field("text", TEXT | STORED); + let schema = schema_builder.build(); + + // Build an index with the custom plugin and persist a segment. + let dir = RamDirectory::create(); + let plugin: Arc = Arc::new(MarkerPlugin); + let index = Index::builder() + .schema(schema) + .register_plugin(plugin) + .create(dir.clone())?; + { + let mut writer: IndexWriter = index.writer_with_num_threads(1, 15_000_000)?; + writer.add_document(crate::doc!(text_field => "hello world"))?; + writer.commit()?; + } + + // The index records, index-wide, that it requires the "marker" extension. + let segment_metas = index.searchable_segment_metas()?; + assert_eq!(segment_metas.len(), 1); + assert_eq!( + index.load_metas()?.persisted_custom_extensions, + vec!["marker".to_string()] + ); + + // Reopen without re-registering the plugin: writing must fail closed + // rather than silently dropping the plugin's data. + let reopened = Index::open(dir.clone())?; + let err = reopened + .writer_with_num_threads::(1, 15_000_000) + .err() + .expect("writer creation should fail when the plugin is not registered"); + assert!( + matches!(err, TantivyError::MissingPlugin(ref exts) if exts.contains("marker")), + "expected MissingPlugin error, got {err:?}" + ); + + // Re-registering the plugin clears the guard. + let mut reopened = reopened; + reopened.register_plugin(Arc::new(MarkerPlugin)); + let _writer: IndexWriter = reopened.writer_with_num_threads(1, 15_000_000)?; + + Ok(()) + } + + #[test] + fn test_add_plugin_to_non_empty_index_fails_closed() -> crate::Result<()> { + use crate::directory::RamDirectory; + use crate::TantivyError; + + let mut schema_builder = Schema::builder(); + let text_field = schema_builder.add_text_field("text", TEXT | STORED); + + // Build an index WITHOUT the plugin and persist a segment. + let dir = RamDirectory::create(); + let index = Index::builder() + .schema(schema_builder.build()) + .create(dir.clone())?; + { + let mut writer: IndexWriter = index.writer_with_num_threads(1, 15_000_000)?; + writer.add_document(crate::doc!(text_field => "hello world"))?; + writer.commit()?; + } + + // Registering the plugin now — after the index has data — must fail closed: the + // existing segment has no "marker" component, so the plugin set can't change. + let mut reopened = Index::open(dir)?; + reopened.register_plugin(Arc::new(MarkerPlugin)); + let err = reopened + .writer_with_num_threads::(1, 15_000_000) + .err() + .expect("writer creation should fail when a plugin is added to a non-empty index"); + assert!( + matches!(err, TantivyError::UnexpectedPlugin(ref exts) if exts.contains("marker")), + "expected UnexpectedPlugin error, got {err:?}" + ); + + Ok(()) + } + + #[test] + fn test_conflicting_plugins_fail_closed() -> crate::Result<()> { + use crate::TantivyError; + + let mut schema_builder = Schema::builder(); + schema_builder.add_text_field("text", TEXT | STORED); + let schema = schema_builder.build(); + + // Two plugins claiming the same "marker" extension would make their writers + // contend for the same segment file. + let index = Index::builder() + .schema(schema) + .register_plugin(Arc::new(MarkerPlugin)) + .register_plugin(Arc::new(MarkerPlugin)) + .create_in_ram()?; + + let err = index + .writer_with_num_threads::(1, 15_000_000) + .err() + .expect("writer creation should fail when two plugins claim one extension"); + assert!( + matches!(err, TantivyError::ConflictingPlugins(ref exts) if exts.contains("marker")), + "expected ConflictingPlugins error, got {err:?}" + ); + Ok(()) + } + + #[test] + fn test_custom_plugin_indexes_custom_document() -> crate::Result<()> { + use std::collections::BTreeMap; + + use crate::indexer::operation::AddOperation; + use crate::indexer::SegmentWriter; + use crate::schema::{Field, OwnedValue}; + + let mut schema_builder = Schema::builder(); + let text_field = schema_builder.add_text_field("text", TEXT | STORED); + let index = Index::builder() + .schema(schema_builder.build()) + .register_plugin(Arc::new(MarkerPlugin)) + .create_in_ram()?; + + let segment = index.new_segment(); + let mut segment_writer = SegmentWriter::for_segment(15_000_000, segment)?; + + // A custom (non-`TantivyDocument`) document type indexes fine even with a custom plugin + // registered: built-ins take it generically, and the custom plugin receives it + // type-erased (`&dyn ErasedDocument`) — no materialization, no error. + let mut document: BTreeMap = BTreeMap::new(); + document.insert(text_field, "hello".into()); + segment_writer.add_document(AddOperation { + opstamp: 0, + document, + })?; + + Ok(()) + } + + #[test] + fn test_reserved_extension_plugins_fail_closed() -> crate::Result<()> { + use crate::TantivyError; + + struct ReservedExtPlugin(&'static str); + + impl SegmentPlugin for ReservedExtPlugin { + fn extensions(&self) -> &[&str] { + std::slice::from_ref(&self.0) + } + + fn create_writer( + &self, + _ctx: &PluginWriterContext, + ) -> crate::Result> { + unreachable!("guard rejects the reserved extension before writer creation") + } + + fn merge(&self, _ctx: PluginMergeContext) -> crate::Result<()> { + unreachable!("guard rejects the reserved extension before merge") + } + } + + // The temp store (`store.temp`) and delete bitset (`del`) are not owned by any + // plugin, so a custom plugin claiming one would contend for the reserved file. + for reserved in ["temp", "store.temp", "del"] { + let mut schema_builder = Schema::builder(); + schema_builder.add_text_field("text", TEXT | STORED); + let index = Index::builder() + .schema(schema_builder.build()) + .register_plugin(Arc::new(ReservedExtPlugin(reserved))) + .create_in_ram()?; + + let err = index + .writer_with_num_threads::(1, 15_000_000) + .err() + .unwrap_or_else(|| { + panic!("writer creation should fail when a plugin claims `{reserved}`") + }); + assert!( + matches!(err, TantivyError::ConflictingPlugins(ref exts) if exts.contains(reserved)), + "expected ConflictingPlugins error for `{reserved}`, got {err:?}" + ); + } + Ok(()) + } + + #[test] + fn test_merge_indices_mismatched_plugins_fails_closed() -> crate::Result<()> { + use crate::directory::RamDirectory; + use crate::indexer::merge_indices; + use crate::TantivyError; + + // Two indices with identical schemas but different plugin sets: only the first + // registers the marker plugin. Merging must fail closed rather than silently drop + // the marker component (or read a missing one), since the merged output would carry + // only the first index's plugin set. + let build = |with_plugin: bool| -> crate::Result { + let mut schema_builder = Schema::builder(); + let text_field = schema_builder.add_text_field("text", TEXT | STORED); + let mut builder = Index::builder().schema(schema_builder.build()); + if with_plugin { + builder = builder.register_plugin(Arc::new(MarkerPlugin)); + } + let index = builder.create_in_ram()?; + let mut writer: IndexWriter = index.writer_with_num_threads(1, 15_000_000)?; + writer.add_document(crate::doc!(text_field => "hello world"))?; + writer.commit()?; + Ok(index) + }; + + let with_marker = build(true)?; + let without_marker = build(false)?; + + let err = merge_indices(&[with_marker, without_marker], RamDirectory::create()) + .err() + .expect("merge should fail when source indices register different plugin sets"); + assert!( + matches!(err, TantivyError::InvalidArgument(ref msg) if msg.contains("plugin sets")), + "expected InvalidArgument about plugin sets, got {err:?}" + ); + + Ok(()) + } +} diff --git a/src/postings/indexing_context.rs b/src/postings/indexing_context.rs index 2675476f37..395050af63 100644 --- a/src/postings/indexing_context.rs +++ b/src/postings/indexing_context.rs @@ -1,6 +1,28 @@ use stacker::{ArenaHashMap, MemoryArena}; use crate::indexer::path_to_unordered_id::PathToUnorderedId; +use crate::postings::compute_table_memory_size; + +/// Computes the initial size of the term hash table. +/// +/// Returns the recommended initial table size as a power of 2. +/// +/// Note this is a very dumb way to compute log2, but it is easier to proofread that way. +pub(crate) fn compute_initial_table_size(per_thread_memory_budget: usize) -> crate::Result { + let table_memory_upper_bound = per_thread_memory_budget / 3; + (10..20) // We cap it at 2^19 = 512K capacity. + // TODO: There are cases where this limit causes a + // reallocation in the hashmap. Check if this affects performance. + .map(|power| 1 << power) + .take_while(|capacity| compute_table_memory_size(*capacity) < table_memory_upper_bound) + .last() + .ok_or_else(|| { + crate::TantivyError::InvalidArgument(format!( + "per thread memory budget (={per_thread_memory_budget}) is too small. Raise the \ + memory budget or lower the number of threads." + )) + }) +} /// IndexingContext contains all of the transient memory arenas /// required for building the inverted index. diff --git a/src/postings/mod.rs b/src/postings/mod.rs index 13b6761cfa..45b6af94af 100644 --- a/src/postings/mod.rs +++ b/src/postings/mod.rs @@ -23,7 +23,7 @@ pub(crate) use loaded_postings::LoadedPostings; pub(crate) use stacker::compute_table_memory_size; pub use self::block_segment_postings::BlockSegmentPostings; -pub(crate) use self::indexing_context::IndexingContext; +pub(crate) use self::indexing_context::{compute_initial_table_size, IndexingContext}; pub(crate) use self::per_field_postings_writer::PerFieldPostingsWriter; pub use self::postings::Postings; pub(crate) use self::postings_writer::{serialize_postings, IndexingPosition, PostingsWriter}; diff --git a/src/postings/serializer.rs b/src/postings/serializer.rs index 726cce03b2..f522dfa93d 100644 --- a/src/postings/serializer.rs +++ b/src/postings/serializer.rs @@ -55,7 +55,7 @@ pub struct InvertedIndexSerializer { impl InvertedIndexSerializer { /// Open a new `InvertedIndexSerializer` for the given segment - pub fn open(segment: &mut Segment) -> crate::Result { + pub fn open(segment: &Segment) -> crate::Result { use crate::index::SegmentComponent::{Positions, Postings, Terms}; let inv_index_serializer = InvertedIndexSerializer { terms_write: CompositeWrite::wrap(segment.open_write(Terms)?), diff --git a/src/schema/document/default_document.rs b/src/schema/document/default_document.rs index 915b685aa7..b012c59146 100644 --- a/src/schema/document/default_document.rs +++ b/src/schema/document/default_document.rs @@ -671,6 +671,10 @@ impl Document for CompactDoc { container: self, } } + + fn into_tantivy_document(self) -> TantivyDocument { + self + } } /// A helper wrapper for creating an iterator over the field values diff --git a/src/schema/document/erased.rs b/src/schema/document/erased.rs new file mode 100644 index 0000000000..03d5298ed0 --- /dev/null +++ b/src/schema/document/erased.rs @@ -0,0 +1,101 @@ +//! Object-safe, type-erased views of documents and their values. +//! +//! [`Document`] and [`Value`] both use GATs, so neither can be turned into a trait object. +//! These mirror types erase the generic parameter — the nested array/object iterators are +//! boxed, while leaf values stay borrowed — so a plugin behind a `dyn` boundary can read an +//! arbitrary document without it being materialized into a concrete type. A plugin that +//! knows the concrete document type can recover it for free via [`ErasedDocument::as_any`] +//! and skip the erased walk entirely. + +use std::any::Any; + +use super::{Document, ReferenceValue, ReferenceValueLeaf, Value}; +use crate::schema::Field; + +/// A type-erased mirror of [`ReferenceValue`]. +/// +/// Leaves are the same concrete, mostly-borrowed enum as `ReferenceValue`, so they cost +/// nothing. Only the nested cases carry a boxed iterator, so a flat document allocates once +/// (for the field iterator) and nested documents allocate once per container. +pub enum ErasedValue<'a> { + /// A leaf value, borrowed from the document. + Leaf(ReferenceValueLeaf<'a>), + /// An array, walked lazily through a boxed iterator. + Array(Box> + 'a>), + /// A nested object, walked lazily through a boxed iterator. + Object(Box)> + 'a>), +} + +/// Object-safe counterpart to [`Document`]. +/// +/// A blanket impl covers every `Document`, so `&document as &dyn ErasedDocument` always +/// works. Consumers either downcast via [`as_any`](ErasedDocument::as_any) (zero cost, when +/// they know the concrete type) or walk [`erased_fields`](ErasedDocument::erased_fields). +pub trait ErasedDocument { + /// Iterates the document's `(field, value)` pairs in erased form. + fn erased_fields<'a>(&'a self) -> Box)> + 'a>; + + /// Recovers the concrete document. `Document: 'static`, so this is free — a consumer that + /// knows the type can `as_any().downcast_ref::()` and avoid the erased walk. + fn as_any(&self) -> &dyn Any; +} + +fn erase_value<'a, V: Value<'a> + 'a>(value: ReferenceValue<'a, V>) -> ErasedValue<'a> { + match value { + ReferenceValue::Leaf(leaf) => ErasedValue::Leaf(leaf), + ReferenceValue::Array(iter) => { + ErasedValue::Array(Box::new(iter.map(|value| erase_value(value.as_value())))) + } + ReferenceValue::Object(iter) => ErasedValue::Object(Box::new( + iter.map(|(key, value)| (key, erase_value(value.as_value()))), + )), + } +} + +impl ErasedDocument for D { + fn erased_fields<'a>(&'a self) -> Box)> + 'a> { + Box::new( + self.iter_fields_and_values() + .map(|(field, value)| (field, erase_value(value.as_value()))), + ) + } + + fn as_any(&self) -> &dyn Any { + self + } +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + + use super::*; + use crate::schema::{Field, OwnedValue}; + + #[test] + fn test_erased_document_walk_and_downcast() { + let field = Field::from_field_id(0); + let mut document: BTreeMap = BTreeMap::new(); + document.insert(field, "hello".into()); + + // Any `Document` is an `ErasedDocument` via the blanket impl. + let erased: &dyn ErasedDocument = &document; + + // The erased walk yields the leaf borrowed straight from the document — no copy. + let fields: Vec<(Field, ErasedValue)> = erased.erased_fields().collect(); + assert_eq!(fields.len(), 1); + match &fields[0] { + (got_field, ErasedValue::Leaf(ReferenceValueLeaf::Str(text))) => { + assert_eq!(*got_field, field); + assert_eq!(*text, "hello"); + } + _ => panic!("expected a borrowed string leaf"), + } + + // A consumer that knows the concrete type recovers it for free, skipping the walk. + assert!(erased + .as_any() + .downcast_ref::>() + .is_some()); + } +} diff --git a/src/schema/document/mod.rs b/src/schema/document/mod.rs index 8168ee8112..e98716ba19 100644 --- a/src/schema/document/mod.rs +++ b/src/schema/document/mod.rs @@ -161,6 +161,7 @@ mod de; mod default_document; +mod erased; mod existing_type_impls; mod owned_value; mod se; @@ -177,6 +178,7 @@ pub use self::de::{ pub use self::default_document::{ CompactDocArrayIter, CompactDocObjectIter, CompactDocValue, DocParsingError, TantivyDocument, }; +pub use self::erased::{ErasedDocument, ErasedValue}; pub use self::owned_value::OwnedValue; pub(crate) use self::se::BinaryDocumentSerializer; pub use self::value::{ReferenceValue, ReferenceValueLeaf, Value}; @@ -195,6 +197,21 @@ pub trait Document: Send + Sync + 'static { /// Get an iterator iterating over all fields and values in a document. fn iter_fields_and_values(&self) -> Self::FieldsValuesIter<'_>; + /// Convert this document into the canonical [`TantivyDocument`] representation. + /// + /// The indexer calls this once per document so it can hand a single concrete type to + /// every [`PluginWriter`](crate::plugin::PluginWriter). The default materializes a new + /// `TantivyDocument`; `TantivyDocument` overrides it to return itself, so indexing the + /// default document type costs nothing extra. + fn into_tantivy_document(self) -> TantivyDocument + where Self: Sized { + let mut doc = TantivyDocument::new(); + for (field, value) in self.iter_fields_and_values() { + doc.add_field_value(field, value); + } + doc + } + /// Sort and groups the field_values by field. /// /// The result of this method is not cached and is diff --git a/src/space_usage/mod.rs b/src/space_usage/mod.rs index 1dc4131569..ab98d4824f 100644 --- a/src/space_usage/mod.rs +++ b/src/space_usage/mod.rs @@ -9,6 +9,7 @@ use std::collections::btree_map::Entry; use std::collections::BTreeMap; +use std::sync::LazyLock; use columnar::ColumnSpaceUsage; use common::ByteCount; @@ -16,7 +17,25 @@ use serde::{Deserialize, Serialize}; use crate::index::SegmentComponent; +/// Component keys used by the built-in plugins when reporting [`SegmentSpaceUsage`]. +/// +/// These are shared between the plugins that produce the entries and the +/// [`SegmentSpaceUsage`] accessors that look them up. +pub(crate) const TERMDICT: &str = "termdict"; +pub(crate) const POSTINGS: &str = "postings"; +pub(crate) const POSITIONS: &str = "positions"; +pub(crate) const FAST_FIELDS: &str = "fast_fields"; +pub(crate) const FIELDNORMS: &str = "fieldnorms"; +pub(crate) const STORE: &str = "store"; +pub(crate) const DELETES: &str = "deletes"; + +static EMPTY_PER_FIELD: LazyLock = + LazyLock::new(|| PerFieldSpaceUsage::new(Vec::new())); +static EMPTY_STORE: LazyLock = + LazyLock::new(|| StoreSpaceUsage::new(ByteCount::default(), ByteCount::default())); + /// Enum containing any of the possible space usage results for segment components. +#[derive(Clone, Debug, Serialize, Deserialize)] pub enum ComponentSpaceUsage { /// Data is stored per field in a uniform way PerField(PerFieldSpaceUsage), @@ -26,6 +45,17 @@ pub enum ComponentSpaceUsage { Basic(ByteCount), } +impl ComponentSpaceUsage { + /// Total bytes used by this component. + pub fn total(&self) -> ByteCount { + match self { + ComponentSpaceUsage::PerField(usage) => usage.total(), + ComponentSpaceUsage::Store(usage) => usage.total(), + ComponentSpaceUsage::Basic(bytes) => *bytes, + } + } +} + /// Represents combined space usage of an entire searcher and its component segments. #[derive(Clone, Debug, Serialize, Deserialize)] pub struct SearcherSpaceUsage { @@ -61,51 +91,31 @@ impl SearcherSpaceUsage { } /// Represents combined space usage for all of the large components comprising a segment. +/// +/// The per-component usage is assembled from each registered [`SegmentPlugin`] (plus the +/// non-plugin `deletes` entry) and keyed by component name. The named accessors below are +/// thin lookups into that map. +/// +/// [`SegmentPlugin`]: crate::plugin::SegmentPlugin #[derive(Clone, Debug, Serialize, Deserialize)] pub struct SegmentSpaceUsage { num_docs: u32, - - termdict: PerFieldSpaceUsage, - postings: PerFieldSpaceUsage, - positions: PerFieldSpaceUsage, - fast_fields: PerFieldSpaceUsage, - fieldnorms: PerFieldSpaceUsage, - - store: StoreSpaceUsage, - - deletes: ByteCount, - + components: BTreeMap, total: ByteCount, } impl SegmentSpaceUsage { - #[expect(clippy::too_many_arguments)] - pub(crate) fn new( + pub(crate) fn from_components( num_docs: u32, - termdict: PerFieldSpaceUsage, - postings: PerFieldSpaceUsage, - positions: PerFieldSpaceUsage, - fast_fields: PerFieldSpaceUsage, - fieldnorms: PerFieldSpaceUsage, - store: StoreSpaceUsage, - deletes: ByteCount, + components: BTreeMap, ) -> SegmentSpaceUsage { - let total = termdict.total() - + postings.total() - + positions.total() - + fast_fields.total() - + fieldnorms.total() - + store.total() - + deletes; + let total = components + .values() + .map(ComponentSpaceUsage::total) + .fold(ByteCount::default(), |acc, bytes| acc + bytes); SegmentSpaceUsage { num_docs, - termdict, - postings, - positions, - fast_fields, - fieldnorms, - store, - deletes, + components, total, } } @@ -115,18 +125,21 @@ impl SegmentSpaceUsage { /// Clones the underlying data. /// Use the components directly if this is somehow in performance critical code. pub fn component(&self, component: SegmentComponent) -> ComponentSpaceUsage { - use self::ComponentSpaceUsage::*; use crate::index::SegmentComponent::*; - match component { - Postings => PerField(self.postings().clone()), - Positions => PerField(self.positions().clone()), - FastFields => PerField(self.fast_fields().clone()), - FieldNorms => PerField(self.fieldnorms().clone()), - Terms => PerField(self.termdict().clone()), - SegmentComponent::Store => ComponentSpaceUsage::Store(self.store().clone()), - SegmentComponent::TempStore => ComponentSpaceUsage::Store(self.store().clone()), - Delete => Basic(self.deletes()), - } + let key: &str = match &component { + Postings => POSTINGS, + Positions => POSITIONS, + FastFields => FAST_FIELDS, + FieldNorms => FIELDNORMS, + Terms => TERMDICT, + Store | TempStore => STORE, + Delete => DELETES, + Custom(ext) => ext.as_str(), + }; + self.components + .get(key) + .cloned() + .unwrap_or_else(|| ComponentSpaceUsage::Basic(ByteCount::default())) } /// Num docs in segment @@ -136,37 +149,50 @@ impl SegmentSpaceUsage { /// Space usage for term dictionary pub fn termdict(&self) -> &PerFieldSpaceUsage { - &self.termdict + self.per_field(TERMDICT) } /// Space usage for postings list pub fn postings(&self) -> &PerFieldSpaceUsage { - &self.postings + self.per_field(POSTINGS) } /// Space usage for positions pub fn positions(&self) -> &PerFieldSpaceUsage { - &self.positions + self.per_field(POSITIONS) } /// Space usage for fast fields pub fn fast_fields(&self) -> &PerFieldSpaceUsage { - &self.fast_fields + self.per_field(FAST_FIELDS) } /// Space usage for field norms pub fn fieldnorms(&self) -> &PerFieldSpaceUsage { - &self.fieldnorms + self.per_field(FIELDNORMS) + } + + fn per_field(&self, key: &str) -> &PerFieldSpaceUsage { + match self.components.get(key) { + Some(ComponentSpaceUsage::PerField(usage)) => usage, + _ => &EMPTY_PER_FIELD, + } } /// Space usage for stored documents pub fn store(&self) -> &StoreSpaceUsage { - &self.store + match self.components.get(STORE) { + Some(ComponentSpaceUsage::Store(usage)) => usage, + _ => &EMPTY_STORE, + } } /// Space usage for document deletions pub fn deletes(&self) -> ByteCount { - self.deletes + match self.components.get(DELETES) { + Some(ComponentSpaceUsage::Basic(bytes)) => *bytes, + _ => ByteCount::default(), + } } /// Total space usage in bytes for this segment. diff --git a/src/store/mod.rs b/src/store/mod.rs index cccf4d8f9b..065eec9258 100644 --- a/src/store/mod.rs +++ b/src/store/mod.rs @@ -33,11 +33,13 @@ mod compressors; mod decompressors; mod footer; mod index; +pub mod plugin; mod reader; mod writer; pub use self::compressors::{Compressor, ZstdCompressor}; pub use self::decompressors::Decompressor; +pub use self::plugin::{StorePlugin, StorePluginWriter}; pub use self::reader::{CacheStats, StoreReader}; pub(crate) use self::reader::{DocStoreVersion, DOCSTORE_CACHE_CAPACITY}; pub use self::writer::StoreWriter; diff --git a/src/store/plugin.rs b/src/store/plugin.rs new file mode 100644 index 0000000000..493bfcc964 --- /dev/null +++ b/src/store/plugin.rs @@ -0,0 +1,237 @@ +//! Store as a [`SegmentPlugin`] implementation. +//! +//! This wraps the existing `StoreWriter` and `StoreReader` types behind the plugin +//! interface so that the document store participates in the unified plugin lifecycle. + +use std::any::Any; +use std::collections::BTreeMap; + +use measure_time::debug_time; + +use crate::directory::Directory; +use crate::index::{SegmentComponent, SegmentReader}; +use crate::indexer::doc_id_mapping::DocIdMapping; +use crate::plugin::{PluginMergeContext, PluginWriter, PluginWriterContext, SegmentPlugin}; +use crate::schema::document::Document; +use crate::schema::Schema; +use crate::space_usage::{ComponentSpaceUsage, STORE}; +use crate::store::{StoreReader, StoreWriter}; +use crate::Segment; + +pub struct StorePlugin; + +impl SegmentPlugin for StorePlugin { + fn extensions(&self) -> &[&str] { + &["store"] + } + + fn create_writer(&self, _ctx: &PluginWriterContext) -> crate::Result> { + unimplemented!("StorePlugin is a built-in; use StorePluginWriter::new") + } + + fn merge(&self, ctx: PluginMergeContext) -> crate::Result<()> { + debug_time!("write-storable-fields"); + debug!("write-storable-fields"); + + let path = ctx.target_segment.relative_path(SegmentComponent::Store); + let store_write = ctx.target_segment.index().directory().open_write(&path)?; + let settings = ctx.settings; + let mut store_writer = StoreWriter::new( + store_write, + settings.docstore_compression, + settings.docstore_blocksize, + settings.docstore_compress_dedicated_thread, + )?; + + if !ctx.doc_id_mapping.is_trivial() { + debug!("non-trivial-doc-id-mapping"); + + let store_readers: Vec<_> = ctx + .readers + .iter() + .map(|reader| reader.get_store_reader(50)) + .collect::>()?; + + let mut document_iterators: Vec<_> = store_readers + .iter() + .enumerate() + .map(|(i, store)| store.iter_raw(ctx.readers[i].alive_bitset())) + .collect(); + + for old_doc_addr in ctx.doc_id_mapping.iter_old_doc_addrs() { + let doc_bytes_it = &mut document_iterators[old_doc_addr.segment_ord as usize]; + if let Some(doc_bytes_res) = doc_bytes_it.next() { + let doc_bytes = doc_bytes_res?; + store_writer.store_bytes(&doc_bytes)?; + } else { + return Err(crate::error::DataCorruption::comment_only(format!( + "unexpected missing document in docstore on merge, doc address \ + {old_doc_addr:?}", + )) + .into()); + } + } + } else { + debug!("trivial-doc-id-mapping"); + for reader in ctx.readers { + let store_reader = reader.get_store_reader(1)?; + if reader.has_deletes() + // If there is not enough data in the store, we avoid stacking in order to + // avoid creating many small blocks in the doc store. Once we have 5 full blocks, + // we start stacking. In the worst case 2/7 of the blocks would be very small. + || store_reader.block_checkpoints().take(7).count() < 6 + || store_reader.decompressor() != store_writer.compressor().into() + { + for doc_bytes_res in store_reader.iter_raw(reader.alive_bitset()) { + let doc_bytes = doc_bytes_res?; + store_writer.store_bytes(&doc_bytes)?; + } + } else { + store_writer.stack(store_reader)?; + } + } + } + store_writer.close()?; + Ok(()) + } + + fn space_usage( + &self, + segment_reader: &SegmentReader, + ) -> crate::Result> { + let store = segment_reader.get_store_reader(0)?; + Ok(BTreeMap::from([( + STORE.to_string(), + ComponentSpaceUsage::Store(store.space_usage()), + )])) + } +} + +pub struct StorePluginWriter { + store_writer: Option, + remapping_required: bool, +} + +impl StorePluginWriter { + pub(crate) fn new(ctx: &PluginWriterContext) -> crate::Result { + let settings = ctx.segment.index().settings(); + let directory = ctx.segment.index().directory(); + let remapping_required = settings.sort_by_field.is_some(); + + let store_writer = if remapping_required { + let path = ctx.segment.relative_path(SegmentComponent::TempStore); + let store_write = directory.open_write(&path)?; + StoreWriter::new( + store_write, + 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. + 16000, + settings.docstore_compress_dedicated_thread, + )? + } else { + let path = ctx.segment.relative_path(SegmentComponent::Store); + let store_write = directory.open_write(&path)?; + StoreWriter::new( + store_write, + settings.docstore_compression, + settings.docstore_blocksize, + settings.docstore_compress_dedicated_thread, + )? + }; + + Ok(StorePluginWriter { + store_writer: Some(store_writer), + remapping_required, + }) + } + + pub fn store(&mut self, document: &D, schema: &Schema) -> crate::Result<()> { + if let Some(ref mut writer) = self.store_writer { + writer + .store(document, schema) + .map_err(|e| crate::TantivyError::InternalError(e.to_string()))?; + } + Ok(()) + } + + pub fn store_bytes(&mut self, serialized_document: &[u8]) -> crate::Result<()> { + if let Some(ref mut writer) = self.store_writer { + writer + .store_bytes(serialized_document) + .map_err(|e| crate::TantivyError::InternalError(e.to_string()))?; + } + Ok(()) + } +} + +impl PluginWriter for StorePluginWriter { + fn serialize( + mut self: Box, + segment: &Segment, + doc_id_map: Option<&DocIdMapping>, + ) -> crate::Result<()> { + // For sorted indexes (remapping_required), we need to close the temp store + // and reorder documents into the final store. + if self.remapping_required { + if let Some(doc_id_map) = doc_id_map { + debug!("resort-docstore"); + // Close the temp store writer first + if let Some(old_writer) = self.store_writer.take() { + old_writer + .close() + .map_err(|e| crate::TantivyError::InternalError(e.to_string()))?; + } + + // Open the final store writer + let settings = segment.index().settings().clone(); + let store_write = segment.open_write(SegmentComponent::Store)?; + let mut store_writer = StoreWriter::new( + store_write, + settings.docstore_compression, + settings.docstore_blocksize, + settings.docstore_compress_dedicated_thread, + ) + .map_err(|e| crate::TantivyError::InternalError(e.to_string()))?; + + // Read from temp store and write in new order + let store_read = StoreReader::open( + segment.open_read(SegmentComponent::TempStore)?, + 1, /* The docstore is configured to have one doc per block, and each doc is + * accessed only once: we don't need caching. */ + )?; + for old_doc_id in doc_id_map.iter_old_doc_ids() { + let doc_bytes = store_read.get_document_bytes(old_doc_id)?; + store_writer + .store_bytes(&doc_bytes) + .map_err(|e| crate::TantivyError::InternalError(e.to_string()))?; + } + + // Replace the writer with the final store writer + self.store_writer = Some(store_writer); + } + } + // For non-sorted indexes, data was already written incrementally. + if let Some(writer) = self.store_writer { + writer + .close() + .map_err(|e| crate::TantivyError::InternalError(e.to_string()))?; + } + Ok(()) + } + + fn mem_usage(&self) -> usize { + self.store_writer + .as_ref() + .map_or(0, |writer| writer.mem_usage()) + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } +}