mentat/
auto_index.rs

1// Copyright 2026 the Mentat authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License"); you may not use
4// this file except in compliance with the License. You may obtain a copy of the
5// License at http://www.apache.org/licenses/LICENSE-2.0
6// Unless required by applicable law or agreed to in writing, software distributed
7// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
8// CONDITIONS OF ANY KIND, either express or implied. See the License for the
9// specific language governing permissions and limitations under the License.
10
11//! Automatic index management for the embedded store.
12//!
13//! `datoms` always has EAVT and AEVT indexes. The schema's `:db/index` and
14//! `:db/unique` add partial AVET indexes (`WHERE index_avet IS NOT 0`), but the
15//! query SQL never states that predicate, so SQLite can't use them: a
16//! value-filtered pattern like `[?u :user/email "x"]` would scan every
17//! `:user/email` datom through AEVT. So every `:db/index` / `:db/unique`
18//! attribute also gets a *schema* value index (below), created and dropped
19//! with the flag by `mentat_db::db::sync_schema_indexes`, in every mode.
20//!
21//! In [`AutoIndex::Adaptive`] mode the store also counts, per attribute, the queries
22//! whose pattern pins that attribute's value (a constant or a bound `:in`
23//! scalar). When an attribute reaches `min_uses` queries within one tuning
24//! period, mentat creates
25//! `idx_auto_avet_<a> ON datoms (a, v, e) WHERE a = <a>`. The planner uses
26//! that without statistics, because the query already says `a = <a>`. There's
27//! no `value_type_tag` column: an attribute has one value type, and the SQL
28//! doesn't pin the tag when the type is known, so a tag column between `a` and
29//! `v` would stop the index serving `v = ?`.
30//! Indexes mentat creates are listed in `mentat_managed_indexes`, with `kind`
31//! `schema` or `adaptive`. An adaptive index is
32//! dropped when a whole period passes without a use *and* its last use is at
33//! least `idle` old; a schema index only when its attribute loses the flag. That gap between "`min_uses` in a period" and "no use for
34//! `idle`" is the hysteresis. Nothing outside that table is ever dropped.
35//! An index whose average value matches over 5% of the attribute's datoms is
36//! dropped straight after it's built: SQLite would probe it and lose to the
37//! AEVT scan (the scale suite's `:issue/state`, 5 values).
38
39use std::collections::{BTreeMap, BTreeSet};
40use std::fmt;
41use std::time::{Duration, SystemTime, UNIX_EPOCH};
42
43use core_traits::Entid;
44use mentat_core::{HasSchema, Keyword, Schema};
45use mentat_db::db;
46use public_traits::errors::Result;
47
48/// How mentat manages secondary indexes on `datoms`.
49#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
50pub enum AutoIndex {
51    /// No adaptive indexes, and [`Store::tune_indexes`] does nothing. Schema
52    /// value indexes still follow `:db/index` / `:db/unique` (they're part of
53    /// the schema, like SQLite's own indexes).
54    ///
55    /// [`Store::tune_indexes`]: crate::Store::tune_indexes
56    Off,
57    /// Only the schema's indexes: a value index per `:db/index` /
58    /// `:db/unique` attribute. [`Store::tune_indexes`] drops any adaptive
59    /// index mentat created earlier. This is the default.
60    ///
61    /// [`Store::tune_indexes`]: crate::Store::tune_indexes
62    #[default]
63    Schema,
64    /// Schema indexes, plus per-attribute value indexes created and dropped
65    /// from the query workload. Tuning runs automatically every
66    /// [`PERIOD`] queries, and when an attribute reaches `min_uses`.
67    Adaptive,
68}
69
70/// One index change made (or, in a dry run, proposed) by `tune_indexes`.
71#[derive(Clone, Debug, PartialEq, Eq)]
72pub enum IndexAction {
73    Create {
74        index: String,
75        attribute: Entid,
76        ident: Option<Keyword>,
77        /// Value-filtered queries on the attribute in this period.
78        uses: u32,
79    },
80    Drop {
81        index: String,
82        attribute: Entid,
83        ident: Option<Keyword>,
84    },
85}
86
87impl fmt::Display for IndexAction {
88    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
89        let who = |a: &Entid, i: &Option<Keyword>| match i {
90            Some(k) => format!("{k}"),
91            None => a.to_string(),
92        };
93        match self {
94            IndexAction::Create {
95                index,
96                attribute,
97                ident,
98                uses,
99            } => write!(
100                f,
101                "create {index} on {} ({uses} value-filtered queries)",
102                who(attribute, ident)
103            ),
104            IndexAction::Drop {
105                index,
106                attribute,
107                ident,
108            } => write!(f, "drop {index} on {} (idle)", who(attribute, ident)),
109        }
110    }
111}
112
113/// Queries between automatic tuning runs in `Adaptive` mode.
114pub const PERIOD: u32 = 1000;
115/// User attributes only: bootstrap attributes are few and cached.
116const USER0: Entid = 0x10000;
117
118#[derive(Debug)]
119pub(crate) struct Advisor {
120    pub(crate) mode: AutoIndex,
121    pub(crate) min_uses: u32,
122    pub(crate) idle: Duration,
123    uses: BTreeMap<Entid, u32>,
124    queries: u32,
125    /// Attributes whose index turned out unselective; not retried by this Store.
126    rejected: BTreeSet<Entid>,
127    /// Attributes with a managed index, as of the last tuning run.
128    indexed: BTreeSet<Entid>,
129}
130
131impl Default for Advisor {
132    fn default() -> Advisor {
133        let mode = match std::env::var("MENTAT_AUTO_INDEX").as_deref() {
134            Ok("off") => AutoIndex::Off,
135            Ok("adaptive") => AutoIndex::Adaptive,
136            _ => AutoIndex::default(),
137        };
138        Advisor {
139            mode,
140            min_uses: 16,
141            idle: Duration::from_secs(24 * 3600),
142            uses: BTreeMap::new(),
143            queries: 0,
144            rejected: BTreeSet::new(),
145            indexed: BTreeSet::new(),
146        }
147    }
148}
149
150fn now() -> i64 {
151    SystemTime::now()
152        .duration_since(UNIX_EPOCH)
153        .map(|d| d.as_secs() as i64)
154        .unwrap_or(0)
155}
156
157/// `(name, attribute, last_used, is_schema)` for every index mentat created.
158fn registry(conn: &rusqlite::Connection) -> Result<Vec<(String, Entid, i64, bool)>> {
159    if !db::has_table(conn, "mentat_managed_indexes")? {
160        return Ok(vec![]);
161    }
162    let mut stmt = conn.prepare("SELECT name, a, last_used, kind FROM mentat_managed_indexes")?;
163    let rows: Vec<(String, Entid, i64, bool)> = stmt
164        .query_map([], |r| {
165            Ok((
166                r.get(0)?,
167                r.get(1)?,
168                r.get(2)?,
169                r.get::<_, String>(3)? == "schema",
170            ))
171        })?
172        .collect::<rusqlite::Result<Vec<_>>>()?;
173    // Only ever act on names mentat generates, whatever the table says.
174    Ok(rows
175        .into_iter()
176        .filter(|(n, a, _, _)| *n == db::value_index_name(*a))
177        .collect())
178}
179
180/// From the index's fresh `sqlite_stat1` row (`nrow rows/a rows/(a,v) 1`): is
181/// the average value rare enough that probing the index beats the AEVT scan?
182// ponytail: fixed 5% cut from the scale suite (a 20%-per-value enum got
183// slower, a unique email 4x faster); per-query costing if it misjudges.
184fn selective(conn: &rusqlite::Connection, index: &str) -> Result<bool> {
185    let stat: Option<String> = conn
186        .query_row(
187            "SELECT stat FROM sqlite_stat1 WHERE idx = ?",
188            [index],
189            |r| r.get(0),
190        )
191        .ok();
192    let n: Vec<u64> = stat
193        .unwrap_or_default()
194        .split_whitespace()
195        .filter_map(|x| x.parse().ok())
196        .collect();
197    Ok(match n[..] {
198        [rows, _, per_value, ..] => per_value * 20 <= rows.max(1),
199        _ => true, // No statistics (e.g. an empty attribute): keep it.
200    })
201}
202
203impl Advisor {
204    /// Count one query's value-filtered attributes. True when a tuning run is
205    /// due: every `PERIOD` queries, or as soon as an attribute without an index
206    /// reaches `min_uses`.
207    pub(crate) fn record(&mut self, attrs: &BTreeSet<Entid>) -> bool {
208        self.queries += 1;
209        let mut due = self.queries >= PERIOD;
210        for a in attrs {
211            let n = self.uses.entry(*a).or_insert(0);
212            *n += 1;
213            due |= *n == self.min_uses && !self.indexed.contains(a) && !self.rejected.contains(a);
214        }
215        due
216    }
217
218    pub(crate) fn tune(
219        &mut self,
220        conn: &rusqlite::Connection,
221        schema: &Schema,
222        dry_run: bool,
223    ) -> Result<Vec<IndexAction>> {
224        if self.mode == AutoIndex::Off {
225            return Ok(vec![]);
226        }
227        let tx = conn.unchecked_transaction()?;
228        if !dry_run {
229            // Existing stores, and anything a transaction missed.
230            db::sync_schema_indexes(&tx, schema)?;
231            db::refresh_value_index_stats(&tx)?;
232        }
233        let managed = registry(&tx)?;
234        let ident = |a: Entid| schema.get_ident(a).cloned();
235        let now = now();
236        let mut actions = vec![];
237        let have: BTreeSet<Entid> = managed.iter().map(|(_, a, _, _)| *a).collect();
238        if self.mode == AutoIndex::Adaptive {
239            for (&a, &uses) in &self.uses {
240                let eligible = a >= USER0
241                    && schema
242                        .attribute_for_entid(a)
243                        .is_some_and(|attr| !attr.fulltext);
244                if uses >= self.min_uses
245                    && eligible
246                    && !have.contains(&a)
247                    && !self.rejected.contains(&a)
248                {
249                    actions.push(IndexAction::Create {
250                        index: db::value_index_name(a),
251                        attribute: a,
252                        ident: ident(a),
253                        uses,
254                    });
255                }
256            }
257        }
258        for (index, a, last_used, is_schema) in &managed {
259            if *is_schema {
260                continue; // Follows the schema flag, not the workload.
261            }
262            let used = self.uses.get(a).copied().unwrap_or(0) > 0;
263            let idle = now.saturating_sub(*last_used) as u64 >= self.idle.as_secs();
264            if self.mode == AutoIndex::Schema || (!used && idle) {
265                actions.push(IndexAction::Drop {
266                    index: index.clone(),
267                    attribute: *a,
268                    ident: ident(*a),
269                });
270            }
271        }
272        if dry_run {
273            return Ok(actions);
274        }
275
276        for (_, a, _, _) in &managed {
277            if self.uses.get(a).copied().unwrap_or(0) > 0 {
278                tx.execute(
279                    "UPDATE mentat_managed_indexes SET last_used = ? WHERE a = ?",
280                    rusqlite::params![now, a],
281                )?;
282            }
283        }
284        let mut rejected = BTreeSet::new();
285        for action in &actions {
286            match action {
287                IndexAction::Create {
288                    index, attribute, ..
289                } => {
290                    // Built and fully ANALYZEd (O(attribute), like the build):
291                    // with no stat1 row the planner may probe it last, after
292                    // joining everything else, and `selective` reads it.
293                    db::create_value_index(&tx, *attribute, "adaptive")?;
294                    if !selective(&tx, index)? {
295                        db::drop_value_index(&tx, *attribute)?;
296                        rejected.insert(*attribute);
297                    }
298                }
299                IndexAction::Drop { attribute, .. } => {
300                    db::drop_value_index(&tx, *attribute)?;
301                }
302            }
303        }
304        tx.commit()?;
305        self.indexed = registry(conn)?.into_iter().map(|(_, a, _, _)| a).collect();
306        actions.retain(
307            |a| !matches!(a, IndexAction::Create { attribute, .. } if rejected.contains(attribute)),
308        );
309        self.rejected.extend(rejected);
310        // A new period.
311        self.uses.clear();
312        self.queries = 0;
313        Ok(actions)
314    }
315}
mentat/auto_index.rs

mentat/
auto_index.rs
