1use crate::backend::StorageBackend;
5#[cfg(feature = "lance-backend")]
6use crate::backend::lance::LanceDbBackend;
7use crate::backend::table_names;
8use crate::backend::types::{ScanRequest, VectorQueryOpts};
9use crate::compaction::{CompactionStats, CompactionStatus, CompactionTask};
10use crate::runtime::WorkingGraph;
11use crate::runtime::context::QueryContext;
12use crate::runtime::l0::L0Buffer;
13use crate::storage::adjacency::AdjacencyDataset;
14use crate::storage::compaction::Compactor;
15use crate::storage::delta::{DeltaDataset, ENTRY_SIZE_ESTIMATE, Op};
16use crate::storage::direction::Direction;
17#[cfg(feature = "lance-backend")]
18use crate::storage::edge::EdgeDataset;
19#[cfg(feature = "lance-backend")]
20use crate::storage::index::UidIndex;
21#[cfg(feature = "lance-backend")]
22use crate::storage::inverted_index::InvertedIndex;
23use crate::storage::main_edge::MainEdgeDataset;
24use crate::storage::main_vertex::MainVertexDataset;
25use crate::storage::vertex::VertexDataset;
26use anyhow::{Result, anyhow};
27use arrow_array::{Array, Float32Array, TimestampNanosecondArray, UInt64Array};
28use object_store::ObjectStore;
29#[cfg(feature = "lance-backend")]
30use object_store::local::LocalFileSystem;
31use parking_lot::RwLock;
32use std::collections::{HashMap, HashSet};
33use std::sync::{Arc, Mutex};
34use std::time::{Duration, SystemTime, UNIX_EPOCH};
35use tracing::warn;
36use uni_common::config::UniConfig;
37#[cfg(feature = "lance-backend")]
38use uni_common::core::id::UniId;
39use uni_common::core::id::{Eid, Vid};
40#[cfg(feature = "lance-backend")]
41use uni_common::core::schema::IndexDefinition;
42use uni_common::core::schema::{DistanceMetric, SchemaManager};
43use uni_common::sync::acquire_mutex;
44
45use crate::snapshot::manager::SnapshotManager;
46use crate::storage::IndexManager;
47use crate::storage::adjacency_manager::AdjacencyManager;
48use crate::storage::resilient_store::ResilientObjectStore;
49
50use uni_common::core::snapshot::SnapshotManifest;
51
52use uni_common::graph::simple_graph::Direction as GraphDirection;
53
54struct EdgeState {
56 neighbor: Vid,
57 version: u64,
58 deleted: bool,
59}
60
61pub struct StorageManager {
62 base_uri: String,
63 store: Arc<dyn ObjectStore>,
64 schema_manager: Arc<SchemaManager>,
65 snapshot_manager: Arc<SnapshotManager>,
66 adjacency_manager: Arc<AdjacencyManager>,
67 pub config: UniConfig,
68 pub compaction_status: Arc<Mutex<CompactionStatus>>,
69 pub flush_in_progress: std::sync::atomic::AtomicUsize,
74 pinned_snapshot: Option<SnapshotManifest>,
76 pinned_version_hwm: Option<u64>,
85 fork_scope: Option<Arc<crate::fork::ForkScope>>,
93 backend: Arc<dyn StorageBackend>,
95 vid_labels_index: Arc<parking_lot::RwLock<crate::storage::vid_labels::VidLabelsIndex>>,
103 plugin_registry: Option<Arc<uni_plugin::PluginRegistry>>,
113}
114
115pub struct FlushInProgressGuard {
124 storage: Arc<StorageManager>,
125}
126
127impl FlushInProgressGuard {
128 pub fn new(storage: &Arc<StorageManager>) -> Self {
129 storage
130 .flush_in_progress
131 .fetch_add(1, std::sync::atomic::Ordering::AcqRel);
132 Self {
133 storage: storage.clone(),
134 }
135 }
136}
137
138impl Drop for FlushInProgressGuard {
139 fn drop(&mut self) {
140 self.storage
142 .flush_in_progress
143 .fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
144 }
145}
146
147fn is_lance_conflict(err: &anyhow::Error) -> bool {
153 let msg = err.to_string();
154 msg.contains("Incompatible transaction") || msg.contains("conflict")
155}
156
157async fn retry_on_lance_conflict<F, Fut>(mut op: F) -> anyhow::Result<()>
162where
163 F: FnMut() -> Fut,
164 Fut: std::future::Future<Output = anyhow::Result<()>>,
165{
166 for attempt in 0u32..10 {
167 match op().await {
168 Ok(()) => return Ok(()),
169 Err(e) => {
170 if !is_lance_conflict(&e) || attempt == 9 {
171 return Err(e);
172 }
173 let backoff_ms = 1u64 << attempt;
174 tokio::time::sleep(std::time::Duration::from_millis(backoff_ms)).await;
175 }
176 }
177 }
178 unreachable!("retry loop exits via Ok or Err")
179}
180
181pub async fn merge_insert_batch_with_lance_conflict_retry(
190 backend: &dyn crate::backend::StorageBackend,
191 table_name: &str,
192 batch: arrow_array::RecordBatch,
193 on: &[&str],
194) -> anyhow::Result<()> {
195 retry_on_lance_conflict(|| async {
201 let exists = backend.table_exists(table_name).await?;
202 if !exists {
203 anyhow::bail!(
204 "merge_insert target table '{}' does not exist (partial writes \
205 require the row to already be present; CREATE goes through Append)",
206 table_name
207 );
208 }
209 backend
210 .merge_insert(table_name, on, vec![batch.clone()])
211 .await
212 })
213 .await
214}
215
216pub async fn write_batch_with_lance_conflict_retry(
225 backend: &dyn crate::backend::StorageBackend,
226 table_name: &str,
227 batch: arrow_array::RecordBatch,
228) -> anyhow::Result<()> {
229 use crate::backend::types::WriteMode;
230 retry_on_lance_conflict(|| async {
231 let exists = backend.table_exists(table_name).await?;
232 if exists {
233 backend
234 .write(table_name, vec![batch.clone()], WriteMode::Append)
235 .await
236 } else {
237 backend.create_table(table_name, vec![batch.clone()]).await
238 }
239 })
240 .await
241}
242
243struct CompactionGuard {
245 status: Arc<Mutex<CompactionStatus>>,
246}
247
248impl CompactionGuard {
249 fn new(status: Arc<Mutex<CompactionStatus>>) -> Option<Self> {
250 let mut s = acquire_mutex(&status, "compaction_status").ok()?;
251 if s.compaction_in_progress {
252 return None;
253 }
254 s.compaction_in_progress = true;
255 Some(Self {
256 status: status.clone(),
257 })
258 }
259}
260
261impl Drop for CompactionGuard {
262 fn drop(&mut self) {
263 match uni_common::sync::acquire_mutex(&self.status, "compaction_status") {
266 Ok(mut s) => {
267 s.compaction_in_progress = false;
268 s.last_compaction = Some(std::time::SystemTime::now());
269 }
270 Err(e) => {
271 log::error!(
275 "CompactionGuard drop failed to acquire poisoned lock: {}. \
276 Compaction status may be inconsistent. Issue #18/#150",
277 e
278 );
279 }
280 }
281 }
282}
283
284impl StorageManager {
285 pub async fn new_with_backend(
287 base_uri: &str,
288 store: Arc<dyn ObjectStore>,
289 backend: Arc<dyn StorageBackend>,
290 schema_manager: Arc<SchemaManager>,
291 config: UniConfig,
292 ) -> Result<Self> {
293 let resilient_store: Arc<dyn ObjectStore> = Arc::new(ResilientObjectStore::new(
294 store,
295 config.object_store.clone(),
296 ));
297
298 let snapshot_manager = Arc::new(SnapshotManager::new(resilient_store.clone()));
299
300 Self::recover_all_staging_tables(backend.as_ref(), &schema_manager).await?;
302
303 let mut sm = Self {
304 base_uri: base_uri.to_string(),
305 store: resilient_store,
306 schema_manager,
307 snapshot_manager,
308 adjacency_manager: Arc::new(AdjacencyManager::new(config.cache_size)),
309 config,
310 compaction_status: Arc::new(Mutex::new(CompactionStatus::default())),
311 flush_in_progress: std::sync::atomic::AtomicUsize::new(0),
312 pinned_snapshot: None,
313 pinned_version_hwm: None,
314 fork_scope: None,
315 backend,
316 vid_labels_index: Arc::new(parking_lot::RwLock::new(
317 crate::storage::vid_labels::VidLabelsIndex::new(),
318 )),
319 plugin_registry: None,
320 };
321
322 if let Err(e) = sm.rebuild_vid_labels_index().await {
326 warn!(
327 "Failed to rebuild VidLabelsIndex on startup: {}. Falling back to storage queries.",
328 e
329 );
330 }
331
332 Ok(sm)
333 }
334
335 #[cfg(feature = "lance-backend")]
337 pub async fn new(base_uri: &str, schema_manager: Arc<SchemaManager>) -> Result<Self> {
338 Self::new_with_config(base_uri, schema_manager, UniConfig::default()).await
339 }
340
341 #[cfg(feature = "lance-backend")]
343 pub async fn new_with_cache(
344 base_uri: &str,
345 schema_manager: Arc<SchemaManager>,
346 adjacency_cache_size: usize,
347 ) -> Result<Self> {
348 let config = UniConfig {
349 cache_size: adjacency_cache_size,
350 ..Default::default()
351 };
352 Self::new_with_config(base_uri, schema_manager, config).await
353 }
354
355 #[cfg(feature = "lance-backend")]
357 pub async fn new_with_config(
358 base_uri: &str,
359 schema_manager: Arc<SchemaManager>,
360 config: UniConfig,
361 ) -> Result<Self> {
362 let store = Self::build_store_from_uri(base_uri)?;
363 Self::new_with_store_and_config(base_uri, store, schema_manager, config).await
364 }
365
366 #[cfg(feature = "lance-backend")]
368 pub async fn new_with_store_and_config(
369 base_uri: &str,
370 store: Arc<dyn ObjectStore>,
371 schema_manager: Arc<SchemaManager>,
372 config: UniConfig,
373 ) -> Result<Self> {
374 Self::new_with_store_and_storage_options(base_uri, store, schema_manager, config, None)
375 .await
376 }
377
378 #[cfg(feature = "lance-backend")]
380 pub async fn new_with_store_and_storage_options(
381 base_uri: &str,
382 store: Arc<dyn ObjectStore>,
383 schema_manager: Arc<SchemaManager>,
384 config: UniConfig,
385 lancedb_storage_options: Option<HashMap<String, String>>,
386 ) -> Result<Self> {
387 let backend = Arc::new(LanceDbBackend::connect(base_uri, lancedb_storage_options).await?);
388 Self::new_with_backend(base_uri, store, backend, schema_manager, config).await
389 }
390
391 async fn recover_all_staging_tables(
396 backend: &dyn StorageBackend,
397 schema_manager: &SchemaManager,
398 ) -> Result<()> {
399 let schema = schema_manager.schema();
400
401 backend
403 .recover_staging(table_names::main_vertex_table_name())
404 .await?;
405 backend
406 .recover_staging(table_names::main_edge_table_name())
407 .await?;
408
409 for label in schema.labels.keys() {
411 let name = table_names::vertex_table_name(label);
412 backend.recover_staging(&name).await?;
413 }
414
415 for edge_type in schema.edge_types.keys() {
417 for direction in &["fwd", "bwd"] {
418 let delta_name = table_names::delta_table_name(edge_type, direction);
420 backend.recover_staging(&delta_name).await?;
421
422 for _label in schema.labels.keys() {
424 let adj_name = table_names::adjacency_table_name(edge_type, direction);
425 backend.recover_staging(&adj_name).await?;
426 }
427 }
428 }
429
430 Ok(())
431 }
432
433 #[cfg(feature = "lance-backend")]
434 fn build_store_from_uri(base_uri: &str) -> Result<Arc<dyn ObjectStore>> {
435 if base_uri.contains("://") {
436 let parsed = url::Url::parse(base_uri).map_err(|e| anyhow!("Invalid base URI: {e}"))?;
437 let (store, _path) = object_store::parse_url(&parsed)
438 .map_err(|e| anyhow!("Failed to parse object store URL: {e}"))?;
439 Ok(Arc::from(store))
440 } else {
441 std::fs::create_dir_all(base_uri)?;
443 Ok(Arc::new(LocalFileSystem::new_with_prefix(base_uri)?))
444 }
445 }
446
447 pub fn local_fs_root(&self) -> Option<std::path::PathBuf> {
453 if self.base_uri.contains("://") {
454 None
455 } else {
456 Some(std::path::PathBuf::from(&self.base_uri))
457 }
458 }
459
460 pub fn pinned(&self, snapshot: SnapshotManifest) -> Self {
461 Self {
469 base_uri: self.base_uri.clone(),
470 store: self.store.clone(),
471 schema_manager: self.schema_manager.clone(),
472 snapshot_manager: self.snapshot_manager.clone(),
473 adjacency_manager: Arc::new(AdjacencyManager::new(self.adjacency_manager.max_bytes())),
477 config: self.config.clone(),
478 compaction_status: Arc::new(Mutex::new(CompactionStatus::default())),
479 flush_in_progress: std::sync::atomic::AtomicUsize::new(0),
480 pinned_snapshot: Some(snapshot),
481 pinned_version_hwm: None,
482 fork_scope: self.fork_scope.clone(),
483 backend: self.backend.clone(),
484 vid_labels_index: Arc::new(parking_lot::RwLock::new(
490 self.vid_labels_index.read().clone(),
491 )),
492 plugin_registry: self.plugin_registry.clone(),
493 }
494 }
495
496 pub fn pinned_at_version(&self, hwm: u64) -> Self {
515 Self {
516 base_uri: self.base_uri.clone(),
517 store: self.store.clone(),
518 schema_manager: self.schema_manager.clone(),
519 snapshot_manager: self.snapshot_manager.clone(),
520 adjacency_manager: self.adjacency_manager.clone(),
521 config: self.config.clone(),
522 compaction_status: Arc::new(Mutex::new(CompactionStatus::default())),
523 flush_in_progress: std::sync::atomic::AtomicUsize::new(0),
524 pinned_snapshot: None,
525 pinned_version_hwm: Some(hwm),
526 fork_scope: self.fork_scope.clone(),
527 backend: self.backend.clone(),
528 vid_labels_index: Arc::new(parking_lot::RwLock::new(
534 self.vid_labels_index.read().clone(),
535 )),
536 plugin_registry: self.plugin_registry.clone(),
537 }
538 }
539
540 pub fn at_fork(&self, scope: Arc<crate::fork::ForkScope>) -> Self {
555 self.at_fork_with_schema(scope, self.schema_manager.clone())
556 }
557
558 pub fn at_fork_with_schema(
566 &self,
567 scope: Arc<crate::fork::ForkScope>,
568 merged_schema: Arc<SchemaManager>,
569 ) -> Self {
570 debug_assert!(
571 self.pinned_snapshot.is_none(),
572 "forking a pinned StorageManager is unsupported in Phase 1"
573 );
574 let branched_backend: Arc<dyn StorageBackend> = Arc::new(
575 crate::backend::branched::BranchedBackend::new(self.backend.clone(), scope.clone()),
576 );
577 let snapshot_manager = Arc::new(SnapshotManager::new_for_fork(
582 self.store.clone(),
583 scope.fork_id(),
584 ));
585 Self {
586 base_uri: self.base_uri.clone(),
587 store: self.store.clone(),
588 schema_manager: merged_schema,
589 snapshot_manager,
590 adjacency_manager: Arc::new(AdjacencyManager::new(self.adjacency_manager.max_bytes())),
591 config: self.config.clone(),
592 compaction_status: Arc::new(Mutex::new(CompactionStatus::default())),
593 flush_in_progress: std::sync::atomic::AtomicUsize::new(0),
594 pinned_snapshot: None,
595 pinned_version_hwm: None,
596 fork_scope: Some(scope),
597 backend: branched_backend,
598 vid_labels_index: Arc::new(parking_lot::RwLock::new(
604 self.vid_labels_index.read().clone(),
605 )),
606 plugin_registry: self.plugin_registry.clone(),
607 }
608 }
609
610 pub fn plugin_registry(&self) -> Option<&Arc<uni_plugin::PluginRegistry>> {
615 self.plugin_registry.as_ref()
616 }
617
618 pub fn set_plugin_registry(&mut self, registry: Arc<uni_plugin::PluginRegistry>) {
624 self.plugin_registry = Some(registry);
625 }
626
627 pub fn fork_scope(&self) -> Option<&Arc<crate::fork::ForkScope>> {
629 self.fork_scope.as_ref()
630 }
631
632 #[must_use]
645 pub fn has_fork_index(
646 &self,
647 label: &str,
648 column: &str,
649 kind: crate::fork::ForkLocalIndexKind,
650 ) -> bool {
651 self.fork_scope
652 .as_ref()
653 .is_some_and(|s| s.has_fork_local_index(label, column, kind))
654 }
655
656 pub fn base_uri(&self) -> &str {
659 &self.base_uri
660 }
661
662 pub fn get_edge_version_by_id(&self, edge_type_id: u32) -> Option<u64> {
663 let schema = self.schema_manager.schema();
664 let name = schema.edge_type_name_by_id(edge_type_id)?;
665 self.pinned_snapshot
666 .as_ref()
667 .and_then(|s| s.edges.get(name).map(|es| es.lance_version))
668 .filter(|v| *v != 0)
673 }
674
675 pub fn version_high_water_mark(&self) -> Option<u64> {
683 self.pinned_snapshot
684 .as_ref()
685 .map(|s| s.version_high_water_mark)
686 .or(self.pinned_version_hwm)
687 }
688
689 pub fn snapshot_version_hwm(&self) -> Option<u64> {
701 self.pinned_snapshot
702 .as_ref()
703 .map(|s| s.version_high_water_mark)
704 }
705
706 pub fn apply_version_filter(&self, base_filter: String) -> String {
711 if let Some(hwm) = self.version_high_water_mark() {
712 format!("({}) AND (_version <= {})", base_filter, hwm)
713 } else {
714 base_filter
715 }
716 }
717
718 fn build_active_filter(user_filter: Option<&str>) -> String {
721 match user_filter {
722 Some(expr) => format!("({}) AND (_deleted = false)", expr),
723 None => "_deleted = false".to_string(),
724 }
725 }
726
727 pub fn store(&self) -> Arc<dyn ObjectStore> {
728 self.store.clone()
729 }
730
731 pub fn compaction_status(
737 &self,
738 ) -> Result<CompactionStatus, uni_common::sync::LockPoisonedError> {
739 let guard = uni_common::sync::acquire_mutex(&self.compaction_status, "compaction_status")?;
740 Ok(guard.clone())
741 }
742
743 pub async fn compact(&self) -> Result<CompactionStats> {
744 let start = std::time::Instant::now();
746 let schema = self.schema_manager.schema();
747 let mut files_compacted = 0;
748
749 for label in schema.labels.keys() {
750 let name = table_names::vertex_table_name(label);
751 if self.backend.table_exists(&name).await? {
752 self.backend.optimize_table(&name).await?;
753 files_compacted += 1;
754 self.backend.invalidate_cache(&name);
755 }
756 }
757
758 Ok(CompactionStats {
759 files_compacted,
760 bytes_before: 0,
761 bytes_after: 0,
762 duration: start.elapsed(),
763 crdt_merges: 0,
764 })
765 }
766
767 pub async fn compact_label(&self, label: &str) -> Result<CompactionStats> {
768 let _guard = CompactionGuard::new(self.compaction_status.clone())
769 .ok_or_else(|| anyhow!("Compaction already in progress"))?;
770
771 let start = std::time::Instant::now();
772 let name = table_names::vertex_table_name(label);
773
774 if self.backend.table_exists(&name).await? {
775 self.backend.optimize_table(&name).await?;
776 self.backend.invalidate_cache(&name);
777 }
778
779 Ok(CompactionStats {
780 files_compacted: 1,
781 bytes_before: 0,
782 bytes_after: 0,
783 duration: start.elapsed(),
784 crdt_merges: 0,
785 })
786 }
787
788 pub async fn compact_edge_type(&self, edge_type: &str) -> Result<CompactionStats> {
789 let _guard = CompactionGuard::new(self.compaction_status.clone())
790 .ok_or_else(|| anyhow!("Compaction already in progress"))?;
791
792 let start = std::time::Instant::now();
793 let mut files_compacted = 0;
794
795 for dir in ["fwd", "bwd"] {
796 let name = table_names::delta_table_name(edge_type, dir);
797 if self.backend.table_exists(&name).await? {
798 self.backend.optimize_table(&name).await?;
799 files_compacted += 1;
800 }
801 }
802
803 Ok(CompactionStats {
804 files_compacted,
805 bytes_before: 0,
806 bytes_after: 0,
807 duration: start.elapsed(),
808 crdt_merges: 0,
809 })
810 }
811
812 pub async fn wait_for_compaction(&self) -> Result<()> {
813 loop {
814 let in_progress = {
815 acquire_mutex(&self.compaction_status, "compaction_status")?.compaction_in_progress
816 };
817 if !in_progress {
818 return Ok(());
819 }
820 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
821 }
822 }
823
824 pub fn start_background_compaction(
825 self: Arc<Self>,
826 mut shutdown_rx: tokio::sync::broadcast::Receiver<()>,
827 ) -> tokio::task::JoinHandle<()> {
828 if !self.config.compaction.enabled {
829 return tokio::spawn(async {});
830 }
831
832 tokio::spawn(async move {
833 let start = tokio::time::Instant::now() + self.config.compaction.check_interval;
839 let mut interval =
840 tokio::time::interval_at(start, self.config.compaction.check_interval);
841
842 loop {
843 tokio::select! {
844 _ = interval.tick() => {
845 if let Err(e) = self.update_compaction_status().await {
846 log::error!("Failed to update compaction status: {}", e);
847 continue;
848 }
849
850 if let Some(task) = self.pick_compaction_task() {
851 log::info!("Triggering background compaction: {:?}", task);
852 if let Err(e) = Self::execute_compaction(Arc::clone(&self), task).await {
853 log::error!("Compaction failed: {}", e);
854 }
855 }
856 }
857 _ = shutdown_rx.recv() => {
858 log::info!("Background compaction shutting down");
859 let _ = self.wait_for_compaction().await;
860 break;
861 }
862 }
863 }
864 })
865 }
866
867 async fn update_compaction_status(&self) -> Result<()> {
868 let schema = self.schema_manager.schema();
869 let backend = self.backend.as_ref();
870 let mut total_rows: usize = 0;
871 let mut oldest_ts: Option<i64> = None;
872
873 for name in schema.edge_types.keys() {
874 for dir in ["fwd", "bwd"] {
875 let tbl_name = table_names::delta_table_name(name, dir);
876 if !backend.table_exists(&tbl_name).await? {
877 continue;
878 }
879 let row_count = backend.count_rows(&tbl_name, None).await.unwrap_or(0);
880 if row_count == 0 {
881 continue;
882 }
883 total_rows += row_count;
884
885 let request =
887 ScanRequest::all(&tbl_name).with_columns(vec!["_created_at".to_string()]);
888 let Ok(batches) = backend.scan(request).await else {
889 continue;
890 };
891 for batch in batches {
892 let Some(col) = batch
893 .column_by_name("_created_at")
894 .and_then(|c| c.as_any().downcast_ref::<TimestampNanosecondArray>())
895 else {
896 continue;
897 };
898 for i in 0..col.len() {
899 if !col.is_null(i) {
900 let ts = col.value(i);
901 oldest_ts = Some(oldest_ts.map_or(ts, |prev| prev.min(ts)));
902 }
903 }
904 }
905 }
906 }
907
908 let oldest_l1_age = oldest_ts
909 .and_then(|ts| {
910 let created = UNIX_EPOCH + Duration::from_nanos(ts as u64);
911 SystemTime::now().duration_since(created).ok()
912 })
913 .unwrap_or(Duration::ZERO);
914
915 let mut status = acquire_mutex(&self.compaction_status, "compaction_status")?;
916 status.l1_size_bytes = (total_rows * ENTRY_SIZE_ESTIMATE) as u64;
919 status.oldest_l1_age = oldest_l1_age;
920 Ok(())
921 }
922
923 fn pick_compaction_task(&self) -> Option<CompactionTask> {
924 let status = acquire_mutex(&self.compaction_status, "compaction_status").ok()?;
925
926 if status.l1_runs >= self.config.compaction.max_l1_runs {
927 return Some(CompactionTask::ByRunCount);
928 }
929 if status.l1_size_bytes >= self.config.compaction.max_l1_size_bytes {
930 return Some(CompactionTask::BySize);
931 }
932 if status.oldest_l1_age >= self.config.compaction.max_l1_age
933 && status.oldest_l1_age > Duration::ZERO
934 {
935 return Some(CompactionTask::ByAge);
936 }
937
938 None
939 }
940
941 async fn try_optimize_table(backend: &dyn StorageBackend, table_name: &str) -> bool {
943 if let Err(e) = backend.optimize_table(table_name).await {
944 log::warn!("Failed to optimize table {}: {}", table_name, e);
945 return false;
946 }
947 true
948 }
949
950 pub fn trigger_async_compaction(self: &Arc<Self>) {
953 let this = Arc::clone(self);
954 tokio::spawn(async move {
955 if let Err(e) = Self::execute_compaction(this, CompactionTask::ByRunCount).await {
956 log::debug!("Post-flush compaction skipped: {}", e);
958 }
959 });
960 }
961
962 pub(crate) async fn execute_compaction(
963 this: Arc<Self>,
964 _task: CompactionTask,
965 ) -> Result<CompactionStats> {
966 let start = std::time::Instant::now();
967 let _guard = CompactionGuard::new(this.compaction_status.clone())
968 .ok_or_else(|| anyhow!("Compaction already in progress"))?;
969
970 let schema = this.schema_manager.schema();
971 let mut files_compacted = 0;
972
973 let compactor = Compactor::new(Arc::clone(&this));
976 let compaction_results = compactor.compact_all().await.unwrap_or_else(|e| {
977 log::error!(
978 "Semantic compaction failed (continuing with backend optimize): {}",
979 e
980 );
981 Vec::new()
982 });
983
984 let am = this.adjacency_manager();
986 for info in &compaction_results {
987 let direction = match info.direction.as_str() {
988 "fwd" => Direction::Outgoing,
989 "bwd" => Direction::Incoming,
990 _ => continue,
991 };
992 if let Some(etid) = schema.edge_type_id_unified_case_insensitive(&info.edge_type)
993 && let Err(e) = am.warm(&this, etid, direction, None).await
994 {
995 log::warn!(
996 "Failed to re-warm adjacency for {}/{}: {}",
997 info.edge_type,
998 info.direction,
999 e
1000 );
1001 }
1002 }
1003
1004 let backend = this.backend.as_ref();
1006
1007 for name in schema.edge_types.keys() {
1009 for dir in ["fwd", "bwd"] {
1010 let delta = table_names::delta_table_name(name, dir);
1011 if Self::try_optimize_table(backend, &delta).await {
1012 files_compacted += 1;
1013 }
1014 let adj = table_names::adjacency_table_name(name, dir);
1015 if Self::try_optimize_table(backend, &adj).await {
1016 files_compacted += 1;
1017 }
1018 }
1019 }
1020
1021 for label in schema.labels.keys() {
1023 let tbl = table_names::vertex_table_name(label);
1024 if Self::try_optimize_table(backend, &tbl).await {
1025 files_compacted += 1;
1026 backend.invalidate_cache(&tbl);
1027 }
1028 }
1029
1030 for tbl in [
1032 table_names::main_vertex_table_name(),
1033 table_names::main_edge_table_name(),
1034 ] {
1035 if Self::try_optimize_table(backend, tbl).await {
1036 files_compacted += 1;
1037 }
1038 }
1039
1040 {
1041 let mut status = acquire_mutex(&this.compaction_status, "compaction_status")?;
1042 status.total_compactions += 1;
1043 status.l1_runs = 0; }
1045
1046 Ok(CompactionStats {
1047 files_compacted,
1048 bytes_before: 0,
1049 bytes_after: 0,
1050 duration: start.elapsed(),
1051 crdt_merges: 0,
1052 })
1053 }
1054
1055 pub fn invalidate_table_cache(&self, label: &str) {
1059 let name = table_names::vertex_table_name(label);
1060 self.backend.invalidate_cache(&name);
1061 }
1062
1063 pub fn base_path(&self) -> &str {
1064 &self.base_uri
1065 }
1066
1067 pub fn schema_manager(&self) -> &SchemaManager {
1068 &self.schema_manager
1069 }
1070
1071 pub fn schema_manager_arc(&self) -> Arc<SchemaManager> {
1072 self.schema_manager.clone()
1073 }
1074
1075 #[must_use]
1083 pub fn schema_manager_arc_ref(&self) -> &Arc<SchemaManager> {
1084 &self.schema_manager
1085 }
1086
1087 pub fn adjacency_manager(&self) -> Arc<AdjacencyManager> {
1089 Arc::clone(&self.adjacency_manager)
1090 }
1091
1092 pub async fn warm_adjacency(
1097 &self,
1098 edge_type_id: u32,
1099 direction: crate::storage::direction::Direction,
1100 version: Option<u64>,
1101 ) -> anyhow::Result<()> {
1102 self.adjacency_manager
1103 .warm(self, edge_type_id, direction, version)
1104 .await
1105 }
1106
1107 pub async fn warm_adjacency_coalesced(
1112 &self,
1113 edge_type_id: u32,
1114 direction: crate::storage::direction::Direction,
1115 version: Option<u64>,
1116 ) -> anyhow::Result<()> {
1117 self.adjacency_manager
1118 .warm_coalesced(self, edge_type_id, direction, version)
1119 .await
1120 }
1121
1122 pub fn has_adjacency_csr(
1124 &self,
1125 edge_type_id: u32,
1126 direction: crate::storage::direction::Direction,
1127 ) -> bool {
1128 self.adjacency_manager.has_csr(edge_type_id, direction)
1129 }
1130
1131 pub fn get_neighbors_at_version(
1133 &self,
1134 vid: uni_common::core::id::Vid,
1135 edge_type: u32,
1136 direction: crate::storage::direction::Direction,
1137 version: u64,
1138 ) -> Vec<(uni_common::core::id::Vid, uni_common::core::id::Eid)> {
1139 self.adjacency_manager
1140 .get_neighbors_at_version(vid, edge_type, direction, version)
1141 }
1142
1143 pub fn backend(&self) -> &dyn StorageBackend {
1145 self.backend.as_ref()
1146 }
1147
1148 pub fn backend_arc(&self) -> Arc<dyn StorageBackend> {
1150 self.backend.clone()
1151 }
1152
1153 async fn rebuild_vid_labels_index(&mut self) -> Result<()> {
1158 use crate::storage::vid_labels::VidLabelsIndex;
1159
1160 let backend = self.backend.as_ref();
1161 let vtable = table_names::main_vertex_table_name();
1162
1163 if !backend.table_exists(vtable).await? {
1165 self.vid_labels_index = Arc::new(parking_lot::RwLock::new(VidLabelsIndex::new()));
1166 return Ok(());
1167 }
1168
1169 let request = ScanRequest::all(vtable)
1171 .with_filter("_deleted = false")
1172 .with_limit(100_000);
1173 let batches = backend
1174 .scan(request)
1175 .await
1176 .map_err(|e| anyhow!("Failed to query main vertex table: {}", e))?;
1177
1178 let mut index = VidLabelsIndex::new();
1179 for batch in batches {
1180 let vid_col = batch
1181 .column_by_name("_vid")
1182 .ok_or_else(|| anyhow!("Missing _vid column"))?
1183 .as_any()
1184 .downcast_ref::<UInt64Array>()
1185 .ok_or_else(|| anyhow!("Invalid _vid column type"))?;
1186
1187 let labels_col = batch
1188 .column_by_name("labels")
1189 .ok_or_else(|| anyhow!("Missing labels column"))?
1190 .as_any()
1191 .downcast_ref::<arrow_array::ListArray>()
1192 .ok_or_else(|| anyhow!("Invalid labels column type"))?;
1193
1194 for row_idx in 0..batch.num_rows() {
1195 let vid = Vid::from(vid_col.value(row_idx));
1196 let labels_array = labels_col.value(row_idx);
1197 let labels_str_array = labels_array
1198 .as_any()
1199 .downcast_ref::<arrow_array::StringArray>()
1200 .ok_or_else(|| anyhow!("Invalid labels array element type"))?;
1201
1202 let labels: Vec<String> = (0..labels_str_array.len())
1203 .map(|i| labels_str_array.value(i).to_string())
1204 .collect();
1205
1206 index.insert(vid, labels);
1207 }
1208 }
1209
1210 self.vid_labels_index = Arc::new(parking_lot::RwLock::new(index));
1211 Ok(())
1212 }
1213
1214 pub fn get_labels_from_index(&self, vid: Vid) -> Option<Vec<String>> {
1219 let index = self.vid_labels_index.read();
1220 index.get_labels(vid).map(|labels| labels.to_vec())
1221 }
1222
1223 pub fn update_vid_labels_index(&self, vid: Vid, labels: Vec<String>) {
1225 let mut index = self.vid_labels_index.write();
1226 index.insert(vid, labels);
1227 }
1228
1229 pub fn remove_from_vid_labels_index(&self, vid: Vid) {
1231 let mut index = self.vid_labels_index.write();
1232 index.remove_vid(vid);
1233 }
1234
1235 pub async fn load_subgraph_cached(
1236 &self,
1237 start_vids: &[Vid],
1238 edge_types: &[u32],
1239 max_hops: usize,
1240 direction: GraphDirection,
1241 _l0: Option<Arc<RwLock<L0Buffer>>>,
1242 ) -> Result<WorkingGraph> {
1243 let mut graph = WorkingGraph::new();
1244
1245 let dir = match direction {
1246 GraphDirection::Outgoing => crate::storage::direction::Direction::Outgoing,
1247 GraphDirection::Incoming => crate::storage::direction::Direction::Incoming,
1248 };
1249
1250 let neighbor_is_dst = matches!(direction, GraphDirection::Outgoing);
1251
1252 let mut frontier: Vec<Vid> = start_vids.to_vec();
1254 let mut visited: HashSet<Vid> = HashSet::new();
1255
1256 for &vid in start_vids {
1258 graph.add_vertex(vid);
1259 }
1260
1261 for _hop in 0..max_hops {
1262 let mut next_frontier = HashSet::new();
1263
1264 for &vid in &frontier {
1265 if visited.contains(&vid) {
1266 continue;
1267 }
1268 visited.insert(vid);
1269 graph.add_vertex(vid);
1270
1271 for &etype_id in edge_types {
1272 let edge_ver = self.snapshot_version_hwm();
1276 self.adjacency_manager
1277 .warm_coalesced(self, etype_id, dir, edge_ver)
1278 .await?;
1279
1280 let edges = self.adjacency_manager.get_neighbors(vid, etype_id, dir);
1282
1283 for (neighbor_vid, eid) in edges {
1284 graph.add_vertex(neighbor_vid);
1285 if !visited.contains(&neighbor_vid) {
1286 next_frontier.insert(neighbor_vid);
1287 }
1288
1289 if neighbor_is_dst {
1290 graph.add_edge(vid, neighbor_vid, eid, etype_id);
1291 } else {
1292 graph.add_edge(neighbor_vid, vid, eid, etype_id);
1293 }
1294 }
1295 }
1296 }
1297 frontier = next_frontier.into_iter().collect();
1298
1299 if frontier.is_empty() {
1301 break;
1302 }
1303 }
1304
1305 Ok(graph)
1306 }
1307
1308 pub fn snapshot_manager(&self) -> &SnapshotManager {
1309 &self.snapshot_manager
1310 }
1311
1312 pub fn index_manager(&self) -> IndexManager {
1313 IndexManager::new(&self.base_uri, self.schema_manager.clone())
1314 .with_backend(self.backend_arc())
1315 }
1316
1317 pub async fn scan_vertex_table(
1326 &self,
1327 label: &str,
1328 columns: &[&str],
1329 additional_filter: Option<&str>,
1330 ) -> Result<Option<arrow_array::RecordBatch>> {
1331 let backend = self.backend();
1332 let table_name = table_names::vertex_table_name(label);
1333
1334 if !backend.table_exists(&table_name).await? {
1335 return Ok(None);
1336 }
1337
1338 let actual_columns =
1340 if let Some(table_schema) = backend.get_table_schema(&table_name).await? {
1341 let table_field_names: HashSet<&str> = table_schema
1342 .fields()
1343 .iter()
1344 .map(|f| f.name().as_str())
1345 .collect();
1346 columns
1347 .iter()
1348 .copied()
1349 .filter(|c| table_field_names.contains(c))
1350 .map(|s| s.to_string())
1351 .collect::<Vec<_>>()
1352 } else {
1353 return Ok(None);
1354 };
1355
1356 let filter = match (self.version_high_water_mark(), additional_filter) {
1358 (Some(hwm), Some(f)) => Some(format!("_version <= {} AND ({})", hwm, f)),
1359 (Some(hwm), None) => Some(format!("_version <= {}", hwm)),
1360 (None, Some(f)) => Some(f.to_string()),
1361 (None, None) => None,
1362 };
1363
1364 let mut request = ScanRequest::all(&table_name).with_columns(actual_columns);
1365 if let Some(f) = filter {
1366 request = request.with_filter(f);
1367 }
1368
1369 let batches = backend.scan(request).await?;
1376 if batches.is_empty() {
1377 Ok(None)
1378 } else {
1379 Ok(Some(arrow::compute::concat_batches(
1380 &batches[0].schema(),
1381 &batches,
1382 )?))
1383 }
1384 }
1385
1386 pub async fn scan_delta_table(
1389 &self,
1390 edge_type: &str,
1391 direction: &str,
1392 columns: &[&str],
1393 additional_filter: Option<&str>,
1394 ) -> Result<Option<arrow_array::RecordBatch>> {
1395 let edge_hwm = self.snapshot_version_hwm();
1401 let backend = self.backend();
1402 let table_name = table_names::delta_table_name(edge_type, direction);
1403
1404 if !backend.table_exists(&table_name).await? {
1405 return Ok(None);
1406 }
1407
1408 let actual_columns =
1410 if let Some(table_schema) = backend.get_table_schema(&table_name).await? {
1411 let table_field_names: HashSet<&str> = table_schema
1412 .fields()
1413 .iter()
1414 .map(|f| f.name().as_str())
1415 .collect();
1416 columns
1417 .iter()
1418 .copied()
1419 .filter(|c| table_field_names.contains(c))
1420 .map(|s| s.to_string())
1421 .collect::<Vec<_>>()
1422 } else {
1423 return Ok(None);
1424 };
1425
1426 let filter = match (edge_hwm, additional_filter) {
1427 (Some(hwm), Some(f)) => Some(format!("_version <= {} AND ({})", hwm, f)),
1428 (Some(hwm), None) => Some(format!("_version <= {}", hwm)),
1429 (None, Some(f)) => Some(f.to_string()),
1430 (None, None) => None,
1431 };
1432
1433 let mut request = ScanRequest::all(&table_name).with_columns(actual_columns);
1434 if let Some(f) = filter {
1435 request = request.with_filter(f);
1436 }
1437
1438 let batches = backend.scan(request).await?;
1445 if batches.is_empty() {
1446 Ok(None)
1447 } else {
1448 Ok(Some(arrow::compute::concat_batches(
1449 &batches[0].schema(),
1450 &batches,
1451 )?))
1452 }
1453 }
1454
1455 pub async fn scan_main_vertex_table(
1460 &self,
1461 columns: &[&str],
1462 filter: Option<&str>,
1463 ) -> Result<Option<arrow_array::RecordBatch>> {
1464 let backend = self.backend();
1465 let table_name = table_names::main_vertex_table_name();
1466
1467 if !backend.table_exists(table_name).await? {
1468 return Ok(None);
1469 }
1470
1471 let full_filter = match (self.version_high_water_mark(), filter) {
1473 (Some(hwm), Some(f)) => Some(format!("_version <= {} AND ({})", hwm, f)),
1474 (Some(hwm), None) => Some(format!("_version <= {}", hwm)),
1475 (None, Some(f)) => Some(f.to_string()),
1476 (None, None) => None,
1477 };
1478
1479 let request = ScanRequest::all(table_name)
1480 .with_columns(columns.iter().map(|s| s.to_string()).collect());
1481 let request = match full_filter.as_deref() {
1482 Some(f) => request.with_filter(f),
1483 None => request,
1484 };
1485
1486 let batches = backend.scan(request).await?;
1493 if batches.is_empty() {
1494 Ok(None)
1495 } else {
1496 Ok(Some(arrow::compute::concat_batches(
1497 &batches[0].schema(),
1498 &batches,
1499 )?))
1500 }
1501 }
1502
1503 pub async fn scan_main_edge_table_stream(
1505 &self,
1506 filter: Option<&str>,
1507 ) -> Result<
1508 Option<
1509 std::pin::Pin<Box<dyn futures::Stream<Item = Result<arrow_array::RecordBatch>> + Send>>,
1510 >,
1511 > {
1512 let backend = self.backend();
1513 let table_name = table_names::main_edge_table_name();
1514
1515 if !backend.table_exists(table_name).await? {
1516 return Ok(None);
1517 }
1518
1519 let mut request = ScanRequest::all(table_name);
1520 if let Some(f) = filter {
1521 request = request.with_filter(f);
1522 }
1523
1524 let stream = backend.scan_stream(request).await?;
1525 Ok(Some(stream))
1526 }
1527
1528 pub async fn scan_vertex_table_stream(
1530 &self,
1531 label: &str,
1532 ) -> Result<
1533 Option<
1534 std::pin::Pin<Box<dyn futures::Stream<Item = Result<arrow_array::RecordBatch>> + Send>>,
1535 >,
1536 > {
1537 let backend = self.backend();
1538 let table_name = table_names::vertex_table_name(label);
1539
1540 if !backend.table_exists(&table_name).await? {
1541 return Ok(None);
1542 }
1543
1544 let stream = backend.scan_stream(ScanRequest::all(&table_name)).await?;
1545 Ok(Some(stream))
1546 }
1547
1548 pub async fn find_vertex_by_ext_id(&self, ext_id: &str) -> Result<Option<Vid>> {
1550 MainVertexDataset::find_by_ext_id(self.backend(), ext_id, self.version_high_water_mark())
1551 .await
1552 }
1553
1554 pub async fn get_vertex_ext_ids(&self) -> Result<std::collections::HashMap<Vid, String>> {
1566 use arrow_array::StringArray;
1567 let backend = self.backend.as_ref();
1568 let vtable = table_names::main_vertex_table_name();
1569 let mut out = std::collections::HashMap::new();
1570 if !backend.table_exists(vtable).await? {
1571 return Ok(out);
1572 }
1573 let request = ScanRequest::all(vtable)
1574 .with_filter("_deleted = false")
1575 .with_columns(vec!["_vid".to_string(), "ext_id".to_string()]);
1576 let batches = backend
1577 .scan(request)
1578 .await
1579 .map_err(|e| anyhow!("get_vertex_ext_ids: {}", e))?;
1580 for batch in batches {
1581 let vids = batch
1582 .column_by_name("_vid")
1583 .and_then(|c| c.as_any().downcast_ref::<UInt64Array>())
1584 .ok_or_else(|| anyhow!("get_vertex_ext_ids: missing/invalid _vid column"))?;
1585 let exts = batch
1586 .column_by_name("ext_id")
1587 .and_then(|c| c.as_any().downcast_ref::<StringArray>())
1588 .ok_or_else(|| anyhow!("get_vertex_ext_ids: missing/invalid ext_id column"))?;
1589 for i in 0..batch.num_rows() {
1590 if exts.is_null(i) {
1591 continue;
1592 }
1593 let ext = exts.value(i);
1594 if !ext.is_empty() {
1595 out.insert(Vid::from(vids.value(i)), ext.to_string());
1596 }
1597 }
1598 }
1599 Ok(out)
1600 }
1601
1602 pub async fn find_vertex_labels_by_vid(&self, vid: Vid) -> Result<Option<Vec<String>>> {
1604 MainVertexDataset::find_labels_by_vid(self.backend(), vid, self.version_high_water_mark())
1605 .await
1606 }
1607
1608 pub async fn find_edges_by_type_names(
1611 &self,
1612 type_names: &[&str],
1613 endpoint_filter: Option<(crate::storage::main_edge::EndpointSide, &[Vid])>,
1614 ) -> Result<Vec<(Eid, Vid, Vid, String, uni_common::Properties)>> {
1615 MainEdgeDataset::find_edges_by_type_names(self.backend(), type_names, endpoint_filter).await
1616 }
1617
1618 pub async fn scan_vertex_candidates(
1620 &self,
1621 label: &str,
1622 filter: Option<&str>,
1623 ) -> Result<Vec<Vid>> {
1624 let backend = self.backend();
1625 let table_name = table_names::vertex_table_name(label);
1626
1627 if !backend.table_exists(&table_name).await? {
1628 return Ok(Vec::new());
1629 }
1630
1631 let full_filter = match filter {
1632 Some(f) => format!("_deleted = false AND ({})", f),
1633 None => "_deleted = false".to_string(),
1634 };
1635
1636 let request = ScanRequest::all(&table_name)
1637 .with_filter(full_filter)
1638 .with_columns(vec!["_vid".to_string()]);
1639
1640 let batches = backend.scan(request).await?;
1641
1642 let mut vids = Vec::new();
1643 for batch in batches {
1644 let vid_col = batch
1645 .column_by_name("_vid")
1646 .ok_or(anyhow!("Missing _vid"))?
1647 .as_any()
1648 .downcast_ref::<UInt64Array>()
1649 .ok_or(anyhow!("Invalid _vid"))?;
1650 for i in 0..batch.num_rows() {
1651 vids.push(Vid::from(vid_col.value(i)));
1652 }
1653 }
1654 Ok(vids)
1655 }
1656
1657 pub fn vertex_dataset(&self, label: &str) -> Result<VertexDataset> {
1663 let schema = self.schema_manager.schema();
1664 let label_meta = schema
1665 .labels
1666 .get(label)
1667 .ok_or_else(|| anyhow!("Label '{}' not found", label))?;
1668 Ok(VertexDataset::new(&self.base_uri, label, label_meta.id))
1669 }
1670
1671 #[cfg(feature = "lance-backend")]
1672 pub fn edge_dataset(
1673 &self,
1674 edge_type: &str,
1675 src_label: &str,
1676 dst_label: &str,
1677 ) -> Result<EdgeDataset> {
1678 let key = format!("edges_{edge_type}");
1679 match self.fork_branch_for(&key) {
1680 Some(branch) => Ok(EdgeDataset::new_branched(
1681 &self.base_uri,
1682 edge_type,
1683 src_label,
1684 dst_label,
1685 branch,
1686 )),
1687 None => Ok(EdgeDataset::new(
1688 &self.base_uri,
1689 edge_type,
1690 src_label,
1691 dst_label,
1692 )),
1693 }
1694 }
1695
1696 pub fn delta_dataset(&self, edge_type: &str, direction: &str) -> Result<DeltaDataset> {
1697 let key = format!("deltas_{edge_type}_{direction}");
1698 match self.fork_branch_for(&key) {
1699 Some(branch) => Ok(DeltaDataset::new_branched(
1700 &self.base_uri,
1701 edge_type,
1702 direction,
1703 branch,
1704 )),
1705 None => Ok(DeltaDataset::new(&self.base_uri, edge_type, direction)),
1706 }
1707 }
1708
1709 pub fn adjacency_dataset(
1710 &self,
1711 edge_type: &str,
1712 label: &str,
1713 direction: &str,
1714 ) -> Result<AdjacencyDataset> {
1715 let key = crate::backend::table_names::adjacency_table_name(edge_type, direction);
1722 match self.fork_branch_for(&key) {
1723 Some(branch) => Ok(AdjacencyDataset::new_branched(
1724 &self.base_uri,
1725 edge_type,
1726 label,
1727 direction,
1728 branch,
1729 )),
1730 None => Ok(AdjacencyDataset::new(
1731 &self.base_uri,
1732 edge_type,
1733 label,
1734 direction,
1735 )),
1736 }
1737 }
1738
1739 fn fork_branch_for(&self, dataset_name: &str) -> Option<String> {
1743 self.fork_scope
1744 .as_ref()
1745 .and_then(|s| s.branch_for(dataset_name))
1746 }
1747
1748 pub fn main_vertex_dataset(&self) -> MainVertexDataset {
1753 MainVertexDataset::new(&self.base_uri)
1754 }
1755
1756 pub fn main_edge_dataset(&self) -> MainEdgeDataset {
1761 MainEdgeDataset::new(&self.base_uri)
1762 }
1763
1764 #[cfg(feature = "lance-backend")]
1765 pub fn uid_index(&self, label: &str) -> Result<UidIndex> {
1766 Ok(UidIndex::new(&self.base_uri, label))
1767 }
1768
1769 #[cfg(feature = "lance-backend")]
1770 pub async fn inverted_index(&self, label: &str, property: &str) -> Result<InvertedIndex> {
1771 let schema = self.schema_manager.schema();
1772 let config = schema
1773 .indexes
1774 .iter()
1775 .find_map(|idx| match idx {
1776 IndexDefinition::Inverted(cfg)
1777 if cfg.label == label && cfg.property == property =>
1778 {
1779 Some(cfg.clone())
1780 }
1781 _ => None,
1782 })
1783 .ok_or_else(|| anyhow!("Inverted index not found for {}.{}", label, property))?;
1784
1785 InvertedIndex::new(&self.base_uri, config).await
1786 }
1787
1788 #[expect(clippy::too_many_arguments)]
1789 pub async fn vector_search(
1790 &self,
1791 label: &str,
1792 property: &str,
1793 query: &[f32],
1794 k: usize,
1795 filter: Option<&str>,
1796 opts: VectorQueryOpts,
1797 ctx: Option<&QueryContext>,
1798 ) -> Result<Vec<(Vid, f32)>> {
1799 use crate::backend::types::{DistanceMetric as BackendMetric, FilterExpr};
1800
1801 let schema = self.schema_manager.schema();
1803
1804 if let Some(meta) = schema
1809 .properties
1810 .get(label)
1811 .and_then(|props| props.get(property))
1812 && let uni_common::core::schema::DataType::Vector { dimensions } = &meta.r#type
1813 && query.len() != *dimensions
1814 {
1815 return Err(anyhow!(
1816 "vector dimension mismatch: query vector has {} dimensions but '{}.{}' is \
1817 declared VECTOR({})",
1818 query.len(),
1819 label,
1820 property,
1821 dimensions
1822 ));
1823 }
1824
1825 let metric = schema
1826 .vector_index_for_property(label, property)
1827 .map(|config| config.metric.clone())
1828 .unwrap_or(DistanceMetric::L2);
1829
1830 let backend = self.backend.as_ref();
1831 let name = table_names::vertex_table_name(label);
1832
1833 let mut results = Vec::new();
1834
1835 if backend.table_exists(&name).await? {
1837 let backend_metric = match &metric {
1838 DistanceMetric::L2 => BackendMetric::L2,
1839 DistanceMetric::Cosine => BackendMetric::Cosine,
1840 DistanceMetric::Dot => BackendMetric::Dot,
1841 DistanceMetric::L1 => BackendMetric::L2,
1846 _ => BackendMetric::L2,
1847 };
1848
1849 let mut filter_parts = vec![Self::build_active_filter(filter)];
1851 if ctx.is_some()
1852 && let Some(hwm) = self.version_high_water_mark()
1853 {
1854 filter_parts.push(format!("_version <= {}", hwm));
1855 }
1856 let combined_filter = FilterExpr::Sql(filter_parts.join(" AND "));
1857
1858 let fetch_k = if matches!(metric, DistanceMetric::L1) {
1862 backend.count_rows(&name, None).await.unwrap_or(k).max(k)
1863 } else {
1864 k
1865 };
1866
1867 let batches = backend
1868 .vector_search(
1869 &name,
1870 property,
1871 query,
1872 fetch_k,
1873 backend_metric,
1874 combined_filter,
1875 opts,
1876 )
1877 .await?;
1878
1879 results = extract_vid_and_vector_pairs(&batches, "_vid", property)?
1888 .into_iter()
1889 .filter(|(_, emb)| emb.len() == query.len())
1890 .map(|(vid, emb)| (vid, metric.compute_distance(&emb, query)))
1891 .collect();
1892 }
1893
1894 if let Some(qctx) = ctx {
1896 merge_l0_into_vector_results(&mut results, qctx, label, property, query, k, &metric);
1897 }
1898
1899 results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
1904 results.truncate(k);
1905
1906 Ok(results)
1907 }
1908
1909 #[expect(clippy::too_many_arguments)]
1920 pub async fn muvera_fde_candidates(
1921 &self,
1922 label: &str,
1923 fde_column: &str,
1924 fde_query: &[f32],
1925 k: usize,
1926 filter: Option<&str>,
1927 opts: VectorQueryOpts,
1928 ctx: Option<&QueryContext>,
1929 ) -> Result<Vec<(Vid, f32)>> {
1930 use crate::backend::types::{DistanceMetric as BackendMetric, FilterExpr};
1931
1932 let backend = self.backend.as_ref();
1933 let name = table_names::vertex_table_name(label);
1934 if !backend.table_exists(&name).await? {
1938 return Ok(Vec::new());
1939 }
1940
1941 let mut filter_parts = vec![Self::build_active_filter(filter)];
1942 if ctx.is_some()
1943 && let Some(hwm) = self.version_high_water_mark()
1944 {
1945 filter_parts.push(format!("_version <= {}", hwm));
1946 }
1947 let combined_filter = FilterExpr::Sql(filter_parts.join(" AND "));
1948
1949 let batches = backend
1950 .vector_search(
1951 &name,
1952 fde_column,
1953 fde_query,
1954 k,
1955 BackendMetric::Dot,
1956 combined_filter,
1957 opts,
1958 )
1959 .await?;
1960 extract_vid_score_pairs(&batches, "_vid", "_distance")
1961 }
1962
1963 #[expect(clippy::too_many_arguments)]
1984 pub async fn multivector_search(
1985 &self,
1986 label: &str,
1987 property: &str,
1988 query: &[Vec<f32>],
1989 k: usize,
1990 filter: Option<&str>,
1991 opts: VectorQueryOpts,
1992 ctx: Option<&QueryContext>,
1993 ) -> Result<Vec<(Vid, f32)>> {
1994 use crate::backend::types::{DistanceMetric as BackendMetric, FilterExpr};
1995
1996 let schema = self.schema_manager.schema();
1997 let metric = schema
1998 .vector_index_for_property(label, property)
1999 .map(|config| config.metric.clone())
2000 .unwrap_or(DistanceMetric::Cosine);
2001
2002 let backend = self.backend.as_ref();
2003 let name = table_names::vertex_table_name(label);
2004
2005 let branched = self
2016 .fork_scope
2017 .as_ref()
2018 .is_some_and(|s| s.branch_for(&name).is_some());
2019 if branched {
2020 if !backend.table_exists(&name).await? {
2023 return Ok(Vec::new());
2024 }
2025 let mut filter_parts = vec![Self::build_active_filter(filter)];
2026 if ctx.is_some()
2027 && let Some(hwm) = self.version_high_water_mark()
2028 {
2029 filter_parts.push(format!("_version <= {}", hwm));
2030 }
2031 let request = ScanRequest::all(&name)
2032 .with_filter(filter_parts.join(" AND "))
2033 .with_columns(vec!["_vid".to_string()]);
2034 let batches = backend.scan(request).await?;
2035 let mut results = Vec::new();
2036 for batch in batches {
2037 let vid_col = batch
2038 .column_by_name("_vid")
2039 .ok_or(anyhow!("Missing _vid"))?
2040 .as_any()
2041 .downcast_ref::<UInt64Array>()
2042 .ok_or(anyhow!("Invalid _vid"))?;
2043 for i in 0..batch.num_rows() {
2044 results.push((Vid::from(vid_col.value(i)), 0.0_f32));
2045 }
2046 }
2047 return Ok(results);
2048 }
2049
2050 let mut results = Vec::new();
2051 if backend.table_exists(&name).await? {
2053 let backend_metric = match &metric {
2054 DistanceMetric::L2 => BackendMetric::L2,
2055 DistanceMetric::Cosine => BackendMetric::Cosine,
2056 DistanceMetric::Dot => BackendMetric::Dot,
2057 _ => BackendMetric::Cosine,
2058 };
2059
2060 let mut filter_parts = vec![Self::build_active_filter(filter)];
2061 if ctx.is_some()
2062 && let Some(hwm) = self.version_high_water_mark()
2063 {
2064 filter_parts.push(format!("_version <= {}", hwm));
2065 }
2066 let combined_filter = FilterExpr::Sql(filter_parts.join(" AND "));
2067
2068 let batches = backend
2069 .multivector_search(
2070 &name,
2071 property,
2072 query,
2073 k,
2074 backend_metric,
2075 combined_filter,
2076 opts,
2077 )
2078 .await?;
2079 results = extract_vid_score_pairs(&batches, "_vid", "_distance")?;
2080 }
2081
2082 Ok(results)
2083 }
2084
2085 pub async fn sparse_search(
2098 &self,
2099 label: &str,
2100 property: &str,
2101 query: &[(u32, f32)],
2102 k: usize,
2103 ) -> Result<Vec<(Vid, f32)>> {
2104 #[cfg(feature = "lance-backend")]
2105 {
2106 let name = table_names::vertex_table_name(label);
2107 let branched = self
2108 .fork_scope
2109 .as_ref()
2110 .is_some_and(|s| s.branch_for(&name).is_some());
2111 if branched {
2112 let backend = self.backend.as_ref();
2134 if !backend.table_exists(&name).await? {
2137 return Ok(Vec::new());
2138 }
2139 let request = ScanRequest::all(&name)
2140 .with_filter(Self::build_active_filter(None))
2141 .with_columns(vec!["_vid".to_string()]);
2142 let batches = backend.scan(request).await?;
2143 let mut results = Vec::new();
2144 for batch in batches {
2145 let vid_col = batch
2146 .column_by_name("_vid")
2147 .ok_or_else(|| anyhow!("Missing _vid"))?
2148 .as_any()
2149 .downcast_ref::<UInt64Array>()
2150 .ok_or_else(|| anyhow!("Invalid _vid"))?;
2151 for i in 0..batch.num_rows() {
2152 results.push((Vid::from(vid_col.value(i)), 0.0_f32));
2153 }
2154 }
2155 return Ok(results);
2156 }
2157 match self
2158 .index_manager()
2159 .sparse_vector_index(label, property)
2160 .await
2161 {
2162 Ok(idx) => idx.query_topk(query, k).await,
2163 Err(_) => Ok(Vec::new()),
2164 }
2165 }
2166 #[cfg(not(feature = "lance-backend"))]
2167 {
2168 let _ = (label, property, query, k);
2169 Ok(Vec::new())
2170 }
2171 }
2172
2173 pub async fn fts_search(
2189 &self,
2190 label: &str,
2191 property: &str,
2192 query: &str,
2193 k: usize,
2194 filter: Option<&str>,
2195 ctx: Option<&QueryContext>,
2196 ) -> Result<Vec<(Vid, f32)>> {
2197 use crate::backend::types::FilterExpr;
2198
2199 let backend = self.backend.as_ref();
2200 let name = table_names::vertex_table_name(label);
2201
2202 let mut results = if backend.table_exists(&name).await? {
2203 let mut filter_parts = vec![Self::build_active_filter(filter)];
2205 if ctx.is_some()
2206 && let Some(hwm) = self.version_high_water_mark()
2207 {
2208 filter_parts.push(format!("_version <= {}", hwm));
2209 }
2210 let combined_filter = FilterExpr::Sql(filter_parts.join(" AND "));
2211
2212 let batches = backend
2213 .full_text_search(&name, property, query, k, combined_filter)
2214 .await?;
2215
2216 let mut fts_results = extract_vid_score_pairs(&batches, "_vid", "_score")?;
2217 fts_results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
2219 fts_results
2220 } else {
2221 Vec::new()
2222 };
2223
2224 if let Some(qctx) = ctx {
2226 merge_l0_into_fts_results(&mut results, qctx, label, property, query, k);
2227 }
2228
2229 Ok(results)
2230 }
2231
2232 #[cfg(feature = "lance-backend")]
2233 pub async fn get_vertex_by_uid(&self, uid: &UniId, label: &str) -> Result<Option<Vid>> {
2234 let index = self.uid_index(label)?;
2235 index.get_vid(uid).await
2236 }
2237
2238 #[cfg(feature = "lance-backend")]
2239 pub async fn insert_vertex_with_uid(&self, label: &str, vid: Vid, uid: UniId) -> Result<()> {
2240 let index = self.uid_index(label)?;
2241 index.write_mapping(&[(uid, vid)]).await
2242 }
2243
2244 pub async fn load_subgraph(
2245 &self,
2246 start_vids: &[Vid],
2247 edge_types: &[u32],
2248 max_hops: usize,
2249 direction: GraphDirection,
2250 l0: Option<&L0Buffer>,
2251 ) -> Result<WorkingGraph> {
2252 let mut graph = WorkingGraph::new();
2253 let schema = self.schema_manager.schema();
2254
2255 let label_map: HashMap<u16, String> = schema
2257 .labels
2258 .values()
2259 .map(|meta| {
2260 (
2261 meta.id,
2262 schema.label_name_by_id(meta.id).unwrap().to_owned(),
2263 )
2264 })
2265 .collect();
2266
2267 let edge_type_map: HashMap<u32, String> = schema
2268 .edge_types
2269 .values()
2270 .map(|meta| {
2271 (
2272 meta.id,
2273 schema.edge_type_name_by_id(meta.id).unwrap().to_owned(),
2274 )
2275 })
2276 .collect();
2277
2278 let target_edge_types: HashSet<u32> = edge_types.iter().copied().collect();
2279
2280 let mut frontier: Vec<Vid> = start_vids.to_vec();
2282 let mut visited: HashSet<Vid> = HashSet::new();
2283
2284 for &vid in start_vids {
2286 graph.add_vertex(vid);
2287 }
2288
2289 for _hop in 0..max_hops {
2290 let mut next_frontier = HashSet::new();
2291
2292 for &vid in &frontier {
2293 if visited.contains(&vid) {
2294 continue;
2295 }
2296 visited.insert(vid);
2297 graph.add_vertex(vid);
2298
2299 for &etype_id in &target_edge_types {
2301 let etype_name = edge_type_map
2302 .get(&etype_id)
2303 .ok_or_else(|| anyhow!("Unknown edge type ID: {}", etype_id))?;
2304
2305 let (dir_str, neighbor_is_dst) = match direction {
2309 GraphDirection::Outgoing => ("fwd", true),
2310 GraphDirection::Incoming => ("bwd", false),
2311 };
2312
2313 let mut edges: HashMap<Eid, EdgeState> = HashMap::new();
2314
2315 let _edge_ver = self
2320 .pinned_snapshot
2321 .as_ref()
2322 .and_then(|s| s.edges.get(etype_name).map(|es| es.lance_version));
2323
2324 let backend = self.backend();
2326 for current_src_label in label_map.values() {
2327 let adj_ds =
2328 match self.adjacency_dataset(etype_name, current_src_label, dir_str) {
2329 Ok(ds) => ds,
2330 Err(_) => continue,
2331 };
2332 if let Some((neighbors, eids)) =
2333 adj_ds.read_adjacency_backend(backend, vid).await?
2334 {
2335 for (n, eid) in neighbors.into_iter().zip(eids) {
2336 edges.insert(
2337 eid,
2338 EdgeState {
2339 neighbor: n,
2340 version: 0,
2341 deleted: false,
2342 },
2343 );
2344 }
2345 break; }
2347 }
2348
2349 let delta_ds = self.delta_dataset(etype_name, dir_str)?;
2351 let delta_entries = delta_ds
2352 .read_deltas(backend, vid, &schema, self.snapshot_version_hwm())
2353 .await?;
2354 Self::apply_delta_to_edges(&mut edges, delta_entries, neighbor_is_dst);
2355
2356 if let Some(l0) = l0 {
2358 Self::apply_l0_to_edges(&mut edges, l0, vid, etype_id, direction);
2359 }
2360
2361 Self::add_edges_to_graph(
2363 &mut graph,
2364 edges,
2365 vid,
2366 etype_id,
2367 neighbor_is_dst,
2368 &visited,
2369 &mut next_frontier,
2370 );
2371 }
2372 }
2373 frontier = next_frontier.into_iter().collect();
2374
2375 if frontier.is_empty() {
2377 break;
2378 }
2379 }
2380
2381 Ok(graph)
2382 }
2383
2384 fn apply_delta_to_edges(
2386 edges: &mut HashMap<Eid, EdgeState>,
2387 delta_entries: Vec<crate::storage::delta::L1Entry>,
2388 neighbor_is_dst: bool,
2389 ) {
2390 for entry in delta_entries {
2391 let neighbor = if neighbor_is_dst {
2392 entry.dst_vid
2393 } else {
2394 entry.src_vid
2395 };
2396 let current_ver = edges.get(&entry.eid).map(|s| s.version).unwrap_or(0);
2397
2398 if entry.version > current_ver {
2399 edges.insert(
2400 entry.eid,
2401 EdgeState {
2402 neighbor,
2403 version: entry.version,
2404 deleted: matches!(entry.op, Op::Delete),
2405 },
2406 );
2407 }
2408 }
2409 }
2410
2411 fn apply_l0_to_edges(
2413 edges: &mut HashMap<Eid, EdgeState>,
2414 l0: &L0Buffer,
2415 vid: Vid,
2416 etype_id: u32,
2417 direction: GraphDirection,
2418 ) {
2419 let l0_neighbors = l0.get_neighbors(vid, etype_id, direction);
2420 for (neighbor, eid, ver) in l0_neighbors {
2421 let current_ver = edges.get(&eid).map(|s| s.version).unwrap_or(0);
2422 if ver > current_ver {
2423 edges.insert(
2424 eid,
2425 EdgeState {
2426 neighbor,
2427 version: ver,
2428 deleted: false,
2429 },
2430 );
2431 }
2432 }
2433
2434 for (eid, state) in edges.iter_mut() {
2436 if l0.is_tombstoned(*eid) {
2437 state.deleted = true;
2438 }
2439 }
2440 }
2441
2442 fn add_edges_to_graph(
2444 graph: &mut WorkingGraph,
2445 edges: HashMap<Eid, EdgeState>,
2446 vid: Vid,
2447 etype_id: u32,
2448 neighbor_is_dst: bool,
2449 visited: &HashSet<Vid>,
2450 next_frontier: &mut HashSet<Vid>,
2451 ) {
2452 for (eid, state) in edges {
2453 if state.deleted {
2454 continue;
2455 }
2456 graph.add_vertex(state.neighbor);
2457
2458 if !visited.contains(&state.neighbor) {
2459 next_frontier.insert(state.neighbor);
2460 }
2461
2462 if neighbor_is_dst {
2463 graph.add_edge(vid, state.neighbor, eid, etype_id);
2464 } else {
2465 graph.add_edge(state.neighbor, vid, eid, etype_id);
2466 }
2467 }
2468 }
2469}
2470
2471fn extract_vid_score_pairs(
2473 batches: &[arrow_array::RecordBatch],
2474 vid_column: &str,
2475 score_column: &str,
2476) -> Result<Vec<(Vid, f32)>> {
2477 let mut results = Vec::new();
2478 for batch in batches {
2479 let vid_col = batch
2480 .column_by_name(vid_column)
2481 .ok_or_else(|| anyhow!("Missing {} column", vid_column))?
2482 .as_any()
2483 .downcast_ref::<UInt64Array>()
2484 .ok_or_else(|| anyhow!("Invalid {} column type", vid_column))?;
2485
2486 let score_col = batch
2487 .column_by_name(score_column)
2488 .ok_or_else(|| anyhow!("Missing {} column", score_column))?
2489 .as_any()
2490 .downcast_ref::<Float32Array>()
2491 .ok_or_else(|| anyhow!("Invalid {} column type", score_column))?;
2492
2493 for i in 0..batch.num_rows() {
2494 results.push((Vid::from(vid_col.value(i)), score_col.value(i)));
2495 }
2496 }
2497 Ok(results)
2498}
2499
2500fn extract_vid_and_vector_pairs(
2509 batches: &[arrow_array::RecordBatch],
2510 vid_column: &str,
2511 vector_column: &str,
2512) -> Result<Vec<(Vid, Vec<f32>)>> {
2513 use arrow_array::{Array, FixedSizeListArray};
2514 let mut out = Vec::new();
2515 for batch in batches {
2516 let vid_col = batch
2517 .column_by_name(vid_column)
2518 .ok_or_else(|| anyhow!("Missing {} column", vid_column))?
2519 .as_any()
2520 .downcast_ref::<UInt64Array>()
2521 .ok_or_else(|| anyhow!("Invalid {} column type", vid_column))?;
2522
2523 let vec_col = batch
2524 .column_by_name(vector_column)
2525 .ok_or_else(|| anyhow!("Missing vector column {}", vector_column))?
2526 .as_any()
2527 .downcast_ref::<FixedSizeListArray>()
2528 .ok_or_else(|| anyhow!("Vector column {} is not FixedSizeList", vector_column))?;
2529
2530 for i in 0..batch.num_rows() {
2531 if vec_col.is_null(i) {
2532 continue;
2533 }
2534 let element = vec_col.value(i);
2535 let floats = element
2536 .as_any()
2537 .downcast_ref::<Float32Array>()
2538 .ok_or_else(|| {
2539 anyhow!("Vector column {} inner type is not Float32", vector_column)
2540 })?;
2541 let emb: Vec<f32> = (0..floats.len()).map(|j| floats.value(j)).collect();
2542 out.push((Vid::from(vid_col.value(i)), emb));
2543 }
2544 }
2545 Ok(out)
2546}
2547
2548fn extract_embedding_from_props(
2562 props: &uni_common::Properties,
2563 property: &str,
2564) -> Option<Vec<f32>> {
2565 Vec::<f32>::try_from(props.get(property)?).ok()
2566}
2567
2568fn merge_l0_into_vector_results(
2578 results: &mut Vec<(Vid, f32)>,
2579 ctx: &QueryContext,
2580 label: &str,
2581 property: &str,
2582 query: &[f32],
2583 k: usize,
2584 metric: &DistanceMetric,
2585) {
2586 let mut buffers: Vec<Arc<parking_lot::RwLock<L0Buffer>>> =
2588 ctx.pending_flush_l0s.iter().map(Arc::clone).collect();
2589 buffers.push(Arc::clone(&ctx.l0));
2590 if let Some(ref txn) = ctx.transaction_l0 {
2591 buffers.push(Arc::clone(txn));
2592 }
2593
2594 let mut l0_candidates: HashMap<Vid, f32> = HashMap::new();
2596 let mut tombstoned: HashSet<Vid> = HashSet::new();
2598
2599 for buf_arc in &buffers {
2600 let buf = buf_arc.read();
2601
2602 for &vid in &buf.vertex_tombstones {
2604 tombstoned.insert(vid);
2605 }
2606
2607 for (&vid, labels) in &buf.vertex_labels {
2609 if !labels.iter().any(|l| l == label) {
2610 continue;
2611 }
2612 if let Some(props) = buf.vertex_properties.get(&vid)
2613 && let Some(emb) = extract_embedding_from_props(props, property)
2614 {
2615 if emb.len() != query.len() {
2616 continue; }
2624 let dist = metric.compute_distance(&emb, query);
2625 l0_candidates.insert(vid, dist);
2627 tombstoned.remove(&vid);
2629 }
2630 }
2631 }
2632
2633 if l0_candidates.is_empty() && tombstoned.is_empty() {
2635 return;
2636 }
2637
2638 results.retain(|(vid, _)| !tombstoned.contains(vid));
2640
2641 for (vid, dist) in &l0_candidates {
2643 if tombstoned.contains(vid) {
2649 continue;
2650 }
2651 if let Some(existing) = results.iter_mut().find(|(v, _)| v == vid) {
2652 existing.1 = *dist;
2653 } else {
2654 results.push((*vid, *dist));
2655 }
2656 }
2657
2658 results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
2660 results.truncate(k);
2661}
2662
2663pub fn collect_l0_label_candidates(ctx: &QueryContext, label: &str) -> (Vec<Vid>, HashSet<Vid>) {
2677 let mut buffers: Vec<Arc<parking_lot::RwLock<L0Buffer>>> =
2679 ctx.pending_flush_l0s.iter().map(Arc::clone).collect();
2680 buffers.push(Arc::clone(&ctx.l0));
2681 if let Some(ref txn) = ctx.transaction_l0 {
2682 buffers.push(Arc::clone(txn));
2683 }
2684
2685 let mut live: HashSet<Vid> = HashSet::new();
2686 let mut tombstoned: HashSet<Vid> = HashSet::new();
2687
2688 for buf_arc in &buffers {
2689 let buf = buf_arc.read();
2690
2691 for &vid in &buf.vertex_tombstones {
2693 tombstoned.insert(vid);
2694 live.remove(&vid);
2695 }
2696
2697 for (&vid, labels) in &buf.vertex_labels {
2700 if !labels.iter().any(|l| l == label) {
2701 continue;
2702 }
2703 if buf.vertex_properties.contains_key(&vid) {
2704 live.insert(vid);
2705 tombstoned.remove(&vid);
2706 }
2707 }
2708 }
2709
2710 (live.into_iter().collect(), tombstoned)
2711}
2712
2713fn compute_text_relevance(query: &str, text: &str) -> f32 {
2718 let query_tokens: HashSet<String> =
2719 query.split_whitespace().map(|t| t.to_lowercase()).collect();
2720 if query_tokens.is_empty() {
2721 return 0.0;
2722 }
2723 let text_tokens: HashSet<String> = text.split_whitespace().map(|t| t.to_lowercase()).collect();
2724 let hits = query_tokens
2725 .iter()
2726 .filter(|t| text_tokens.contains(t.as_str()))
2727 .count();
2728 hits as f32 / query_tokens.len() as f32
2729}
2730
2731fn extract_text_from_props<'a>(
2733 props: &'a uni_common::Properties,
2734 property: &str,
2735) -> Option<&'a str> {
2736 props.get(property)?.as_str()
2737}
2738
2739fn merge_l0_into_fts_results(
2749 results: &mut Vec<(Vid, f32)>,
2750 ctx: &QueryContext,
2751 label: &str,
2752 property: &str,
2753 query: &str,
2754 k: usize,
2755) {
2756 let mut buffers: Vec<Arc<parking_lot::RwLock<L0Buffer>>> =
2758 ctx.pending_flush_l0s.iter().map(Arc::clone).collect();
2759 buffers.push(Arc::clone(&ctx.l0));
2760 if let Some(ref txn) = ctx.transaction_l0 {
2761 buffers.push(Arc::clone(txn));
2762 }
2763
2764 let mut l0_candidates: HashMap<Vid, f32> = HashMap::new();
2766 let mut tombstoned: HashSet<Vid> = HashSet::new();
2768
2769 for buf_arc in &buffers {
2770 let buf = buf_arc.read();
2771
2772 for &vid in &buf.vertex_tombstones {
2774 tombstoned.insert(vid);
2775 }
2776
2777 for (&vid, labels) in &buf.vertex_labels {
2779 if !labels.iter().any(|l| l == label) {
2780 continue;
2781 }
2782 if let Some(props) = buf.vertex_properties.get(&vid)
2783 && let Some(text) = extract_text_from_props(props, property)
2784 {
2785 let score = compute_text_relevance(query, text);
2786 if score > 0.0 {
2787 l0_candidates.insert(vid, score);
2789 }
2790 tombstoned.remove(&vid);
2792 }
2793 }
2794 }
2795
2796 if l0_candidates.is_empty() && tombstoned.is_empty() {
2798 return;
2799 }
2800
2801 results.retain(|(vid, _)| !tombstoned.contains(vid));
2803
2804 for (vid, score) in &l0_candidates {
2806 if tombstoned.contains(vid) {
2810 continue;
2811 }
2812 if let Some(existing) = results.iter_mut().find(|(v, _)| v == vid) {
2813 existing.1 = *score;
2814 } else {
2815 results.push((*vid, *score));
2816 }
2817 }
2818
2819 results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
2821 results.truncate(k);
2822}