Skip to main content

kevy_text/
segment_stats.rs

1//! Read-only counters and accessors for [`TextSegment`], split from
2//! `segment.rs` for the 500-LOC house rule. A child module (declared in
3//! `segment.rs`), so it reaches the segment's private fields.
4//!
5//! `stats` is the measured side of the documented memory formula
6//! (`bench/textgate.sh` clamps this against real RSS growth). The three
7//! accessors below are the pass-1 API a cross-shard query sums for
8//! global BM25 (step 4b).
9
10use super::{TextSegment, TextStats};
11use crate::buckets::Buckets;
12use crate::docvalues::DocValues;
13use crate::fields::FieldStats;
14use crate::positions::Positions;
15
16impl TextSegment {
17    /// Live counters — O(1): every term is a running
18    /// counter maintained at the mutation sites (the per-tick stat
19    /// walk of these structures was a consumer's measured tiering
20    /// idle/write-load CPU term, F16a). [`Self::recompute_stats`] is
21    /// the walking reference the tests hold these to.
22    pub fn stats(&self) -> TextStats {
23        TextStats {
24            docs: self.docs.len() as u64,
25            tokens: self.postings.len() as u64,
26            postings: self.postings_total,
27            approx_bytes: self.token_bytes
28                + self.many_slots * 30
29                + self.doc_bytes
30                + self.positions.as_ref().map_or(0, Positions::approx_bytes)
31                + self.fields.as_ref().map_or(0, FieldStats::approx_bytes)
32                + self.values.as_ref().map_or(0, DocValues::approx_bytes),
33        }
34    }
35
36    /// The walking reference — recomputes every counter from the live
37    /// structures. Test-only: production reads the running counters.
38    #[cfg(test)]
39    pub(crate) fn recompute_stats(&self) -> TextStats {
40        let postings: u64 = self.postings.values().map(|l| l.len() as u64).sum();
41        TextStats {
42            docs: self.docs.len() as u64,
43            tokens: self.postings.len() as u64,
44            postings,
45            approx_bytes: self.recompute_approx_bytes(),
46        }
47    }
48
49    /// The measured heap estimate: token keys + per-`Many`-posting
50    /// structure + the docs/id tables. Hapax (`One`) lists are inline,
51    /// so only `Many` lists pay the band-vec + index cost.
52    #[cfg(test)]
53    fn recompute_approx_bytes(&self) -> u64 {
54        let many_postings: u64 = self
55            .postings
56            .values()
57            .map(|l| match l {
58                Buckets::One { .. } => 0,
59                Buckets::Many(m) => m.index.len() as u64,
60            })
61            .sum();
62        let token_bytes: u64 = self.postings.keys().map(|t| (t.len() + 48) as u64).sum();
63        // docs table + the id→key / id→dl tables (key stored twice);
64        // docs keep each field's text so an update re-derives tokens.
65        let doc_bytes: u64 = self
66            .docs
67            .iter()
68            .map(|(k, (_, _, fields))| {
69                let text: usize = fields.iter().map(|(t, _)| t.len() + 4).sum();
70                (2 * k.len() + text + 110) as u64
71            })
72            .sum();
73        // per-Many-posting ≈ 4B band-vec slot + ~26B list-index entry.
74        // The side-channels add their own terms when present (`WITH
75        // POSITIONS`, the per-field breakdown of a multi-field index,
76        // and the declared stored values); absent, each contributes
77        // nothing, so a plain single-field segment's formula is
78        // byte-identical.
79        let position_bytes = self.positions.as_ref().map_or(0, Positions::recompute_bytes);
80        let field_bytes = self.fields.as_ref().map_or(0, FieldStats::recompute_bytes);
81        let value_bytes = self.values.as_ref().map_or(0, DocValues::recompute_bytes);
82        token_bytes + many_postings * 30 + doc_bytes + position_bytes + field_bytes + value_bytes
83    }
84
85    /// Verify hook: is `key` indexed here?
86    pub fn contains(&self, key: &[u8]) -> bool {
87        self.docs.contains_key(key)
88    }
89
90    /// This shard's live document count — the Σ n_docs half of the same
91    /// global-BM25 sum [`Self::total_len`] feeds.
92    ///
93    /// Its own accessor because the number is a `len()` and [`Self::stats`]
94    /// is not: `stats` also computes `approx_bytes`, which walks every
95    /// token, every posting and every positional blob. Pass 1 of a
96    /// cross-shard query read `stats().docs` and dropped the rest, so a
97    /// phrase query over a million documents spent 82% of its CPU
98    /// re-measuring the index's memory footprint once per shard per query.
99    pub fn docs(&self) -> u64 {
100        self.docs.len() as u64
101    }
102
103    /// This shard's total document length in tokens (unweighted) — one
104    /// of the three numbers a cross-shard query sums for global BM25
105    /// (step 4b, pass 1): global avgdl = Σ total_len / Σ n_docs.
106    pub fn total_len(&self) -> u64 {
107        self.total_len
108    }
109
110    /// This shard's document frequency for `token` — how many local
111    /// documents contain it. Summed across shards for one query token's
112    /// global df. `0` when the token is absent here.
113    pub fn local_df(&self, token: &[u8]) -> u32 {
114        self.postings.get(token).map_or(0, Buckets::len) as u32
115    }
116}