added payload per term

src/codec/postings/mod.rs

@@ -93,6 +93,19 @@ pub trait PostingsSerializer {
     /// blockwand is disabled), the term_doc_freq passed will be set 0.
     fn new_term(&mut self, term_doc_freq: u32, record_term_freq: bool);
 
+    /// Codec-specific per-term payload.
+    ///
+    /// It is supplied right after `new_term` and before any `write_doc`, so the
+    /// codec can let it influence how the postings list is encoded.
+    ///
+    /// Hidden contract: `new_term` MUST reset any per-term payload state to its
+    /// default. This method is only called for terms that actually have a
+    /// payload registered, so a codec cannot rely on it being called for every
+    /// term.
+    ///
+    /// The default implementation ignores the payload.
+    fn set_term_payload(&mut self, _payload: &dyn std::any::Any) {}
+
     /// Records a new document id for the current term.
     /// The serializer may ignore it.
     fn write_doc(&mut self, doc_id: DocId, term_freq: u32);

src/docset.rs

@@ -1,8 +1,6 @@
 use std::ops::{Deref as _, DerefMut as _};
 
-use common::BitSet;
-
-use common::TinySet;
+use common::{BitSet, TinySet};
 
 use crate::fastfield::AliveBitSet;
 use crate::DocId;

src/indexer/segment_writer.rs

@@ -1,3 +1,5 @@
+use std::any::Any;
+
 use columnar::MonotonicallyMappableToU64;
 use common::JsonPathWriter;
 use itertools::Itertools;
@@ -16,7 +18,7 @@ use crate::postings::{
     PerFieldPostingsWriter, PostingsWriter, PostingsWriterEnum,
 };
 use crate::schema::document::{Document, Value};
-use crate::schema::{FieldEntry, FieldType, Schema, DATE_TIME_PRECISION_INDEXED};
+use crate::schema::{Field, FieldEntry, FieldType, Schema, DATE_TIME_PRECISION_INDEXED};
 use crate::tokenizer::{FacetTokenizer, PreTokenizedStream, TextAnalyzer, Tokenizer};
 use crate::{DocId, Opstamp, TantivyError};
 
@@ -49,7 +51,7 @@ fn compute_initial_table_size(per_thread_memory_budget: usize) -> crate::Result<
 pub struct SegmentWriter<Codec: crate::codec::Codec> {
     pub(crate) max_doc: DocId,
     pub(crate) ctx: IndexingContext,
-    pub(crate) per_field_postings_writers: PerFieldPostingsWriter,
+    pub per_field_postings_writers: PerFieldPostingsWriter,
     pub(crate) segment_serializer: SegmentSerializer<Codec>,
     pub(crate) fast_field_writers: FastFieldsWriter,
     pub(crate) fieldnorms_writer: FieldNormsWriter,
@@ -148,6 +150,67 @@ impl<Codec: crate::codec::Codec> SegmentWriter<Codec> {
             + self.segment_serializer.mem_usage()
     }
 
+    /// Attaches or updates a codec-specific payload on a term of a regular
+    /// (non-JSON) field.
+    ///
+    /// `value_bytes` is the serialized term value, i.e. exactly what would be
+    /// appended after the field id (the raw text bytes for a str field, or the
+    /// big-endian bytes for a numeric field).
+    ///
+    /// If the term does not exist yet, it is inserted with an empty recorder so
+    /// that it still gets serialized even though it belongs to no document.
+    /// `updater` receives the previously registered payload (`None` if absent)
+    /// and returns the payload to store. The payload is handed to the codec's
+    /// postings serializer (via `set_term_payload`) at the beginning of the
+    /// term during serialization.
+    pub(crate) fn update_term_payload(
+        &mut self,
+        field: Field,
+        value_bytes: &[u8],
+        updater: impl FnOnce(Option<Box<dyn Any + Send>>) -> Box<dyn Any + Send>,
+    ) {
+        let mut term = IndexingTerm::with_capacity(value_bytes.len());
+        term.set_field(field);
+        term.append_bytes(value_bytes);
+        self.update_term_payload_for_serialized_term(field, term.serialized_term(), updater);
+    }
+
+    /// Same as [`Self::update_term_payload`] for a JSON field.
+    ///
+    /// `value_bytes` must be the type-tagged value (`[type code][value]`), the
+    /// representation that follows the path within a JSON term.
+    pub(crate) fn update_json_term_payload(
+        &mut self,
+        field: Field,
+        json_path: &str,
+        value_bytes: &[u8],
+        updater: impl FnOnce(Option<Box<dyn Any + Send>>) -> Box<dyn Any + Send>,
+    ) {
+        let unordered_id = self
+            .ctx
+            .path_to_unordered_id
+            .get_or_allocate_unordered_id(json_path);
+        // JSON term key layout: `[field:4][unordered_path_id:4][type code][value]`.
+        let mut serialized_term = Vec::with_capacity(8 + value_bytes.len());
+        serialized_term.extend_from_slice(&field.field_id().to_be_bytes());
+        serialized_term.extend_from_slice(&unordered_id.to_be_bytes());
+        serialized_term.extend_from_slice(value_bytes);
+        self.update_term_payload_for_serialized_term(field, &serialized_term, updater);
+    }
+
+    fn update_term_payload_for_serialized_term(
+        &mut self,
+        field: Field,
+        serialized_term: &[u8],
+        updater: impl FnOnce(Option<Box<dyn Any + Send>>) -> Box<dyn Any + Send>,
+    ) {
+        let postings_writer = self.per_field_postings_writers.get_for_field(field);
+        let addr = postings_writer.ensure_term(serialized_term, &mut self.ctx);
+        let previous_payload = self.ctx.codec_term_payloads.remove(&addr);
+        let new_payload = updater(previous_payload);
+        self.ctx.codec_term_payloads.insert(addr, new_payload);
+    }
+
     fn index_document<D: Document>(&mut self, doc: &D) -> crate::Result<()> {
         let doc_id = self.max_doc;
 

src/indexer/single_segment_index_writer.rs

@@ -1,3 +1,4 @@
+use std::any::Any;
 use std::marker::PhantomData;
 
 use crate::codec::StandardCodec;
@@ -6,6 +7,7 @@ use crate::indexer::operation::AddOperation;
 use crate::indexer::segment_updater::save_metas;
 use crate::indexer::SegmentWriter;
 use crate::schema::document::Document;
+use crate::schema::Field;
 use crate::{Directory, Index, IndexMeta, Opstamp, Segment, TantivyDocument};
 
 #[doc(hidden)]
@@ -42,6 +44,43 @@ impl<Codec: crate::codec::Codec, D: Document> SingleSegmentIndexWriter<Codec, D>
             .add_document(AddOperation { opstamp, document })
     }
 
+    /// Attaches or updates a codec-specific payload on a term of a regular
+    /// (non-JSON) field.
+    ///
+    /// `value_bytes` is the serialized term value, i.e. exactly what would be
+    /// appended after the field id (the raw text bytes for a str field, or the
+    /// big-endian bytes for a numeric field).
+    ///
+    /// The term does not need to belong to any document: if it does not exist
+    /// yet, it is created with an empty recorder so it still gets serialized.
+    /// `updater` receives the previously registered payload (`None` if absent)
+    /// and returns the payload to store. The payload is handed to the codec at
+    /// the beginning of the term during serialization.
+    pub fn update_term_payload(
+        &mut self,
+        field: Field,
+        value_bytes: &[u8],
+        updater: impl FnOnce(Option<Box<dyn Any + Send>>) -> Box<dyn Any + Send>,
+    ) {
+        self.segment_writer
+            .update_term_payload(field, value_bytes, updater);
+    }
+
+    /// Same as [`Self::update_term_payload`] for a JSON field.
+    ///
+    /// `value_bytes` must be the type-tagged value (`[type code][value]`), the
+    /// representation that follows the path within a JSON term.
+    pub fn update_json_term_payload(
+        &mut self,
+        field: Field,
+        json_path: &str,
+        value_bytes: &[u8],
+        updater: impl FnOnce(Option<Box<dyn Any + Send>>) -> Box<dyn Any + Send>,
+    ) {
+        self.segment_writer
+            .update_json_term_payload(field, json_path, value_bytes, updater);
+    }
+
     pub fn finalize(self) -> crate::Result<Index<Codec>> {
         let max_doc = self.segment_writer.max_doc();
         self.segment_writer.finalize()?;
@@ -60,3 +99,231 @@ impl<Codec: crate::codec::Codec, D: Document> SingleSegmentIndexWriter<Codec, D>
         Ok(segment.index().clone())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::any::Any;
+    use std::cell::RefCell;
+    use std::io;
+
+    use super::SingleSegmentIndexWriter;
+    use crate::codec::postings::{PostingsCodec, PostingsSerializer};
+    use crate::codec::standard::postings::{
+        SegmentPostings, StandardPostingsCodec, StandardPostingsSerializer,
+    };
+    use crate::codec::Codec;
+    use crate::fieldnorm::FieldNormReader;
+    use crate::schema::{IndexRecordOption, Schema, Type, STRING};
+    use crate::{DocId, Score, Term};
+
+    // The codec is round-tripped through `from_json_props` when the index is
+    // opened, so it cannot carry the capture sink itself. We use a thread-local
+    // sink instead: the `SingleSegmentIndexWriter` is single-threaded, so
+    // serialization runs on the test thread, and each test owns its own
+    // thread-local (clear it at the start of the test).
+    thread_local! {
+        static CAPTURED_PAYLOADS: RefCell<Vec<u64>> = const { RefCell::new(Vec::new()) };
+    }
+
+    fn reset_captured() {
+        CAPTURED_PAYLOADS.with(|captured| captured.borrow_mut().clear());
+    }
+
+    fn captured_payloads() -> Vec<u64> {
+        CAPTURED_PAYLOADS.with(|captured| captured.borrow().clone())
+    }
+
+    /// A postings serializer that delegates to the standard one, but records
+    /// the `u64` payload value of every term that carries a codec payload.
+    struct CapturingPostingsSerializer {
+        inner: StandardPostingsSerializer,
+    }
+
+    impl PostingsSerializer for CapturingPostingsSerializer {
+        fn new_term(&mut self, term_doc_freq: u32, record_term_freq: bool) {
+            self.inner.new_term(term_doc_freq, record_term_freq);
+        }
+
+        fn set_term_payload(&mut self, payload: &dyn Any) {
+            let value = *payload
+                .downcast_ref::<u64>()
+                .expect("payload should be a u64");
+            CAPTURED_PAYLOADS.with(|captured| captured.borrow_mut().push(value));
+        }
+
+        fn write_doc(&mut self, doc_id: DocId, term_freq: u32) {
+            self.inner.write_doc(doc_id, term_freq);
+        }
+
+        fn close_term(&mut self, doc_freq: u32, wrt: &mut impl io::Write) -> io::Result<()> {
+            self.inner.close_term(doc_freq, wrt)
+        }
+    }
+
+    #[derive(Clone, Debug)]
+    struct CapturingPostingsCodec;
+
+    impl PostingsCodec for CapturingPostingsCodec {
+        type PostingsSerializer = CapturingPostingsSerializer;
+        type Postings = SegmentPostings;
+
+        fn new_serializer(
+            &self,
+            avg_fieldnorm: Score,
+            mode: IndexRecordOption,
+            fieldnorm_reader: Option<FieldNormReader>,
+        ) -> Self::PostingsSerializer {
+            CapturingPostingsSerializer {
+                inner: StandardPostingsCodec.new_serializer(avg_fieldnorm, mode, fieldnorm_reader),
+            }
+        }
+
+        fn load_postings(
+            &self,
+            doc_freq: u32,
+            postings_data: common::OwnedBytes,
+            record_option: IndexRecordOption,
+            requested_option: IndexRecordOption,
+            positions_data: Option<common::OwnedBytes>,
+        ) -> io::Result<Self::Postings> {
+            StandardPostingsCodec.load_postings(
+                doc_freq,
+                postings_data,
+                record_option,
+                requested_option,
+                positions_data,
+            )
+        }
+    }
+
+    #[derive(Clone, Debug, Default)]
+    struct CapturingCodec;
+
+    impl Codec for CapturingCodec {
+        type PostingsCodec = CapturingPostingsCodec;
+
+        const ID: &'static str = "test-capturing-codec";
+
+        fn from_json_props(_json_value: &serde_json::Value) -> crate::Result<Self> {
+            Ok(CapturingCodec)
+        }
+
+        fn to_json_props(&self) -> serde_json::Value {
+            serde_json::Value::Null
+        }
+
+        fn postings_codec(&self) -> &Self::PostingsCodec {
+            &CapturingPostingsCodec
+        }
+    }
+
+    fn build_writer(schema: Schema) -> SingleSegmentIndexWriter<CapturingCodec> {
+        let index = crate::IndexBuilder::default()
+            .codec(CapturingCodec)
+            .schema(schema)
+            .create_in_ram()
+            .unwrap();
+        SingleSegmentIndexWriter::new(index, 15_000_000).unwrap()
+    }
+
+    #[test]
+    fn test_update_term_payload_regular_field() {
+        reset_captured();
+        let mut schema_builder = Schema::builder();
+        let text = schema_builder.add_text_field("text", STRING);
+        let schema = schema_builder.build();
+        let mut writer = build_writer(schema);
+
+        writer.add_document(crate::doc!(text => "alpha")).unwrap();
+        writer.add_document(crate::doc!(text => "beta")).unwrap();
+        writer.add_document(crate::doc!(text => "gamma")).unwrap();
+
+        // Existing term that belongs to a document.
+        writer.update_term_payload(text, b"beta", |previous_payload| {
+            assert!(previous_payload.is_none());
+            Box::new(100u64)
+        });
+        // Updating the same term: the previous payload is handed back.
+        writer.update_term_payload(text, b"beta", |previous_payload| {
+            let previous = previous_payload.expect("expected the previous payload");
+            assert_eq!(*previous.downcast::<u64>().unwrap(), 100u64);
+            Box::new(101u64)
+        });
+        // Brand-new term that belongs to no document: an empty recorder is
+        // created so it still lands in the term dictionary.
+        writer.update_term_payload(text, b"zeta", |previous_payload| {
+            assert!(previous_payload.is_none());
+            Box::new(200u64)
+        });
+
+        let index = writer.finalize().unwrap();
+
+        // Terms are serialized in sorted order: alpha, beta, gamma, zeta.
+        // Only beta and zeta carry a payload.
+        assert_eq!(captured_payloads(), vec![101u64, 200u64]);
+
+        let searcher = index.reader().unwrap().searcher();
+        let segment_reader = searcher.segment_reader(0);
+        let inverted_index = segment_reader.inverted_index(text).unwrap();
+
+        let beta_info = inverted_index
+            .get_term_info(&Term::from_field_text(text, "beta"))
+            .unwrap()
+            .expect("beta should be in the dictionary");
+        assert_eq!(beta_info.doc_freq, 1);
+
+        let zeta_info = inverted_index
+            .get_term_info(&Term::from_field_text(text, "zeta"))
+            .unwrap()
+            .expect("zeta (no document) should still be in the dictionary");
+        assert_eq!(zeta_info.doc_freq, 0);
+    }
+
+    #[test]
+    fn test_update_json_term_payload() {
+        reset_captured();
+        let mut schema_builder = Schema::builder();
+        let json_field = schema_builder.add_json_field("json", STRING);
+        let schema = schema_builder.build();
+        let mut writer = build_writer(schema);
+
+        writer
+            .add_document(crate::doc!(json_field => serde_json::json!({"name": "hello"})))
+            .unwrap();
+
+        let str_value = |value: &str| {
+            let mut bytes = vec![Type::Str.to_code()];
+            bytes.extend_from_slice(value.as_bytes());
+            bytes
+        };
+
+        // Existing str JSON term (path "name", value "hello").
+        writer.update_json_term_payload(json_field, "name", &str_value("hello"), |previous| {
+            assert!(previous.is_none());
+            Box::new(1u64)
+        });
+        // Brand-new str JSON term with no document.
+        writer.update_json_term_payload(json_field, "name", &str_value("world"), |previous| {
+            assert!(previous.is_none());
+            Box::new(2u64)
+        });
+        // Brand-new non-str (numeric) JSON term with no document: exercises the
+        // DocIdRecorder branch of `ensure_term`.
+        let numeric_value = {
+            let mut bytes = vec![Type::I64.to_code()];
+            bytes.extend_from_slice(&[0u8; 8]);
+            bytes
+        };
+        writer.update_json_term_payload(json_field, "count", &numeric_value, |previous| {
+            assert!(previous.is_none());
+            Box::new(3u64)
+        });
+
+        // Should not panic and should serialize cleanly.
+        let _index = writer.finalize().unwrap();
+
+        let mut got = captured_payloads();
+        got.sort_unstable();
+        assert_eq!(got, vec![1u64, 2u64, 3u64]);
+    }
+}

src/lib.rs

@@ -221,7 +221,7 @@ pub mod snippet;
 use std::fmt;
 
 pub use census::{Inventory, TrackedObject};
-pub use common::{f64_to_u64, i64_to_u64, u64_to_f64, u64_to_i64, HasLen};
+pub use common::{self, f64_to_u64, i64_to_u64, u64_to_f64, u64_to_i64, HasLen};
 use once_cell::sync::Lazy;
 use serde::{Deserialize, Serialize};
 

src/postings/indexing_context.rs

@@ -1,4 +1,7 @@
-use stacker::{ArenaHashMap, MemoryArena};
+use std::any::Any;
+
+use fnv::FnvHashMap;
+use stacker::{Addr, ArenaHashMap, MemoryArena};
 
 use crate::indexer::path_to_unordered_id::PathToUnorderedId;
 
@@ -11,6 +14,15 @@ pub(crate) struct IndexingContext {
     /// Arena is a memory arena that stores posting lists / term frequencies / positions.
     pub arena: MemoryArena,
     pub path_to_unordered_id: PathToUnorderedId,
+    /// Optional codec-specific payload attached to a term, keyed by the value
+    /// `Addr` of the term's recorder in `term_index`.
+    ///
+    /// Hidden contract: keying on `Addr` is sound because a term's recorder
+    /// address never changes once allocated (the arena only appends, and
+    /// `subscribe` updates the recorder in place). The payload is therefore
+    /// looked up by `Addr` at serialization time and fed to the codec's
+    /// postings serializer at the beginning of the term.
+    pub codec_term_payloads: FnvHashMap<Addr, Box<dyn Any + Send>>,
 }
 
 impl IndexingContext {
@@ -21,6 +33,7 @@ impl IndexingContext {
             arena: MemoryArena::default(),
             term_index,
             path_to_unordered_id: PathToUnorderedId::default(),
+            codec_term_payloads: FnvHashMap::default(),
         }
     }
 

src/postings/json_postings_writer.rs

@@ -96,6 +96,20 @@ impl<Rec: Recorder> PostingsWriter for JsonPostingsWriter<Rec> {
         Ok(())
     }
 
+    fn ensure_term(&self, serialized_term: &[u8], ctx: &mut IndexingContext) -> Addr {
+        // JSON term key layout: `[field:4][unordered_path_id:4][type code][value]`.
+        // Str values are recorded with `Rec`, all other types with `DocIdRecorder`
+        // (mirroring the dispatch in `serialize`).
+        let typ = Type::from_code(serialized_term[8]).expect("Invalid type code in JSON term");
+        if typ == Type::Str {
+            ctx.term_index
+                .get_or_create_value_addr::<Rec>(serialized_term, Rec::default)
+        } else {
+            ctx.term_index
+                .get_or_create_value_addr::<DocIdRecorder>(serialized_term, DocIdRecorder::default)
+        }
+    }
+
     fn total_num_tokens(&self) -> u64 {
         self.str_posting_writer.total_num_tokens() + self.non_str_posting_writer.total_num_tokens()
     }

src/postings/mod.rs

@@ -21,7 +21,8 @@ pub use postings::DocFreq;
 pub(crate) use stacker::compute_table_memory_size;
 
 pub(crate) use self::indexing_context::IndexingContext;
-pub(crate) use self::per_field_postings_writer::PerFieldPostingsWriter;
+#[doc(hidden)]
+pub use self::per_field_postings_writer::PerFieldPostingsWriter;
 pub use self::postings::Postings;
 pub(crate) use self::postings_writer::{
     serialize_postings, IndexingPosition, PostingsWriter, PostingsWriterEnum,

src/postings/per_field_postings_writer.rs

@@ -3,7 +3,7 @@ use crate::postings::postings_writer::{PostingsWriterEnum, SpecializedPostingsWr
 use crate::postings::recorder::{DocIdRecorder, TermFrequencyRecorder, TfAndPositionRecorder};
 use crate::schema::{Field, FieldEntry, FieldType, IndexRecordOption, Schema};
 
-pub(crate) struct PerFieldPostingsWriter {
+pub struct PerFieldPostingsWriter {
     per_field_postings_writers: Vec<PostingsWriterEnum>,
 }
 
@@ -22,7 +22,7 @@ impl PerFieldPostingsWriter {
         &self.per_field_postings_writers[field.field_id() as usize]
     }
 
-    pub(crate) fn get_for_field_mut(&mut self, field: Field) -> &mut PostingsWriterEnum {
+    pub fn get_for_field_mut(&mut self, field: Field) -> &mut PostingsWriterEnum {
         &mut self.per_field_postings_writers[field.field_id() as usize]
     }
 }

src/postings/postings_writer.rs

@@ -196,6 +196,21 @@ impl PostingsWriter for PostingsWriterEnum {
         }
     }
 
+    fn ensure_term(&self, serialized_term: &[u8], ctx: &mut IndexingContext) -> Addr {
+        match self {
+            PostingsWriterEnum::DocId(writer) => writer.ensure_term(serialized_term, ctx),
+            PostingsWriterEnum::DocIdTf(writer) => writer.ensure_term(serialized_term, ctx),
+            PostingsWriterEnum::DocTfAndPosition(writer) => {
+                writer.ensure_term(serialized_term, ctx)
+            }
+            PostingsWriterEnum::JsonDocId(writer) => writer.ensure_term(serialized_term, ctx),
+            PostingsWriterEnum::JsonDocIdTf(writer) => writer.ensure_term(serialized_term, ctx),
+            PostingsWriterEnum::JsonDocTfAndPosition(writer) => {
+                writer.ensure_term(serialized_term, ctx)
+            }
+        }
+    }
+
     /// Tokenize a text and subscribe all of its token.
     fn index_text(
         &mut self,
@@ -263,6 +278,15 @@ pub(crate) trait PostingsWriter: Send + Sync {
         serializer: &mut FieldSerializer<C>,
     ) -> io::Result<()>;
 
+    /// Ensures `serialized_term` has an entry in the term index, creating an
+    /// empty recorder (matching this writer's indexing option) if the term is
+    /// not present yet, and returns the value `Addr` of its recorder.
+    ///
+    /// An existing recorder is never overwritten, so the term keeps any
+    /// posting data already recorded for it. This is used to attach a
+    /// codec-specific payload to a term that may belong to no document.
+    fn ensure_term(&self, serialized_term: &[u8], ctx: &mut IndexingContext) -> Addr;
+
     /// Tokenize a text and subscribe all of its token.
     fn index_text(
         &mut self,
@@ -322,6 +346,10 @@ impl<Rec: Recorder> SpecializedPostingsWriter<Rec> {
         let recorder: Rec = ctx.term_index.read(addr);
         let term_doc_freq = recorder.term_doc_freq().unwrap_or(0u32);
         serializer.new_term(term, term_doc_freq, recorder.has_term_freq())?;
+        if let Some(payload) = ctx.codec_term_payloads.get(&addr) {
+            // `&(dyn Any + Send)` upcasts to `&dyn Any`.
+            serializer.set_term_payload(payload.as_ref());
+        }
         recorder.serialize(&ctx.arena, serializer, buffer_lender);
         serializer.close_term()?;
         Ok(())
@@ -372,6 +400,11 @@ impl<Rec: Recorder> PostingsWriter for SpecializedPostingsWriter<Rec> {
         Ok(())
     }
 
+    fn ensure_term(&self, serialized_term: &[u8], ctx: &mut IndexingContext) -> Addr {
+        ctx.term_index
+            .get_or_create_value_addr::<Rec>(serialized_term, Rec::default)
+    }
+
     fn total_num_tokens(&self) -> u64 {
         self.total_num_tokens
     }

src/postings/serializer.rs

@@ -203,6 +203,13 @@ impl<'a, C: Codec> FieldSerializer<'a, C> {
         self.new_term(term, 0, false)
     }
 
+    /// Forwards a codec-specific per-term payload to the postings serializer.
+    ///
+    /// Must be called after `new_term` and before any `write_doc`.
+    pub fn set_term_payload(&mut self, payload: &dyn std::any::Any) {
+        self.postings_serializer.set_term_payload(payload);
+    }
+
     /// Serialize the information that a document contains for the current term:
     /// its term frequency, and the position deltas.
     ///

src/query/boolean_query/boolean_weight.rs

@@ -279,7 +279,6 @@ impl<TScoreCombiner: ScoreCombiner> BooleanWeight<TScoreCombiner> {
             Some(exclude_scorers_union)
         };
 
-
         let include_scorer = match (should_scorers, must_scorers) {
             (ShouldScorersCombinationMethod::Ignored, must_scorers) => {
                 // No SHOULD clauses (or they were absorbed into MUST).

src/query/boolean_query/mod.rs

@@ -1,7 +1,6 @@
 mod boolean_query;
 mod boolean_weight;
 
-
 pub use self::boolean_query::BooleanQuery;
 pub use self::boolean_weight::BooleanWeight;
 

stacker/src/arena_hashmap.rs

@@ -83,6 +83,28 @@ impl ArenaHashMap {
         self.shared_arena_hashmap
             .mutate_or_create(key, &mut self.memory_arena, updater);
     }
+
+    /// Returns the address of the value associated to `key`, creating an entry
+    /// with `make_default()` if the key is not present yet (an existing value
+    /// is left untouched).
+    ///
+    /// See [`SharedArenaHashMap::get_or_create_value_addr`] for the address
+    /// stability guarantees.
+    #[inline]
+    pub fn get_or_create_value_addr<V>(
+        &mut self,
+        key: &[u8],
+        make_default: impl FnOnce() -> V,
+    ) -> Addr
+    where
+        V: Copy + 'static,
+    {
+        self.shared_arena_hashmap.get_or_create_value_addr(
+            key,
+            &mut self.memory_arena,
+            make_default,
+        )
+    }
 }
 
 #[cfg(test)]
@@ -121,6 +143,28 @@ mod tests {
         assert_eq!(hash_map.get::<u32>(b"abc"), None);
     }
 
+    #[test]
+    fn test_get_or_create_value_addr() {
+        let mut hash_map: ArenaHashMap = ArenaHashMap::default();
+        // Creates the entry with the default value.
+        let addr_abc = hash_map.get_or_create_value_addr(b"abc", || 7u32);
+        assert_eq!(hash_map.read::<u32>(addr_abc), 7u32);
+        // Returns the same address and does NOT overwrite an existing value.
+        let addr_abc_again = hash_map.get_or_create_value_addr(b"abc", || 99u32);
+        assert_eq!(addr_abc_again, addr_abc);
+        assert_eq!(hash_map.read::<u32>(addr_abc), 7u32);
+        // A different key gets its own entry.
+        let addr_def = hash_map.get_or_create_value_addr(b"def", || 5u32);
+        assert_ne!(addr_def, addr_abc);
+        assert_eq!(hash_map.read::<u32>(addr_def), 5u32);
+        // The address matches the one yielded by `iter`.
+        for (key, addr) in hash_map.iter() {
+            if key == b"abc" {
+                assert_eq!(addr, addr_abc);
+            }
+        }
+    }
+
     #[test]
     fn test_many_terms() {
         let mut terms: Vec<String> = (0..20_000).map(|val| val.to_string()).collect();

stacker/src/memory_arena.rs

@@ -36,7 +36,7 @@ const PAGE_SIZE: usize = 1 << NUM_BITS_PAGE_ADDR; // pages are 1 MB large
 /// page of memory.
 ///
 /// The last 20 bits are an address within this page of memory.
-#[derive(Copy, Clone, Debug)]
+#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
 pub struct Addr(u32);
 
 impl Addr {

stacker/src/shared_arena_hashmap.rs

@@ -347,6 +347,67 @@ impl SharedArenaHashMap {
             kv = self.table[bucket];
         }
     }
+
+    /// Returns the address of the value associated to `key`.
+    ///
+    /// If the key is not present yet, a new entry is created with the value
+    /// returned by `make_default()`. If the key is already present, the stored
+    /// value is left untouched and its address is returned.
+    ///
+    /// The returned `Addr` is the value address, i.e. the same address yielded
+    /// by [`Self::iter`] and consumed by [`MemoryArena::read`]. It remains valid
+    /// for the lifetime of the arena: arena allocations only ever append, and a
+    /// table resize relocates buckets, not the arena-backed key/value data.
+    ///
+    /// The key will be truncated to `u16::MAX` bytes.
+    #[inline]
+    pub fn get_or_create_value_addr<V>(
+        &mut self,
+        key: &[u8],
+        memory_arena: &mut MemoryArena,
+        make_default: impl FnOnce() -> V,
+    ) -> Addr
+    where
+        V: Copy + 'static,
+    {
+        if self.is_saturated() {
+            self.resize();
+        }
+        // Limit the key size to u16::MAX
+        let key = &key[..std::cmp::min(key.len(), u16::MAX as usize)];
+        let hash = self.get_hash(key);
+        let mut probe = self.probe(hash);
+        let mut bucket = probe.next_probe();
+        let mut kv: KeyValue = self.table[bucket];
+        loop {
+            if kv.is_empty() {
+                // The key does not exist yet: create it with the default value.
+                let val = make_default();
+                let num_bytes = std::mem::size_of::<u16>() + key.len() + std::mem::size_of::<V>();
+                let key_addr = memory_arena.allocate_space(num_bytes);
+                {
+                    let data = memory_arena.slice_mut(key_addr, num_bytes);
+                    let key_len_bytes: [u8; 2] = (key.len() as u16).to_le_bytes();
+                    data[..2].copy_from_slice(&key_len_bytes);
+                    let stop = 2 + key.len();
+                    fast_short_slice_copy(key, &mut data[2..stop]);
+                    store(&mut data[stop..], val);
+                }
+                self.set_bucket(hash, key_addr, bucket);
+                return key_addr.offset(2 + key.len() as u32);
+            }
+            if kv.hash == hash
+                && let Some(val_addr) =
+                    self.get_value_addr_if_key_match(key, kv.key_value_addr, memory_arena)
+            {
+                // The key already exists: leave its value untouched.
+                return val_addr;
+            }
+            // This allows fetching the next bucket before the loop jmp
+            bucket = probe.next_probe();
+            kv = self.table[bucket];
+        }
+    }
 }
 
 #[cfg(test)]