1use std::collections::HashMap;
51use std::fmt::{self, Write as _};
52use std::path::{Path, PathBuf};
53use std::sync::atomic::{AtomicU64, Ordering};
54use std::sync::{Arc, Mutex, OnceLock};
55use std::time::{Duration, Instant};
56
57use datafusion::arrow::buffer::BooleanBuffer;
58use datafusion::catalog::default_table_source::source_as_provider;
59use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion};
60use datafusion::datasource::file_format::arrow::ArrowFormat;
61use datafusion::datasource::file_format::csv::CsvFormat;
62use datafusion::datasource::file_format::json::JsonFormat;
63use datafusion::datasource::file_format::parquet::ParquetFormat;
64use datafusion::datasource::file_format::FileFormat as DataFusionFileFormat;
65use datafusion::datasource::listing::ListingTable;
66use datafusion::datasource::physical_plan::FileScanConfig;
67use datafusion::datasource::source::DataSourceExec;
68use datafusion::error::Result;
69use datafusion::execution::TaskContext;
70use datafusion::logical_expr::LogicalPlan;
71use datafusion::physical_plan::ExecutionPlan;
72use datafusion::prelude::{SessionConfig, SessionContext};
73use futures::future::BoxFuture;
74use object_store::ObjectStoreExt;
75use ursa_core::{IdMap, Topology};
76
77use crate::database::{database_of_config, IndexKey, IndexStore};
78use crate::notices::{self, NoticeLevel, Notifier};
79
80#[derive(Debug)]
88pub struct GraphIndex {
89 pub topology: Arc<Topology>,
90 pub ids: Arc<IdMap>,
91 n_dropped: u64,
93 kept: Option<BooleanBuffer>,
96 rows: OnceLock<Vec<u64>>,
98 weights: WeightCache,
100 stored: Option<StoredWeights>,
102}
103
104pub type WeightReader = Box<dyn Fn(&str) -> Option<Vec<Vec<f64>>> + Send + Sync>;
107
108struct StoredWeights {
110 keys: Vec<String>,
111 read: WeightReader,
112}
113
114impl fmt::Debug for StoredWeights {
115 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
116 write!(f, "StoredWeights({} sets)", self.keys.len())
117 }
118}
119
120impl GraphIndex {
121 pub fn new(topology: Arc<Topology>, ids: Arc<IdMap>) -> Self {
122 Self {
123 topology,
124 ids,
125 n_dropped: 0,
126 kept: None,
127 rows: OnceLock::new(),
128 weights: WeightCache::default(),
129 stored: None,
130 }
131 }
132
133 pub fn with_dropped(mut self, n_dropped: u64, kept: Option<BooleanBuffer>) -> Self {
136 self.n_dropped = n_dropped;
137 self.kept = kept.filter(|_| n_dropped > 0);
138 self
139 }
140
141 pub fn n_dropped(&self) -> u64 {
144 self.n_dropped
145 }
146
147 pub fn kept_rows(&self) -> Option<&BooleanBuffer> {
150 self.kept.as_ref()
151 }
152
153 pub fn edge_row(&self, edge: u32) -> u64 {
156 match &self.kept {
157 None => u64::from(edge),
158 Some(kept) => self
159 .rows
160 .get_or_init(|| kept.set_indices().map(|i| i as u64).collect())[edge as usize],
161 }
162 }
163
164 pub fn with_stored_weights(mut self, keys: Vec<String>, read: WeightReader) -> Self {
167 if !keys.is_empty() {
168 self.stored = Some(StoredWeights { keys, read });
169 }
170 self
171 }
172
173 pub fn weights(&self, key: &str) -> Option<Weights> {
175 self.weights.get(key)
176 }
177
178 pub fn has_stored_weights(&self, key: &str) -> bool {
180 self.stored
181 .as_ref()
182 .is_some_and(|s| s.keys.iter().any(|k| k == key))
183 }
184
185 pub fn stored_weights(&self, key: &str, n_columns: usize) -> Option<Weights> {
189 if let Some(weights) = self.weights(key) {
190 return Some(weights);
191 }
192 if !self.has_stored_weights(key) {
193 return None;
194 }
195 let columns = (self.stored.as_ref()?.read)(key)?;
196 let n_edges = self.topology.n_edges();
197 if columns.len() != n_columns || columns.iter().any(|c| c.len() != n_edges) {
198 return None;
199 }
200 let weights: Weights = Arc::new(columns.into_iter().map(Arc::new).collect());
201 self.store_weights(key.to_string(), Arc::clone(&weights));
202 Some(weights)
203 }
204
205 pub fn weight_sets(&self) -> Vec<(String, Weights)> {
207 self.weights.all()
208 }
209
210 pub fn store_weights(&self, key: String, weights: Weights) {
213 self.weights.put(key, weights)
214 }
215}
216
217pub type Weights = Arc<Vec<Arc<Vec<f64>>>>;
220
221const WEIGHT_CAPACITY: usize = 4;
225
226#[derive(Default)]
230struct WeightCache {
231 entries: Mutex<Vec<(String, Weights)>>,
232}
233
234impl fmt::Debug for WeightCache {
235 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
236 let entries = self.entries.lock().unwrap_or_else(|e| e.into_inner());
237 write!(f, "WeightCache({} entries)", entries.len())
238 }
239}
240
241impl WeightCache {
242 fn get(&self, key: &str) -> Option<Weights> {
243 let mut entries = self.entries.lock().unwrap_or_else(|e| e.into_inner());
244 let i = entries.iter().position(|(k, _)| k == key)?;
245 let entry = entries.remove(i);
246 let weights = Arc::clone(&entry.1);
247 entries.push(entry);
248 Some(weights)
249 }
250
251 fn all(&self) -> Vec<(String, Weights)> {
252 let entries = self.entries.lock().unwrap_or_else(|e| e.into_inner());
253 entries
254 .iter()
255 .map(|(k, w)| (k.clone(), Arc::clone(w)))
256 .collect()
257 }
258
259 fn put(&self, key: String, weights: Weights) {
260 let mut entries = self.entries.lock().unwrap_or_else(|e| e.into_inner());
261 entries.retain(|(k, _)| *k != key);
262 if entries.len() >= WEIGHT_CAPACITY {
263 entries.remove(0);
264 }
265 entries.push((key, weights));
266 }
267}
268
269#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
271pub struct FileIdentity {
272 pub store: String,
273 pub location: String,
274 pub size: u64,
275 pub modified: i64,
277 pub e_tag: Option<String>,
278}
279
280pub fn fingerprint(plan: &LogicalPlan) -> Option<String> {
283 let mut rendered = plan.display_indent_schema().to_string();
284 let mut cacheable = true;
285 let mut locations: Vec<String> = Vec::new();
286 let mut scans: Vec<(String, String)> = Vec::new();
287 let _ = plan.apply_with_subqueries(|node| {
288 if node.expressions().iter().any(|e| e.is_volatile()) {
289 cacheable = false;
290 return Ok(TreeNodeRecursion::Stop);
291 }
292 if let LogicalPlan::TableScan(scan) = node {
293 let listing = source_as_provider(&scan.source)
294 .ok()
295 .and_then(|p| p.downcast_ref::<ListingTable>().map(listing_key));
296 match listing {
297 Some(listing) => {
298 let name = scan.table_name.to_string();
299 scans.push((name.clone(), listing));
300 if crate::paths::looks_like_location(&name) && !locations.contains(&name) {
301 locations.push(name);
302 }
303 }
304 None => {
305 cacheable = false;
306 return Ok(TreeNodeRecursion::Stop);
307 }
308 }
309 }
310 Ok(TreeNodeRecursion::Continue)
311 });
312 if !cacheable {
313 return None;
314 }
315 let placeholder = |name: &str| {
321 locations
322 .iter()
323 .position(|l| l == name)
324 .map(|i| format!("\u{0}file {i}\u{0}"))
325 };
326 let mut by_length: Vec<&String> = locations.iter().collect();
327 by_length.sort_by_key(|name| std::cmp::Reverse(name.len()));
328 for name in by_length {
329 if let Some(placeholder) = placeholder(name) {
330 rendered = rendered.replace(name.as_str(), &placeholder);
331 }
332 }
333 for (name, listing) in &scans {
334 let name = placeholder(name).unwrap_or_else(|| name.clone());
335 let _ = write!(rendered, "\n{name}: {listing}");
336 }
337 Some(rendered)
338}
339
340fn listing_key(table: &ListingTable) -> String {
347 let paths: Vec<String> = table.table_paths().iter().map(|p| p.to_string()).collect();
348 let options = table.options();
349 let partitions: Vec<String> = options
350 .table_partition_cols
351 .iter()
352 .map(|(name, data_type)| format!("{name} {data_type}"))
353 .collect();
354 let order: Vec<String> = options
355 .file_sort_order
356 .iter()
357 .map(|exprs| {
358 exprs
359 .iter()
360 .map(ToString::to_string)
361 .collect::<Vec<_>>()
362 .join(", ")
363 })
364 .collect();
365 format!(
366 "paths {paths:?}; extension {:?}; {}; partitions {partitions:?}; order {order:?}",
367 options.file_extension,
368 format_key(options.format.as_ref())
369 )
370}
371
372fn format_key(format: &dyn DataFusionFileFormat) -> String {
374 if let Some(csv) = format.downcast_ref::<CsvFormat>() {
375 let o = csv.options();
376 return format!(
377 "csv (header {:?}, delimiter {}, quote {}, terminator {:?}, escape {:?}, \
378 double_quote {:?}, comment {:?}, null_regex {:?}, null_value {:?}, \
379 truncated_rows {:?}, compression {:?}, formats {:?})",
380 o.has_header,
381 o.delimiter,
382 o.quote,
383 o.terminator,
384 o.escape,
385 o.double_quote,
386 o.comment,
387 o.null_regex,
388 o.null_value,
389 o.truncated_rows,
390 o.compression,
391 [
392 &o.date_format,
393 &o.datetime_format,
394 &o.timestamp_format,
395 &o.timestamp_tz_format,
396 &o.time_format
397 ],
398 );
399 }
400 if let Some(json) = format.downcast_ref::<JsonFormat>() {
401 let o = json.options();
402 return format!(
403 "json (newline_delimited {}, compression {:?})",
404 o.newline_delimited, o.compression
405 );
406 }
407 if let Some(parquet) = format.downcast_ref::<ParquetFormat>() {
408 let o = &parquet.options().global;
409 return format!(
410 "parquet (binary_as_string {}, schema_force_view_types {}, coerce_int96 {:?} {:?})",
411 o.binary_as_string, o.schema_force_view_types, o.coerce_int96, o.coerce_int96_tz
412 );
413 }
414 if format.downcast_ref::<ArrowFormat>().is_some() {
415 return "arrow".to_string();
416 }
417 format!("{} {format:?}", format.get_ext())
419}
420
421pub(crate) const DATABASE_DIR: &str = "\u{0}database\u{0}";
426
427fn dir_forms(dir: &Path) -> Vec<String> {
432 let mut forms = Vec::new();
433 let canonical = std::fs::canonicalize(dir).ok();
434 for dir in std::iter::once(dir).chain(canonical.as_deref()) {
435 let text = dir.to_string_lossy().replace('\\', "/");
436 let text = text.trim_end_matches('/');
437 if text.is_empty() {
438 continue;
439 }
440 let url = match text.strip_prefix('/') {
441 Some(rest) => format!("file:///{rest}/"),
442 None => format!("file:///{text}/"),
443 };
444 for form in [url, format!("{text}/")] {
445 if !forms.contains(&form) {
446 forms.push(form);
447 }
448 }
449 }
450 forms.sort_by_key(|f| std::cmp::Reverse(f.len()));
451 forms
452}
453
454pub(crate) fn portable(text: &str, dir: &Path) -> String {
457 let mut text = text.to_string();
458 for form in dir_forms(dir) {
459 text = text.replace(&form, &format!("{DATABASE_DIR}/"));
460 }
461 text
462}
463
464pub(crate) fn portable_in(text: &str, config: &SessionConfig) -> String {
467 match database_of_config(config) {
468 Some(db) => portable(text, db.dir()),
469 None => text.to_string(),
470 }
471}
472
473fn local_path(location: &str) -> String {
476 if cfg!(windows) || location.starts_with('/') {
477 location.to_string()
478 } else {
479 format!("/{location}")
480 }
481}
482
483pub async fn file_identities(
486 plan: &Arc<dyn ExecutionPlan>,
487 context: &TaskContext,
488) -> Result<Vec<FileIdentity>> {
489 let mut files = Vec::new();
490 let _ = plan.apply(|node| {
491 if let Some(config) = node
492 .downcast_ref::<DataSourceExec>()
493 .and_then(|exec| exec.data_source().downcast_ref::<FileScanConfig>())
494 {
495 for group in &config.file_groups {
496 for file in group.files() {
497 files.push((
498 config.object_store_url.clone(),
499 file.object_meta.location.clone(),
500 ));
501 }
502 }
503 }
504 Ok(TreeNodeRecursion::Continue)
505 });
506 files.sort();
507 files.dedup();
508 let database = database_of_config(context.session_config());
509 let mut identities = Vec::with_capacity(files.len());
510 for (url, location) in files {
511 let store = context.runtime_env().object_store(&url)?;
512 let meta = store.head(&location).await?;
513 let store = url.to_string();
514 let location = if store == LOCAL_STORE {
517 let path = local_path(location.as_ref());
518 match &database {
519 Some(db) => portable(&path, db.dir()),
520 None => path,
521 }
522 } else {
523 location.to_string()
524 };
525 identities.push(FileIdentity {
526 store,
527 location,
528 size: meta.size,
529 modified: meta.last_modified.timestamp_nanos_opt().unwrap_or(i64::MAX),
530 e_tag: meta.e_tag,
531 });
532 }
533 identities.sort();
534 Ok(identities)
535}
536
537const LOCAL_STORE: &str = "file:///";
539
540const CAPACITY: usize = 16;
543
544#[derive(Debug, Default)]
546pub struct IndexCache {
547 slots: Mutex<HashMap<String, Slot>>,
548 clock: AtomicU64,
550 builds: AtomicU64,
551 loads: AtomicU64,
552}
553
554#[derive(Debug)]
555struct Slot {
556 files: Vec<FileIdentity>,
557 cell: Arc<tokio::sync::OnceCell<Arc<GraphIndex>>>,
558 used: u64,
559}
560
561impl IndexCache {
562 pub fn builds(&self) -> u64 {
564 self.builds.load(Ordering::Relaxed)
565 }
566
567 pub fn loads(&self) -> u64 {
570 self.loads.load(Ordering::Relaxed)
571 }
572
573 pub fn record_build(&self) {
575 self.builds.fetch_add(1, Ordering::Relaxed);
576 }
577
578 pub fn get(&self, key: &(String, Vec<FileIdentity>)) -> Option<Arc<GraphIndex>> {
580 let mut slots = self.slots.lock().unwrap_or_else(|e| e.into_inner());
581 let slot = slots.get_mut(&key.0).filter(|s| s.files == key.1)?;
582 let index = slot.cell.get().cloned()?;
583 slot.used = self.clock.fetch_add(1, Ordering::Relaxed);
584 Some(index)
585 }
586
587 pub async fn get_or_build(
592 &self,
593 key: Option<(String, Vec<FileIdentity>)>,
594 build: BoxFuture<'_, Result<GraphIndex>>,
595 ) -> Result<Arc<GraphIndex>> {
596 self.get_or_build_in(None, &Notifier::silent(), &BuildInfo::default(), key, build)
597 .await
598 }
599
600 pub async fn get_or_build_stored(
605 &self,
606 config: &SessionConfig,
607 key: Option<(String, Vec<FileIdentity>)>,
608 build: BoxFuture<'_, Result<GraphIndex>>,
609 ) -> Result<Arc<GraphIndex>> {
610 self.get_or_build_for(config, key, &BuildInfo::default(), build)
611 .await
612 }
613
614 pub async fn get_or_build_for(
617 &self,
618 config: &SessionConfig,
619 key: Option<(String, Vec<FileIdentity>)>,
620 info: &BuildInfo,
621 build: BoxFuture<'_, Result<GraphIndex>>,
622 ) -> Result<Arc<GraphIndex>> {
623 let stored =
624 database_of_config(config).map(|db| (db.index_store(), db.dir().to_path_buf()));
625 let notifier = Notifier::of(config);
626 self.get_or_build_in(stored, ¬ifier, info, key, build)
627 .await
628 }
629
630 async fn get_or_build_in(
631 &self,
632 stored: Option<(Arc<IndexStore>, PathBuf)>,
633 notifier: &Notifier,
634 info: &BuildInfo,
635 key: Option<(String, Vec<FileIdentity>)>,
636 build: BoxFuture<'_, Result<GraphIndex>>,
637 ) -> Result<Arc<GraphIndex>> {
638 let Some((fingerprint, files)) = key else {
639 self.builds.fetch_add(1, Ordering::Relaxed);
640 return build_noticed(notifier, info, build, Kept::Never).await;
641 };
642 let stored = stored.map(|(store, dir)| {
645 (
646 store,
647 IndexKey::new(portable(&fingerprint, &dir), files.clone()),
648 )
649 });
650 let cell = {
651 let mut slots = self.slots.lock().unwrap_or_else(|e| e.into_inner());
652 let used = self.clock.fetch_add(1, Ordering::Relaxed);
653 match slots.get_mut(&fingerprint) {
654 Some(slot) if slot.files == files => {
655 slot.used = used;
656 Arc::clone(&slot.cell)
657 }
658 _ => {
659 if !slots.contains_key(&fingerprint) && slots.len() >= CAPACITY {
660 if let Some(oldest) = slots
661 .iter()
662 .min_by_key(|(_, s)| s.used)
663 .map(|(k, _)| k.clone())
664 {
665 slots.remove(&oldest);
666 }
667 }
668 let cell = Arc::new(tokio::sync::OnceCell::new());
669 slots.insert(
670 fingerprint,
671 Slot {
672 files,
673 cell: Arc::clone(&cell),
674 used,
675 },
676 );
677 cell
678 }
679 }
680 };
681 let (builds, loads) = (&self.builds, &self.loads);
683 cell.get_or_try_init(|| async move {
684 if let Some((store, key)) = &stored {
685 let (store, key) = (Arc::clone(store), key.clone());
686 let started = clock(notifier, NoticeLevel::All);
687 let reader = notifier.clone();
688 let loaded = crate::run_blocking(move || {
690 let loaded = store.load(&key, &reader);
691 loaded.map(|(index, size)| (index, size, store))
692 })
693 .await
694 .ok()
695 .flatten();
696 if let Some((index, size, store)) = loaded {
697 loads.fetch_add(1, Ordering::Relaxed);
698 notifier.say(NoticeLevel::All, || {
699 format!(
700 "loaded the index for {} ({} edges) from {} ({}){}",
701 info.graph(),
702 notices::count(index.topology.n_edges() as u64),
703 store.database(),
704 notices::size(size),
705 took(started)
706 )
707 });
708 return Ok(index);
709 }
710 }
711 builds.fetch_add(1, Ordering::Relaxed);
712 let kept = match &stored {
713 Some((store, key)) => Kept::Stored(Arc::clone(store), key.clone()),
714 None => Kept::InMemory,
715 };
716 build_noticed(notifier, info, build, kept).await
717 })
718 .await
719 .cloned()
720 }
721}
722
723#[derive(Debug, Clone, Default)]
725pub struct BuildInfo {
726 pub graph: Option<String>,
728 pub edges: Option<(u64, bool)>,
731}
732
733impl BuildInfo {
734 fn graph(&self) -> &str {
735 self.graph.as_deref().unwrap_or("a graph")
736 }
737}
738
739enum Kept {
741 Stored(Arc<IndexStore>, IndexKey),
744 InMemory,
746 Never,
748}
749
750fn clock(notifier: &Notifier, level: NoticeLevel) -> Option<Instant> {
753 (notifier.level() >= level).then(Instant::now)
754}
755
756fn took(started: Option<Instant>) -> String {
758 started
759 .map(|s| format!(" in {}", notices::duration(s.elapsed())))
760 .unwrap_or_default()
761}
762
763async fn build_noticed(
769 notifier: &Notifier,
770 info: &BuildInfo,
771 build: BoxFuture<'_, Result<GraphIndex>>,
772 kept: Kept,
773) -> Result<Arc<GraphIndex>> {
774 let started = clock(notifier, NoticeLevel::On);
775 let estimate = match info.edges {
776 Some((n, true)) if n > 0 => format!(" ({} edges)", notices::count(n)),
777 Some((n, false)) if n > 0 => format!(" (about {} edges)", notices::count(n)),
778 _ => String::new(),
779 };
780 let delay = match notifier.level() {
781 NoticeLevel::All => Duration::ZERO,
782 _ => notices::BUILD_NOTICE_AFTER,
783 };
784 let announcement = notifier.announce_after(
785 delay,
786 format!("building the index for {}{estimate}…", info.graph()),
787 );
788 let built = build.await;
789 let announced = announcement.finish();
790 let index = Arc::new(built?);
791 let (outcome, loses_work) = match kept {
792 Kept::Stored(store, key) => {
793 let saved = store.save(key, Arc::clone(&index), info.graph(), notifier);
794 (
795 format!("; {}", saved.describe(store.database())),
796 saved.loses_work(),
797 )
798 }
799 Kept::InMemory => (
800 "; kept for this session only (open a .ursa database to store it)".to_string(),
801 false,
802 ),
803 Kept::Never => (
804 "; not kept: what its edges read can change (an in-memory table, or a volatile \
805 function such as random())"
806 .to_string(),
807 false,
808 ),
809 };
810 let level = if announced || loses_work {
811 NoticeLevel::On
812 } else {
813 NoticeLevel::All
814 };
815 notifier.say(level, || {
816 format!(
817 "built the index for {} ({} edges){}{outcome}",
818 info.graph(),
819 notices::count(index.topology.n_edges() as u64),
820 took(started)
821 )
822 });
823 Ok(index)
824}
825
826pub(crate) fn with_index_cache(config: SessionConfig) -> SessionConfig {
828 config.with_extension(Arc::new(IndexCache::default()))
829}
830
831pub fn index_builds(ctx: &SessionContext) -> u64 {
833 ctx.state()
834 .config()
835 .get_extension::<IndexCache>()
836 .map(|c| c.builds())
837 .unwrap_or(0)
838}
839
840pub fn index_loads(ctx: &SessionContext) -> u64 {
843 ctx.state()
844 .config()
845 .get_extension::<IndexCache>()
846 .map(|c| c.loads())
847 .unwrap_or(0)
848}
849
850#[cfg(test)]
851mod tests {
852 use super::*;
853
854 #[test]
855 fn paths_inside_the_database_directory_are_portable() {
856 let dir = Path::new("/home/me/proj");
857 let text =
858 "TableScan: /home/me/proj/edges.csv\npaths [\"file:///home/me/proj/edges.csv\"] \
859 /home/me/project2/x.csv /home/me/proj";
860 assert_eq!(
861 portable(text, dir),
862 format!(
863 "TableScan: {DATABASE_DIR}/edges.csv\npaths [\"{DATABASE_DIR}/edges.csv\"] \
864 /home/me/project2/x.csv /home/me/proj"
865 )
866 );
867 assert_eq!(portable("/a/b.csv", Path::new("/")), "/a/b.csv");
869 assert_eq!(
871 portable("/home/me/proj/edges.csv", dir),
872 portable("/mnt/elsewhere/edges.csv", Path::new("/mnt/elsewhere/"))
873 );
874 }
875
876 #[test]
877 fn local_locations_keep_their_leading_slash() {
878 #[cfg(not(windows))]
879 assert_eq!(local_path("home/me/a.csv"), "/home/me/a.csv");
880 assert_eq!(local_path("/home/me/a.csv"), "/home/me/a.csv");
881 }
882}