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

ursa_engine/graph/
index.rs
