ursa_engine/graph/
index.rs

1//! Graph indexes, and the session's cache of them.
2//!
3//! A graph node builds its index (the CSR topology and the id map) from the
4//! rows of its edge plan when it executes. The cache keeps built indexes for
5//! the session, keyed by what the edges are:
6//!
7//! - **the edge plan's fingerprint**: the optimized logical plan under the
8//!   node, rendered with its schema, plus, for each file table it scans, the
9//!   table's paths and format options (two CSV tables over one file with
10//!   different delimiters render alike but read differently);
11//! - **the identity of every file the plan reads**: its location, size,
12//!   modification time and e-tag, from the object store's metadata when the
13//!   node executes (a `HEAD` per file, not DataFusion's listing cache, which
14//!   can be stale);
15//! - for a graph with a vertex table, whose keys are its nodes (#311), **the
16//!   keys' plan's fingerprint and the identity of the files it reads** too,
17//!   so a change to the vertex table builds afresh.
18//!
19//! Rewriting a file changes its identity, so the next query builds afresh.
20//! An edge plan that reads anything other than files, literal rows or nothing
21//! (an in-memory table, which `INSERT` can change, or a volatile function such
22//! as `random()`) has no stable identity and is never cached.
23//!
24//! Concurrent requests for one key wait on one build, so several algorithms
25//! over one graph in one query share one index even when DataFusion runs
26//! their nodes at the same time.
27//!
28//! In a session with a `.ursa` database, a miss here looks in the database's
29//! stored indexes under the same full key before building
30//! ([`IndexCache::get_or_build_stored`]), and an index built in a writable
31//! session is stored there (`crate::database::index_store`). Loads from disk
32//! are counted apart from builds ([`index_loads`], [`index_builds`]).
33//! Evaluated weights are kept on the index in memory, and stored with it in
34//! the database (`crate::database::index_store`), so a weighted query over a
35//! stored index in a later process reads no edge rows.
36
37use std::collections::HashMap;
38use std::fmt::{self, Write as _};
39use std::sync::atomic::{AtomicU64, Ordering};
40use std::sync::{Arc, Mutex, OnceLock};
41
42use datafusion::arrow::buffer::BooleanBuffer;
43use datafusion::catalog::default_table_source::source_as_provider;
44use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion};
45use datafusion::datasource::listing::ListingTable;
46use datafusion::datasource::physical_plan::FileScanConfig;
47use datafusion::datasource::source::DataSourceExec;
48use datafusion::error::Result;
49use datafusion::execution::TaskContext;
50use datafusion::logical_expr::LogicalPlan;
51use datafusion::physical_plan::ExecutionPlan;
52use datafusion::prelude::{SessionConfig, SessionContext};
53use futures::future::BoxFuture;
54use object_store::ObjectStoreExt;
55use ursa_core::{IdMap, Topology};
56
57use crate::database::{database_of_config, IndexKey, IndexStore};
58
59/// A built graph: its topology, and the map from dense node indexes back to
60/// the edges' ids.
61///
62/// A graph with a vertex table has the vertex table's keys as its nodes, in
63/// the table's order, and only the edges whose endpoints are both keys (see
64/// `ursa_core::EdgeInterner`): its edge ids number the edges kept, and
65/// [`Self::edge_row`] maps them back to rows of the edge table.
66#[derive(Debug)]
67pub struct GraphIndex {
68    pub topology: Arc<Topology>,
69    pub ids: Arc<IdMap>,
70    /// Edge rows dropped because an endpoint is not a node.
71    n_dropped: u64,
72    /// Which edge rows were kept, one bit per row of the edge plan, when any
73    /// was dropped.
74    kept: Option<BooleanBuffer>,
75    /// Per edge id, its row in the edge plan, when any row was dropped.
76    rows: OnceLock<Vec<u64>>,
77    /// Weights evaluated over the edge rows this index was built from.
78    weights: WeightCache,
79    /// Weights stored with the index on disk, read on first use.
80    stored: Option<StoredWeights>,
81}
82
83/// Reads the weight columns stored under a key, one `f64` per edge row each
84/// (blocking: it reads a file). `None` if they cannot be read.
85pub type WeightReader = Box<dyn Fn(&str) -> Option<Vec<Vec<f64>>> + Send + Sync>;
86
87/// The weights stored with an index on disk: their keys, and how to read them.
88struct StoredWeights {
89    keys: Vec<String>,
90    read: WeightReader,
91}
92
93impl fmt::Debug for StoredWeights {
94    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
95        write!(f, "StoredWeights({} sets)", self.keys.len())
96    }
97}
98
99impl GraphIndex {
100    pub fn new(topology: Arc<Topology>, ids: Arc<IdMap>) -> Self {
101        Self {
102            topology,
103            ids,
104            n_dropped: 0,
105            kept: None,
106            rows: OnceLock::new(),
107            weights: WeightCache::default(),
108            stored: None,
109        }
110    }
111
112    /// This index, built from edge rows of which `kept` says which were kept
113    /// (`None`: all of them), `n_dropped` having been dropped.
114    pub fn with_dropped(mut self, n_dropped: u64, kept: Option<BooleanBuffer>) -> Self {
115        self.n_dropped = n_dropped;
116        self.kept = kept.filter(|_| n_dropped > 0);
117        self
118    }
119
120    /// Edge rows dropped because an endpoint is not a node of the graph's
121    /// vertex table.
122    pub fn n_dropped(&self) -> u64 {
123        self.n_dropped
124    }
125
126    /// Which rows of the edge plan were kept, one bit per row, if any was
127    /// dropped.
128    pub fn kept_rows(&self) -> Option<&BooleanBuffer> {
129        self.kept.as_ref()
130    }
131
132    /// The row of the edge plan edge `edge` was read from: the edge id itself
133    /// when no row was dropped.
134    pub fn edge_row(&self, edge: u32) -> u64 {
135        match &self.kept {
136            None => u64::from(edge),
137            Some(kept) => self
138                .rows
139                .get_or_init(|| kept.set_indices().map(|i| i as u64).collect())[edge as usize],
140        }
141    }
142
143    /// This index, with the weight sets stored on disk under `keys`, which
144    /// `read` reads when one is first asked for.
145    pub fn with_stored_weights(mut self, keys: Vec<String>, read: WeightReader) -> Self {
146        if !keys.is_empty() {
147            self.stored = Some(StoredWeights { keys, read });
148        }
149        self
150    }
151
152    /// The weight columns stored under `key` in memory, if they are.
153    pub fn weights(&self, key: &str) -> Option<Weights> {
154        self.weights.get(key)
155    }
156
157    /// Whether weights are stored under `key` on disk.
158    pub fn has_stored_weights(&self, key: &str) -> bool {
159        self.stored
160            .as_ref()
161            .is_some_and(|s| s.keys.iter().any(|k| k == key))
162    }
163
164    /// The weight columns stored under `key`, in memory or else on disk, if
165    /// they are and have `n_columns` columns of one value per edge row. Read
166    /// from disk, they are kept in memory too. Blocking: it may read a file.
167    pub fn stored_weights(&self, key: &str, n_columns: usize) -> Option<Weights> {
168        if let Some(weights) = self.weights(key) {
169            return Some(weights);
170        }
171        if !self.has_stored_weights(key) {
172            return None;
173        }
174        let columns = (self.stored.as_ref()?.read)(key)?;
175        let n_edges = self.topology.n_edges();
176        if columns.len() != n_columns || columns.iter().any(|c| c.len() != n_edges) {
177            return None;
178        }
179        let weights: Weights = Arc::new(columns.into_iter().map(Arc::new).collect());
180        self.store_weights(key.to_string(), Arc::clone(&weights));
181        Some(weights)
182    }
183
184    /// Every set of weight columns held in memory, with its key.
185    pub fn weight_sets(&self) -> Vec<(String, Weights)> {
186        self.weights.all()
187    }
188
189    /// Store weight columns under `key`. They must be one per edge row of this
190    /// index, in its edge-row order.
191    pub fn store_weights(&self, key: String, weights: Weights) {
192        self.weights.put(key, weights)
193    }
194}
195
196/// Weight columns, one `f64` per edge row each, in the order of a node's
197/// weight expressions.
198pub type Weights = Arc<Vec<Arc<Vec<f64>>>>;
199
200/// How many sets of weight columns an index keeps. Each column is one `f64`
201/// per edge (80 MB at 10M edges), so the bound is small, as in ursa-plan's
202/// `WeightCache`: enough for a few weightings alternating over one graph.
203const WEIGHT_CAPACITY: usize = 4;
204
205/// Evaluated weights, held on the index so they live exactly as long as the
206/// edge rows they were evaluated over. Keyed by the weighted edge plan's
207/// fingerprint and file identities (see `exec`), most recently used last.
208#[derive(Default)]
209struct WeightCache {
210    entries: Mutex<Vec<(String, Weights)>>,
211}
212
213impl fmt::Debug for WeightCache {
214    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
215        let entries = self.entries.lock().unwrap_or_else(|e| e.into_inner());
216        write!(f, "WeightCache({} entries)", entries.len())
217    }
218}
219
220impl WeightCache {
221    fn get(&self, key: &str) -> Option<Weights> {
222        let mut entries = self.entries.lock().unwrap_or_else(|e| e.into_inner());
223        let i = entries.iter().position(|(k, _)| k == key)?;
224        let entry = entries.remove(i);
225        let weights = Arc::clone(&entry.1);
226        entries.push(entry);
227        Some(weights)
228    }
229
230    fn all(&self) -> Vec<(String, Weights)> {
231        let entries = self.entries.lock().unwrap_or_else(|e| e.into_inner());
232        entries
233            .iter()
234            .map(|(k, w)| (k.clone(), Arc::clone(w)))
235            .collect()
236    }
237
238    fn put(&self, key: String, weights: Weights) {
239        let mut entries = self.entries.lock().unwrap_or_else(|e| e.into_inner());
240        entries.retain(|(k, _)| *k != key);
241        if entries.len() >= WEIGHT_CAPACITY {
242            entries.remove(0);
243        }
244        entries.push((key, weights));
245    }
246}
247
248/// The identity of one source file, as its store reports it.
249#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
250pub struct FileIdentity {
251    pub store: String,
252    pub location: String,
253    pub size: u64,
254    /// Modification time, in nanoseconds since the epoch.
255    pub modified: i64,
256    pub e_tag: Option<String>,
257}
258
259/// The fingerprint of an edge plan, or `None` if what it reads has no stable
260/// identity (see the module docs).
261pub fn fingerprint(plan: &LogicalPlan) -> Option<String> {
262    let mut rendered = plan.display_indent_schema().to_string();
263    let mut cacheable = true;
264    let _ = plan.apply_with_subqueries(|node| {
265        if node.expressions().iter().any(|e| e.is_volatile()) {
266            cacheable = false;
267            return Ok(TreeNodeRecursion::Stop);
268        }
269        if let LogicalPlan::TableScan(scan) = node {
270            let listing = source_as_provider(&scan.source).ok().and_then(|p| {
271                p.downcast_ref::<ListingTable>().map(|t| {
272                    let paths: Vec<String> =
273                        t.table_paths().iter().map(|p| p.to_string()).collect();
274                    format!("{paths:?} {:?}", t.options())
275                })
276            });
277            match listing {
278                Some(listing) => {
279                    let _ = write!(rendered, "\n{}: {listing}", scan.table_name);
280                }
281                None => {
282                    cacheable = false;
283                    return Ok(TreeNodeRecursion::Stop);
284                }
285            }
286        }
287        Ok(TreeNodeRecursion::Continue)
288    });
289    cacheable.then_some(rendered)
290}
291
292/// The identity of every file `plan` reads, sorted, looked up in the object
293/// store now.
294pub async fn file_identities(
295    plan: &Arc<dyn ExecutionPlan>,
296    context: &TaskContext,
297) -> Result<Vec<FileIdentity>> {
298    let mut files = Vec::new();
299    let _ = plan.apply(|node| {
300        if let Some(config) = node
301            .downcast_ref::<DataSourceExec>()
302            .and_then(|exec| exec.data_source().downcast_ref::<FileScanConfig>())
303        {
304            for group in &config.file_groups {
305                for file in group.files() {
306                    files.push((
307                        config.object_store_url.clone(),
308                        file.object_meta.location.clone(),
309                    ));
310                }
311            }
312        }
313        Ok(TreeNodeRecursion::Continue)
314    });
315    files.sort();
316    files.dedup();
317    let mut identities = Vec::with_capacity(files.len());
318    for (url, location) in files {
319        let store = context.runtime_env().object_store(&url)?;
320        let meta = store.head(&location).await?;
321        identities.push(FileIdentity {
322            store: url.to_string(),
323            location: location.to_string(),
324            size: meta.size,
325            modified: meta.last_modified.timestamp_nanos_opt().unwrap_or(i64::MAX),
326            e_tag: meta.e_tag,
327        });
328    }
329    Ok(identities)
330}
331
332/// How many fingerprints the cache keeps. A newer identity of a fingerprint
333/// replaces the older one, so this bounds the number of distinct edge plans.
334const CAPACITY: usize = 16;
335
336/// The session's graph indexes (a session config extension).
337#[derive(Debug, Default)]
338pub struct IndexCache {
339    slots: Mutex<HashMap<String, Slot>>,
340    /// Bumped on every use, to find the least recently used slot.
341    clock: AtomicU64,
342    builds: AtomicU64,
343    loads: AtomicU64,
344}
345
346#[derive(Debug)]
347struct Slot {
348    files: Vec<FileIdentity>,
349    cell: Arc<tokio::sync::OnceCell<Arc<GraphIndex>>>,
350    used: u64,
351}
352
353impl IndexCache {
354    /// How many indexes this cache has built.
355    pub fn builds(&self) -> u64 {
356        self.builds.load(Ordering::Relaxed)
357    }
358
359    /// How many indexes this cache has loaded from the database's stored
360    /// indexes instead of building them.
361    pub fn loads(&self) -> u64 {
362        self.loads.load(Ordering::Relaxed)
363    }
364
365    /// Count a build made outside the cache (an index that cannot be shared).
366    pub fn record_build(&self) {
367        self.builds.fetch_add(1, Ordering::Relaxed);
368    }
369
370    /// The index for `key`, if it is built and cached; nothing is built.
371    pub fn get(&self, key: &(String, Vec<FileIdentity>)) -> Option<Arc<GraphIndex>> {
372        let mut slots = self.slots.lock().unwrap_or_else(|e| e.into_inner());
373        let slot = slots.get_mut(&key.0).filter(|s| s.files == key.1)?;
374        let index = slot.cell.get().cloned()?;
375        slot.used = self.clock.fetch_add(1, Ordering::Relaxed);
376        Some(index)
377    }
378
379    /// The index for `key` (a fingerprint and file identities), built by
380    /// `build` if it is not cached. With no fingerprint, always built. The
381    /// database's stored indexes are not consulted: see
382    /// [`Self::get_or_build_stored`].
383    pub async fn get_or_build(
384        &self,
385        key: Option<(String, Vec<FileIdentity>)>,
386        build: BoxFuture<'_, Result<GraphIndex>>,
387    ) -> Result<Arc<GraphIndex>> {
388        self.get_or_build_in(None, key, build).await
389    }
390
391    /// As [`Self::get_or_build`], and on a miss, the index stored under the
392    /// same full key in the database of the session `config` belongs to, if
393    /// it has one. An index built in a writable database is stored there, in
394    /// the background, within the session's `ursa.index_storage_limit`.
395    pub async fn get_or_build_stored(
396        &self,
397        config: &SessionConfig,
398        key: Option<(String, Vec<FileIdentity>)>,
399        build: BoxFuture<'_, Result<GraphIndex>>,
400    ) -> Result<Arc<GraphIndex>> {
401        let stored = database_of_config(config).map(|db| {
402            (
403                db.index_store(),
404                crate::options::index_storage_limit(config),
405            )
406        });
407        self.get_or_build_in(stored, key, build).await
408    }
409
410    async fn get_or_build_in(
411        &self,
412        stored: Option<(Arc<IndexStore>, u64)>,
413        key: Option<(String, Vec<FileIdentity>)>,
414        build: BoxFuture<'_, Result<GraphIndex>>,
415    ) -> Result<Arc<GraphIndex>> {
416        let Some((fingerprint, files)) = key else {
417            self.builds.fetch_add(1, Ordering::Relaxed);
418            return build.await.map(Arc::new);
419        };
420        let stored = stored.map(|(store, limit)| {
421            (
422                store,
423                limit,
424                IndexKey::new(fingerprint.clone(), files.clone()),
425            )
426        });
427        let cell = {
428            let mut slots = self.slots.lock().unwrap_or_else(|e| e.into_inner());
429            let used = self.clock.fetch_add(1, Ordering::Relaxed);
430            match slots.get_mut(&fingerprint) {
431                Some(slot) if slot.files == files => {
432                    slot.used = used;
433                    Arc::clone(&slot.cell)
434                }
435                _ => {
436                    if !slots.contains_key(&fingerprint) && slots.len() >= CAPACITY {
437                        if let Some(oldest) = slots
438                            .iter()
439                            .min_by_key(|(_, s)| s.used)
440                            .map(|(k, _)| k.clone())
441                        {
442                            slots.remove(&oldest);
443                        }
444                    }
445                    let cell = Arc::new(tokio::sync::OnceCell::new());
446                    slots.insert(
447                        fingerprint,
448                        Slot {
449                            files,
450                            cell: Arc::clone(&cell),
451                            used,
452                        },
453                    );
454                    cell
455                }
456            }
457        };
458        // A failed build leaves the cell empty, so the next query tries again.
459        let (builds, loads) = (&self.builds, &self.loads);
460        cell.get_or_try_init(|| async move {
461            if let Some((store, limit, key)) = &stored {
462                let (store, key, limit) = (Arc::clone(store), key.clone(), *limit);
463                // Reading a file and checking it is blocking work.
464                let loaded = crate::run_blocking(move || store.load(&key, limit))
465                    .await
466                    .ok()
467                    .flatten();
468                if let Some(index) = loaded {
469                    loads.fetch_add(1, Ordering::Relaxed);
470                    return Ok(index);
471                }
472            }
473            builds.fetch_add(1, Ordering::Relaxed);
474            let index = Arc::new(build.await?);
475            if let Some((store, limit, key)) = stored {
476                store.save(key, Arc::clone(&index), limit);
477            }
478            Ok(index)
479        })
480        .await
481        .cloned()
482    }
483}
484
485/// Install an empty index cache on `config`.
486pub(crate) fn with_index_cache(config: SessionConfig) -> SessionConfig {
487    config.with_extension(Arc::new(IndexCache::default()))
488}
489
490/// How many graph indexes `ctx` has built: for tests of the cache.
491pub fn index_builds(ctx: &SessionContext) -> u64 {
492    ctx.state()
493        .config()
494        .get_extension::<IndexCache>()
495        .map(|c| c.builds())
496        .unwrap_or(0)
497}
498
499/// How many graph indexes `ctx` has loaded from its database's stored
500/// indexes instead of building them: for tests of stored indexes.
501pub fn index_loads(ctx: &SessionContext) -> u64 {
502    ctx.state()
503        .config()
504        .get_extension::<IndexCache>()
505        .map(|c| c.loads())
506        .unwrap_or(0)
507}
ursa_engine/graph/index.rs

ursa_engine/graph/
index.rs
