1use crate::backend::StorageBackend;
5#[cfg(feature = "lance-backend")]
6use crate::backend::lance::LanceDbBackend;
7use crate::backend::table_names;
8use crate::backend::types::{FilterExpr, 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 requires_adjacency_revalidation: bool,
70 pub config: UniConfig,
71 pub compaction_status: Arc<Mutex<CompactionStatus>>,
72 pub flush_in_progress: std::sync::atomic::AtomicUsize,
77 pinned_snapshot: Option<SnapshotManifest>,
79 pinned_version_hwm: Option<u64>,
88 #[expect(
96 dead_code,
97 reason = "RAII guard: held solely so its Drop deregisters this view's pinned version. Never read."
98 )]
99 pin_guard: Option<crate::storage::adjacency_manager::PinGuard>,
100 fork_scope: Option<Arc<crate::fork::ForkScope>>,
108 backend: Arc<dyn StorageBackend>,
110 vid_labels_index: Arc<parking_lot::RwLock<crate::storage::vid_labels::VidLabelsIndex>>,
118 plugin_registry: Option<Arc<uni_plugin::PluginRegistry>>,
128}
129
130pub struct FlushInProgressGuard {
139 storage: Arc<StorageManager>,
140}
141
142impl FlushInProgressGuard {
143 pub fn new(storage: &Arc<StorageManager>) -> Self {
144 storage
145 .flush_in_progress
146 .fetch_add(1, std::sync::atomic::Ordering::AcqRel);
147 Self {
148 storage: storage.clone(),
149 }
150 }
151}
152
153impl Drop for FlushInProgressGuard {
154 fn drop(&mut self) {
155 self.storage
157 .flush_in_progress
158 .fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
159 }
160}
161
162fn is_lance_conflict(err: &anyhow::Error) -> bool {
168 let msg = err.to_string();
169 msg.contains("Incompatible transaction") || msg.contains("conflict")
170}
171
172async fn retry_on_lance_conflict<F, Fut>(mut op: F) -> anyhow::Result<()>
177where
178 F: FnMut() -> Fut,
179 Fut: std::future::Future<Output = anyhow::Result<()>>,
180{
181 for attempt in 0u32..10 {
182 match op().await {
183 Ok(()) => return Ok(()),
184 Err(e) => {
185 if !is_lance_conflict(&e) || attempt == 9 {
186 return Err(e);
187 }
188 let backoff_ms = 1u64 << attempt;
189 tokio::time::sleep(std::time::Duration::from_millis(backoff_ms)).await;
190 }
191 }
192 }
193 unreachable!("retry loop exits via Ok or Err")
194}
195
196fn combine_hwm_filter(hwm: Option<u64>, extra: Option<&FilterExpr>) -> Option<FilterExpr> {
201 match (hwm, extra) {
202 (Some(hwm), Some(f)) => Some(FilterExpr::all([
203 FilterExpr::version_at_most(hwm),
204 f.clone(),
205 ])),
206 (Some(hwm), None) => Some(FilterExpr::version_at_most(hwm)),
207 (None, Some(f)) => Some(f.clone()),
208 (None, None) => None,
209 }
210}
211
212async fn existing_columns(
218 backend: &dyn StorageBackend,
219 table_name: &str,
220 columns: &[&str],
221) -> Result<Option<Vec<String>>> {
222 let Some(table_schema) = backend.get_table_schema(table_name).await? else {
223 return Ok(None);
224 };
225 let table_field_names: HashSet<&str> = table_schema
226 .fields()
227 .iter()
228 .map(|f| f.name().as_str())
229 .collect();
230 Ok(Some(
231 columns
232 .iter()
233 .copied()
234 .filter(|c| table_field_names.contains(c))
235 .map(|s| s.to_string())
236 .collect(),
237 ))
238}
239
240async fn scan_concat(
249 backend: &dyn StorageBackend,
250 request: ScanRequest,
251) -> Result<Option<arrow_array::RecordBatch>> {
252 let batches = backend.scan(request).await?;
253 if batches.is_empty() {
254 Ok(None)
255 } else {
256 Ok(Some(arrow::compute::concat_batches(
257 &batches[0].schema(),
258 &batches,
259 )?))
260 }
261}
262
263pub async fn merge_insert_batch_with_lance_conflict_retry(
272 backend: &dyn crate::backend::StorageBackend,
273 table_name: &str,
274 batch: arrow_array::RecordBatch,
275 on: &[&str],
276) -> anyhow::Result<()> {
277 retry_on_lance_conflict(|| async {
283 let exists = backend.table_exists(table_name).await?;
284 if !exists {
285 anyhow::bail!(
286 "merge_insert target table '{}' does not exist (partial writes \
287 require the row to already be present; CREATE goes through Append)",
288 table_name
289 );
290 }
291 backend
292 .merge_insert(table_name, on, vec![batch.clone()])
293 .await
294 })
295 .await
296}
297
298pub async fn write_batch_with_lance_conflict_retry(
307 backend: &dyn crate::backend::StorageBackend,
308 table_name: &str,
309 batch: arrow_array::RecordBatch,
310) -> anyhow::Result<()> {
311 use crate::backend::types::WriteMode;
312 retry_on_lance_conflict(|| async {
313 let exists = backend.table_exists(table_name).await?;
314 if exists {
315 backend
316 .write(table_name, vec![batch.clone()], WriteMode::Append)
317 .await
318 } else {
319 backend.create_table(table_name, vec![batch.clone()]).await
320 }
321 })
322 .await
323}
324
325struct CompactionGuard {
327 status: Arc<Mutex<CompactionStatus>>,
328}
329
330impl CompactionGuard {
331 fn new(status: Arc<Mutex<CompactionStatus>>) -> Option<Self> {
332 let mut s = acquire_mutex(&status, "compaction_status").ok()?;
333 if s.compaction_in_progress {
334 return None;
335 }
336 s.compaction_in_progress = true;
337 Some(Self {
338 status: status.clone(),
339 })
340 }
341}
342
343impl Drop for CompactionGuard {
344 fn drop(&mut self) {
345 match uni_common::sync::acquire_mutex(&self.status, "compaction_status") {
348 Ok(mut s) => {
349 s.compaction_in_progress = false;
350 s.last_compaction = Some(std::time::SystemTime::now());
351 }
352 Err(e) => {
353 log::error!(
357 "CompactionGuard drop failed to acquire poisoned lock: {}. \
358 Compaction status may be inconsistent. Issue #18/#150",
359 e
360 );
361 }
362 }
363 }
364}
365
366impl StorageManager {
367 pub async fn new_with_backend(
369 base_uri: &str,
370 store: Arc<dyn ObjectStore>,
371 backend: Arc<dyn StorageBackend>,
372 schema_manager: Arc<SchemaManager>,
373 config: UniConfig,
374 ) -> Result<Self> {
375 let resilient_store: Arc<dyn ObjectStore> = Arc::new(ResilientObjectStore::new(
376 store,
377 config.object_store.clone(),
378 ));
379
380 let snapshot_manager = Arc::new(SnapshotManager::new(resilient_store.clone()));
381
382 Self::recover_all_staging_tables(backend.as_ref(), &schema_manager).await?;
384
385 let mut sm = Self {
386 requires_adjacency_revalidation: false,
387 base_uri: base_uri.to_string(),
388 store: resilient_store,
389 schema_manager,
390 snapshot_manager,
391 adjacency_manager: Arc::new(AdjacencyManager::new(config.cache_size)),
392 config,
393 compaction_status: Arc::new(Mutex::new(CompactionStatus::default())),
394 flush_in_progress: std::sync::atomic::AtomicUsize::new(0),
395 pinned_snapshot: None,
396 pinned_version_hwm: None,
397 pin_guard: None,
398 fork_scope: None,
399 backend,
400 vid_labels_index: Arc::new(parking_lot::RwLock::new(
401 crate::storage::vid_labels::VidLabelsIndex::new(),
402 )),
403 plugin_registry: None,
404 };
405
406 if let Err(e) = sm.rebuild_vid_labels_index().await {
410 warn!(
411 "Failed to rebuild VidLabelsIndex on startup: {}. Falling back to storage queries.",
412 e
413 );
414 }
415
416 Ok(sm)
417 }
418
419 #[cfg(feature = "lance-backend")]
421 pub async fn new(base_uri: &str, schema_manager: Arc<SchemaManager>) -> Result<Self> {
422 Self::new_with_config(base_uri, schema_manager, UniConfig::default()).await
423 }
424
425 #[cfg(feature = "lance-backend")]
427 pub async fn new_with_cache(
428 base_uri: &str,
429 schema_manager: Arc<SchemaManager>,
430 adjacency_cache_size: usize,
431 ) -> Result<Self> {
432 let config = UniConfig {
433 cache_size: adjacency_cache_size,
434 ..Default::default()
435 };
436 Self::new_with_config(base_uri, schema_manager, config).await
437 }
438
439 #[cfg(feature = "lance-backend")]
441 pub async fn new_with_config(
442 base_uri: &str,
443 schema_manager: Arc<SchemaManager>,
444 config: UniConfig,
445 ) -> Result<Self> {
446 let store = Self::build_store_from_uri(base_uri)?;
447 Self::new_with_store_and_config(base_uri, store, schema_manager, config).await
448 }
449
450 #[cfg(feature = "lance-backend")]
452 pub async fn new_with_store_and_config(
453 base_uri: &str,
454 store: Arc<dyn ObjectStore>,
455 schema_manager: Arc<SchemaManager>,
456 config: UniConfig,
457 ) -> Result<Self> {
458 Self::new_with_store_and_storage_options(base_uri, store, schema_manager, config, None)
459 .await
460 }
461
462 #[cfg(feature = "lance-backend")]
464 pub async fn new_with_store_and_storage_options(
465 base_uri: &str,
466 store: Arc<dyn ObjectStore>,
467 schema_manager: Arc<SchemaManager>,
468 config: UniConfig,
469 lancedb_storage_options: Option<HashMap<String, String>>,
470 ) -> Result<Self> {
471 let backend = Arc::new(LanceDbBackend::connect(base_uri, lancedb_storage_options).await?);
472 Self::new_with_backend(base_uri, store, backend, schema_manager, config).await
473 }
474
475 async fn recover_all_staging_tables(
480 backend: &dyn StorageBackend,
481 schema_manager: &SchemaManager,
482 ) -> Result<()> {
483 let schema = schema_manager.schema();
484
485 backend
487 .recover_staging(table_names::main_vertex_table_name())
488 .await?;
489 backend
490 .recover_staging(table_names::main_edge_table_name())
491 .await?;
492
493 for label in schema.labels.keys() {
495 let name = table_names::vertex_table_name(label);
496 backend.recover_staging(&name).await?;
497 }
498
499 for edge_type in schema.edge_types.keys() {
501 for direction in &["fwd", "bwd"] {
502 let delta_name = table_names::delta_table_name(edge_type, direction);
504 backend.recover_staging(&delta_name).await?;
505
506 for _label in schema.labels.keys() {
508 let adj_name = table_names::adjacency_table_name(edge_type, direction);
509 backend.recover_staging(&adj_name).await?;
510 }
511 }
512 }
513
514 Ok(())
515 }
516
517 #[cfg(feature = "lance-backend")]
518 fn build_store_from_uri(base_uri: &str) -> Result<Arc<dyn ObjectStore>> {
519 if base_uri.contains("://") {
520 let parsed = url::Url::parse(base_uri).map_err(|e| anyhow!("Invalid base URI: {e}"))?;
521 let (store, _path) = object_store::parse_url(&parsed)
522 .map_err(|e| anyhow!("Failed to parse object store URL: {e}"))?;
523 Ok(Arc::from(store))
524 } else {
525 std::fs::create_dir_all(base_uri)?;
527 Ok(Arc::new(LocalFileSystem::new_with_prefix(base_uri)?))
528 }
529 }
530
531 pub fn local_fs_root(&self) -> Option<std::path::PathBuf> {
537 if self.base_uri.contains("://") {
538 None
539 } else {
540 Some(std::path::PathBuf::from(&self.base_uri))
541 }
542 }
543
544 pub fn pinned(&self, snapshot: SnapshotManifest) -> Self {
545 Self {
553 requires_adjacency_revalidation: self.requires_adjacency_revalidation,
554 base_uri: self.base_uri.clone(),
555 store: self.store.clone(),
556 schema_manager: self.schema_manager.clone(),
557 snapshot_manager: self.snapshot_manager.clone(),
558 adjacency_manager: Arc::new(AdjacencyManager::new(self.adjacency_manager.max_bytes())),
562 config: self.config.clone(),
563 compaction_status: Arc::new(Mutex::new(CompactionStatus::default())),
564 flush_in_progress: std::sync::atomic::AtomicUsize::new(0),
565 pinned_snapshot: Some(snapshot),
566 pinned_version_hwm: None,
567 pin_guard: None,
568 fork_scope: self.fork_scope.clone(),
569 backend: self.backend.clone(),
570 vid_labels_index: Arc::new(parking_lot::RwLock::new(
576 self.vid_labels_index.read().clone(),
577 )),
578 plugin_registry: self.plugin_registry.clone(),
579 }
580 }
581
582 pub fn pinned_at_version(&self, hwm: u64) -> Self {
601 let pin_guard = Some(self.adjacency_manager.pinned_versions().pin(hwm));
604 Self {
605 requires_adjacency_revalidation: self.requires_adjacency_revalidation,
606 base_uri: self.base_uri.clone(),
607 store: self.store.clone(),
608 schema_manager: self.schema_manager.clone(),
609 snapshot_manager: self.snapshot_manager.clone(),
610 adjacency_manager: self.adjacency_manager.clone(),
611 config: self.config.clone(),
612 compaction_status: Arc::new(Mutex::new(CompactionStatus::default())),
613 flush_in_progress: std::sync::atomic::AtomicUsize::new(0),
614 pinned_snapshot: None,
615 pinned_version_hwm: Some(hwm),
616 pin_guard,
617 fork_scope: self.fork_scope.clone(),
618 backend: self.backend.clone(),
619 vid_labels_index: Arc::new(parking_lot::RwLock::new(
625 self.vid_labels_index.read().clone(),
626 )),
627 plugin_registry: self.plugin_registry.clone(),
628 }
629 }
630
631 pub fn at_fork(&self, scope: Arc<crate::fork::ForkScope>) -> Self {
646 self.at_fork_with_schema(scope, self.schema_manager.clone())
647 }
648
649 pub fn at_fork_with_schema(
657 &self,
658 scope: Arc<crate::fork::ForkScope>,
659 merged_schema: Arc<SchemaManager>,
660 ) -> Self {
661 debug_assert!(
662 self.pinned_snapshot.is_none(),
663 "forking a pinned StorageManager is unsupported in Phase 1"
664 );
665 let branched_backend: Arc<dyn StorageBackend> = Arc::new(
666 crate::backend::branched::BranchedBackend::new(self.backend.clone(), scope.clone()),
667 );
668 let snapshot_manager = Arc::new(SnapshotManager::new_for_fork(
673 self.store.clone(),
674 scope.fork_id(),
675 ));
676 Self {
677 requires_adjacency_revalidation: self.requires_adjacency_revalidation,
678 base_uri: self.base_uri.clone(),
679 store: self.store.clone(),
680 schema_manager: merged_schema,
681 snapshot_manager,
682 adjacency_manager: Arc::new(AdjacencyManager::new(self.adjacency_manager.max_bytes())),
683 config: self.config.clone(),
684 compaction_status: Arc::new(Mutex::new(CompactionStatus::default())),
685 flush_in_progress: std::sync::atomic::AtomicUsize::new(0),
686 pinned_snapshot: None,
687 pinned_version_hwm: None,
688 pin_guard: None,
689 fork_scope: Some(scope),
690 backend: branched_backend,
691 vid_labels_index: Arc::new(parking_lot::RwLock::new(
697 self.vid_labels_index.read().clone(),
698 )),
699 plugin_registry: self.plugin_registry.clone(),
700 }
701 }
702
703 pub fn plugin_registry(&self) -> Option<&Arc<uni_plugin::PluginRegistry>> {
708 self.plugin_registry.as_ref()
709 }
710
711 pub fn set_plugin_registry(&mut self, registry: Arc<uni_plugin::PluginRegistry>) {
717 self.plugin_registry = Some(registry);
718 }
719
720 pub fn fork_scope(&self) -> Option<&Arc<crate::fork::ForkScope>> {
722 self.fork_scope.as_ref()
723 }
724
725 #[must_use]
738 pub fn has_fork_index(
739 &self,
740 label: &str,
741 column: &str,
742 kind: crate::fork::ForkLocalIndexKind,
743 ) -> bool {
744 self.fork_scope
745 .as_ref()
746 .is_some_and(|s| s.has_fork_local_index(label, column, kind))
747 }
748
749 pub fn base_uri(&self) -> &str {
752 &self.base_uri
753 }
754
755 pub fn get_edge_version_by_id(&self, edge_type_id: u32) -> Option<u64> {
756 let schema = self.schema_manager.schema();
757 let name = schema.edge_type_name_by_id(edge_type_id)?;
758 self.pinned_snapshot
759 .as_ref()
760 .and_then(|s| s.edges.get(name).map(|es| es.lance_version))
761 .filter(|v| *v != 0)
766 }
767
768 pub fn version_high_water_mark(&self) -> Option<u64> {
776 self.pinned_snapshot
777 .as_ref()
778 .map(|s| s.version_high_water_mark)
779 .or(self.pinned_version_hwm)
780 }
781
782 pub fn snapshot_version_hwm(&self) -> Option<u64> {
794 self.pinned_snapshot
795 .as_ref()
796 .map(|s| s.version_high_water_mark)
797 }
798
799 pub fn apply_version_filter(&self, base_filter: FilterExpr) -> FilterExpr {
804 match self.version_high_water_mark() {
805 Some(hwm) => FilterExpr::all([base_filter, FilterExpr::version_at_most(hwm)]),
806 None => base_filter,
807 }
808 }
809
810 fn build_active_filter(user_filter: Option<&str>) -> FilterExpr {
816 match user_filter {
817 Some(expr) => {
818 FilterExpr::all([FilterExpr::Raw(expr.to_string()), FilterExpr::not_deleted()])
819 }
820 None => FilterExpr::not_deleted(),
821 }
822 }
823
824 pub fn store(&self) -> Arc<dyn ObjectStore> {
825 self.store.clone()
826 }
827
828 pub fn compaction_status(
834 &self,
835 ) -> Result<CompactionStatus, uni_common::sync::LockPoisonedError> {
836 let guard = uni_common::sync::acquire_mutex(&self.compaction_status, "compaction_status")?;
837 Ok(guard.clone())
838 }
839
840 pub async fn compact(&self) -> Result<CompactionStats> {
841 let start = std::time::Instant::now();
843 let schema = self.schema_manager.schema();
844 let mut files_compacted = 0;
845
846 for label in schema.labels.keys() {
847 let name = table_names::vertex_table_name(label);
848 if self.backend.table_exists(&name).await? {
849 self.backend.optimize_table(&name).await?;
850 files_compacted += 1;
851 self.backend.invalidate_cache(&name);
852 }
853 }
854
855 Ok(CompactionStats {
856 files_compacted,
857 bytes_before: 0,
858 bytes_after: 0,
859 duration: start.elapsed(),
860 crdt_merges: 0,
861 })
862 }
863
864 pub async fn compact_label(&self, label: &str) -> Result<CompactionStats> {
865 let _guard = CompactionGuard::new(self.compaction_status.clone())
866 .ok_or_else(|| anyhow!("Compaction already in progress"))?;
867
868 let start = std::time::Instant::now();
869 let name = table_names::vertex_table_name(label);
870
871 if self.backend.table_exists(&name).await? {
872 self.backend.optimize_table(&name).await?;
873 self.backend.invalidate_cache(&name);
874 }
875
876 Ok(CompactionStats {
877 files_compacted: 1,
878 bytes_before: 0,
879 bytes_after: 0,
880 duration: start.elapsed(),
881 crdt_merges: 0,
882 })
883 }
884
885 pub async fn compact_edge_type(&self, edge_type: &str) -> Result<CompactionStats> {
886 let _guard = CompactionGuard::new(self.compaction_status.clone())
887 .ok_or_else(|| anyhow!("Compaction already in progress"))?;
888
889 let start = std::time::Instant::now();
890 let mut files_compacted = 0;
891
892 for dir in ["fwd", "bwd"] {
893 let name = table_names::delta_table_name(edge_type, dir);
894 if self.backend.table_exists(&name).await? {
895 self.backend.optimize_table(&name).await?;
896 files_compacted += 1;
897 }
898 }
899
900 Ok(CompactionStats {
901 files_compacted,
902 bytes_before: 0,
903 bytes_after: 0,
904 duration: start.elapsed(),
905 crdt_merges: 0,
906 })
907 }
908
909 pub async fn wait_for_compaction(&self) -> Result<()> {
910 loop {
911 let in_progress = {
912 acquire_mutex(&self.compaction_status, "compaction_status")?.compaction_in_progress
913 };
914 if !in_progress {
915 return Ok(());
916 }
917 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
918 }
919 }
920
921 pub fn start_background_compaction(
922 self: Arc<Self>,
923 mut shutdown_rx: tokio::sync::broadcast::Receiver<()>,
924 ) -> tokio::task::JoinHandle<()> {
925 if !self.config.compaction.enabled {
926 return tokio::spawn(async {});
927 }
928
929 tokio::spawn(async move {
930 let start = tokio::time::Instant::now() + self.config.compaction.check_interval;
936 let mut interval =
937 tokio::time::interval_at(start, self.config.compaction.check_interval);
938
939 loop {
940 tokio::select! {
941 _ = interval.tick() => {
942 if let Err(e) = self.update_compaction_status().await {
943 log::error!("Failed to update compaction status: {}", e);
944 continue;
945 }
946
947 if let Some(task) = self.pick_compaction_task() {
948 log::info!("Triggering background compaction: {:?}", task);
949 if let Err(e) = Self::execute_compaction(Arc::clone(&self), task).await {
950 log::error!("Compaction failed: {}", e);
951 }
952 }
953 }
954 _ = shutdown_rx.recv() => {
955 log::info!("Background compaction shutting down");
956 let _ = self.wait_for_compaction().await;
957 break;
958 }
959 }
960 }
961 })
962 }
963
964 async fn update_compaction_status(&self) -> Result<()> {
965 let schema = self.schema_manager.schema();
966 let backend = self.backend.as_ref();
967 let mut total_rows: usize = 0;
968 let mut oldest_ts: Option<i64> = None;
969
970 for name in schema.edge_types.keys() {
971 for dir in ["fwd", "bwd"] {
972 let tbl_name = table_names::delta_table_name(name, dir);
973 if !backend.table_exists(&tbl_name).await? {
974 continue;
975 }
976 let row_count = backend.count_rows(&tbl_name, None).await.unwrap_or(0);
977 if row_count == 0 {
978 continue;
979 }
980 total_rows += row_count;
981
982 let request =
984 ScanRequest::all(&tbl_name).with_columns(vec!["_created_at".to_string()]);
985 let Ok(batches) = backend.scan(request).await else {
986 continue;
987 };
988 for batch in batches {
989 let Some(col) = batch
990 .column_by_name("_created_at")
991 .and_then(|c| c.as_any().downcast_ref::<TimestampNanosecondArray>())
992 else {
993 continue;
994 };
995 for i in 0..col.len() {
996 if !col.is_null(i) {
997 let ts = col.value(i);
998 oldest_ts = Some(oldest_ts.map_or(ts, |prev| prev.min(ts)));
999 }
1000 }
1001 }
1002 }
1003 }
1004
1005 let oldest_l1_age = oldest_ts
1006 .and_then(|ts| {
1007 let created = UNIX_EPOCH + Duration::from_nanos(ts as u64);
1008 SystemTime::now().duration_since(created).ok()
1009 })
1010 .unwrap_or(Duration::ZERO);
1011
1012 let mut status = acquire_mutex(&self.compaction_status, "compaction_status")?;
1013 status.l1_size_bytes = (total_rows * ENTRY_SIZE_ESTIMATE) as u64;
1016 status.oldest_l1_age = oldest_l1_age;
1017 Ok(())
1018 }
1019
1020 fn pick_compaction_task(&self) -> Option<CompactionTask> {
1021 let status = acquire_mutex(&self.compaction_status, "compaction_status").ok()?;
1022
1023 if status.l1_runs >= self.config.compaction.max_l1_runs {
1024 return Some(CompactionTask::ByRunCount);
1025 }
1026 if status.l1_size_bytes >= self.config.compaction.max_l1_size_bytes {
1027 return Some(CompactionTask::BySize);
1028 }
1029 if status.oldest_l1_age >= self.config.compaction.max_l1_age
1030 && status.oldest_l1_age > Duration::ZERO
1031 {
1032 return Some(CompactionTask::ByAge);
1033 }
1034
1035 None
1036 }
1037
1038 async fn try_optimize_table(backend: &dyn StorageBackend, table_name: &str) -> bool {
1040 if let Err(e) = backend.optimize_table(table_name).await {
1041 log::warn!("Failed to optimize table {}: {}", table_name, e);
1042 return false;
1043 }
1044 true
1045 }
1046
1047 pub fn trigger_async_compaction(self: &Arc<Self>) {
1050 let this = Arc::clone(self);
1051 tokio::spawn(async move {
1052 if let Err(e) = Self::execute_compaction(this, CompactionTask::ByRunCount).await {
1053 log::debug!("Post-flush compaction skipped: {}", e);
1055 }
1056 });
1057 }
1058
1059 pub(crate) async fn execute_compaction(
1060 this: Arc<Self>,
1061 _task: CompactionTask,
1062 ) -> Result<CompactionStats> {
1063 let start = std::time::Instant::now();
1064 let _guard = CompactionGuard::new(this.compaction_status.clone())
1065 .ok_or_else(|| anyhow!("Compaction already in progress"))?;
1066
1067 let schema = this.schema_manager.schema();
1068 let mut files_compacted = 0;
1069
1070 let compactor = Compactor::new(Arc::clone(&this));
1073 let compaction_results = compactor.compact_all().await.unwrap_or_else(|e| {
1074 log::error!(
1075 "Semantic compaction failed (continuing with backend optimize): {}",
1076 e
1077 );
1078 Vec::new()
1079 });
1080
1081 let am = this.adjacency_manager();
1083 for info in &compaction_results {
1084 let direction = match info.direction.as_str() {
1085 "fwd" => Direction::Outgoing,
1086 "bwd" => Direction::Incoming,
1087 _ => continue,
1088 };
1089 if let Some(etid) = schema.edge_type_id_unified_case_insensitive(&info.edge_type)
1090 && let Err(e) = am.warm(&this, etid, direction, None).await
1091 {
1092 log::warn!(
1093 "Failed to re-warm adjacency for {}/{}: {}",
1094 info.edge_type,
1095 info.direction,
1096 e
1097 );
1098 }
1099 }
1100
1101 let backend = this.backend.as_ref();
1103
1104 for name in schema.edge_types.keys() {
1106 for dir in ["fwd", "bwd"] {
1107 let delta = table_names::delta_table_name(name, dir);
1108 if Self::try_optimize_table(backend, &delta).await {
1109 files_compacted += 1;
1110 }
1111 let adj = table_names::adjacency_table_name(name, dir);
1112 if Self::try_optimize_table(backend, &adj).await {
1113 files_compacted += 1;
1114 }
1115 }
1116 }
1117
1118 for label in schema.labels.keys() {
1120 let tbl = table_names::vertex_table_name(label);
1121 if Self::try_optimize_table(backend, &tbl).await {
1122 files_compacted += 1;
1123 backend.invalidate_cache(&tbl);
1124 }
1125 }
1126
1127 for tbl in [
1129 table_names::main_vertex_table_name(),
1130 table_names::main_edge_table_name(),
1131 ] {
1132 if Self::try_optimize_table(backend, tbl).await {
1133 files_compacted += 1;
1134 }
1135 }
1136
1137 {
1138 let mut status = acquire_mutex(&this.compaction_status, "compaction_status")?;
1139 status.total_compactions += 1;
1140 status.l1_runs = 0; }
1142
1143 Ok(CompactionStats {
1144 files_compacted,
1145 bytes_before: 0,
1146 bytes_after: 0,
1147 duration: start.elapsed(),
1148 crdt_merges: 0,
1149 })
1150 }
1151
1152 pub fn invalidate_table_cache(&self, label: &str) {
1156 let name = table_names::vertex_table_name(label);
1157 self.backend.invalidate_cache(&name);
1158 }
1159
1160 pub fn base_path(&self) -> &str {
1161 &self.base_uri
1162 }
1163
1164 pub fn schema_manager(&self) -> &SchemaManager {
1165 &self.schema_manager
1166 }
1167
1168 pub fn schema_manager_arc(&self) -> Arc<SchemaManager> {
1169 self.schema_manager.clone()
1170 }
1171
1172 #[must_use]
1180 pub fn schema_manager_arc_ref(&self) -> &Arc<SchemaManager> {
1181 &self.schema_manager
1182 }
1183
1184 pub fn adjacency_manager(&self) -> Arc<AdjacencyManager> {
1186 Arc::clone(&self.adjacency_manager)
1187 }
1188
1189 pub async fn warm_adjacency(
1194 &self,
1195 edge_type_id: u32,
1196 direction: crate::storage::direction::Direction,
1197 version: Option<u64>,
1198 ) -> anyhow::Result<()> {
1199 self.adjacency_manager
1200 .warm(self, edge_type_id, direction, version)
1201 .await
1202 }
1203
1204 pub async fn warm_adjacency_coalesced(
1209 &self,
1210 edge_type_id: u32,
1211 direction: crate::storage::direction::Direction,
1212 version: Option<u64>,
1213 ) -> anyhow::Result<()> {
1214 self.adjacency_manager
1215 .warm_coalesced(self, edge_type_id, direction, version)
1216 .await
1217 }
1218
1219 pub fn set_requires_adjacency_revalidation(&mut self, yes: bool) {
1229 self.requires_adjacency_revalidation = yes;
1230 }
1231
1232 #[must_use]
1234 pub fn requires_adjacency_revalidation(&self) -> bool {
1235 self.requires_adjacency_revalidation
1236 }
1237
1238 pub async fn adjacency_csr_is_fresh(
1243 &self,
1244 edge_type_id: u32,
1245 direction: crate::storage::direction::Direction,
1246 ) -> bool {
1247 self.adjacency_manager
1248 .csr_is_fresh(self, edge_type_id, direction)
1249 .await
1250 }
1251
1252 pub fn invalidate_adjacency_csr(
1254 &self,
1255 edge_type_id: u32,
1256 direction: crate::storage::direction::Direction,
1257 ) {
1258 self.adjacency_manager
1259 .invalidate_csr(edge_type_id, direction);
1260 }
1261
1262 pub fn has_adjacency_csr(
1264 &self,
1265 edge_type_id: u32,
1266 direction: crate::storage::direction::Direction,
1267 ) -> bool {
1268 self.adjacency_manager.has_csr(edge_type_id, direction)
1269 }
1270
1271 pub fn get_neighbors_at_version(
1273 &self,
1274 vid: uni_common::core::id::Vid,
1275 edge_type: u32,
1276 direction: crate::storage::direction::Direction,
1277 version: u64,
1278 ) -> Vec<(uni_common::core::id::Vid, uni_common::core::id::Eid)> {
1279 self.adjacency_manager
1280 .get_neighbors_at_version(vid, edge_type, direction, version)
1281 }
1282
1283 pub fn backend(&self) -> &dyn StorageBackend {
1285 self.backend.as_ref()
1286 }
1287
1288 pub fn backend_arc(&self) -> Arc<dyn StorageBackend> {
1290 self.backend.clone()
1291 }
1292
1293 async fn rebuild_vid_labels_index(&mut self) -> Result<()> {
1298 use crate::storage::vid_labels::VidLabelsIndex;
1299
1300 let backend = self.backend.as_ref();
1301 let vtable = table_names::main_vertex_table_name();
1302
1303 if !backend.table_exists(vtable).await? {
1305 self.vid_labels_index = Arc::new(parking_lot::RwLock::new(VidLabelsIndex::new()));
1306 return Ok(());
1307 }
1308
1309 let request = ScanRequest::all(vtable)
1311 .with_filter(FilterExpr::not_deleted())
1312 .with_limit(100_000);
1313 let batches = backend
1314 .scan(request)
1315 .await
1316 .map_err(|e| anyhow!("Failed to query main vertex table: {}", e))?;
1317
1318 let mut index = VidLabelsIndex::new();
1319 for batch in batches {
1320 let vid_col = batch
1321 .column_by_name("_vid")
1322 .ok_or_else(|| anyhow!("Missing _vid column"))?
1323 .as_any()
1324 .downcast_ref::<UInt64Array>()
1325 .ok_or_else(|| anyhow!("Invalid _vid column type"))?;
1326
1327 let labels_col = batch
1328 .column_by_name("labels")
1329 .ok_or_else(|| anyhow!("Missing labels column"))?
1330 .as_any()
1331 .downcast_ref::<arrow_array::ListArray>()
1332 .ok_or_else(|| anyhow!("Invalid labels column type"))?;
1333
1334 for row_idx in 0..batch.num_rows() {
1335 let vid = Vid::from(vid_col.value(row_idx));
1336 let labels_array = labels_col.value(row_idx);
1337 let labels_str_array = labels_array
1338 .as_any()
1339 .downcast_ref::<arrow_array::StringArray>()
1340 .ok_or_else(|| anyhow!("Invalid labels array element type"))?;
1341
1342 let labels: Vec<String> = (0..labels_str_array.len())
1343 .map(|i| labels_str_array.value(i).to_string())
1344 .collect();
1345
1346 index.insert(vid, labels);
1347 }
1348 }
1349
1350 self.vid_labels_index = Arc::new(parking_lot::RwLock::new(index));
1351 Ok(())
1352 }
1353
1354 pub fn get_labels_from_index(&self, vid: Vid) -> Option<Vec<String>> {
1359 let index = self.vid_labels_index.read();
1360 index.get_labels(vid).map(|labels| labels.to_vec())
1361 }
1362
1363 pub fn physical_vertex_label_names(&self) -> Vec<String> {
1370 self.vid_labels_index
1371 .read()
1372 .labels()
1373 .map(str::to_owned)
1374 .collect()
1375 }
1376
1377 pub fn update_vid_labels_index(&self, vid: Vid, labels: Vec<String>) {
1379 let mut index = self.vid_labels_index.write();
1380 index.insert(vid, labels);
1381 }
1382
1383 pub fn remove_from_vid_labels_index(&self, vid: Vid) {
1385 let mut index = self.vid_labels_index.write();
1386 index.remove_vid(vid);
1387 }
1388
1389 pub async fn load_subgraph_cached(
1390 &self,
1391 start_vids: &[Vid],
1392 edge_types: &[u32],
1393 max_hops: usize,
1394 direction: GraphDirection,
1395 _l0: Option<Arc<RwLock<L0Buffer>>>,
1396 ) -> Result<WorkingGraph> {
1397 let mut graph = WorkingGraph::new();
1398
1399 let dir = match direction {
1400 GraphDirection::Outgoing => crate::storage::direction::Direction::Outgoing,
1401 GraphDirection::Incoming => crate::storage::direction::Direction::Incoming,
1402 };
1403
1404 let neighbor_is_dst = matches!(direction, GraphDirection::Outgoing);
1405
1406 let mut frontier: Vec<Vid> = start_vids.to_vec();
1408 let mut visited: HashSet<Vid> = HashSet::new();
1409
1410 for &vid in start_vids {
1412 graph.add_vertex(vid);
1413 }
1414
1415 for _hop in 0..max_hops {
1416 let mut next_frontier = HashSet::new();
1417
1418 for &vid in &frontier {
1419 if visited.contains(&vid) {
1420 continue;
1421 }
1422 visited.insert(vid);
1423 graph.add_vertex(vid);
1424
1425 for &etype_id in edge_types {
1426 let edge_ver = self.snapshot_version_hwm();
1430 self.adjacency_manager
1431 .warm_coalesced(self, etype_id, dir, edge_ver)
1432 .await?;
1433
1434 let edges = self.adjacency_manager.get_neighbors(vid, etype_id, dir);
1436
1437 for (neighbor_vid, eid) in edges {
1438 graph.add_vertex(neighbor_vid);
1439 if !visited.contains(&neighbor_vid) {
1440 next_frontier.insert(neighbor_vid);
1441 }
1442
1443 if neighbor_is_dst {
1444 graph.add_edge(vid, neighbor_vid, eid, etype_id);
1445 } else {
1446 graph.add_edge(neighbor_vid, vid, eid, etype_id);
1447 }
1448 }
1449 }
1450 }
1451 frontier = next_frontier.into_iter().collect();
1452
1453 if frontier.is_empty() {
1455 break;
1456 }
1457 }
1458
1459 Ok(graph)
1460 }
1461
1462 pub fn snapshot_manager(&self) -> &SnapshotManager {
1463 &self.snapshot_manager
1464 }
1465
1466 pub fn index_manager(&self) -> IndexManager {
1467 IndexManager::new(&self.base_uri, self.schema_manager.clone())
1468 .with_backend(self.backend_arc())
1469 }
1470
1471 pub async fn scan_vertex_table(
1480 &self,
1481 label: &str,
1482 columns: &[&str],
1483 additional_filter: Option<&FilterExpr>,
1484 ) -> Result<Option<arrow_array::RecordBatch>> {
1485 let backend = self.backend();
1486 let table_name = table_names::vertex_table_name(label);
1487
1488 if !backend.table_exists(&table_name).await? {
1489 return Ok(None);
1490 }
1491
1492 let Some(actual_columns) = existing_columns(backend, &table_name, columns).await? else {
1494 return Ok(None);
1495 };
1496
1497 let filter = combine_hwm_filter(self.version_high_water_mark(), additional_filter);
1499
1500 let mut request = ScanRequest::all(&table_name).with_columns(actual_columns);
1501 if let Some(f) = filter {
1502 request = request.with_filter(f);
1503 }
1504
1505 scan_concat(backend, request).await
1506 }
1507
1508 pub async fn scan_delta_table(
1511 &self,
1512 edge_type: &str,
1513 direction: &str,
1514 columns: &[&str],
1515 additional_filter: Option<&FilterExpr>,
1516 ) -> Result<Option<arrow_array::RecordBatch>> {
1517 let edge_hwm = self.snapshot_version_hwm();
1523 let backend = self.backend();
1524 let table_name = table_names::delta_table_name(edge_type, direction);
1525
1526 if !backend.table_exists(&table_name).await? {
1527 return Ok(None);
1528 }
1529
1530 let Some(actual_columns) = existing_columns(backend, &table_name, columns).await? else {
1532 return Ok(None);
1533 };
1534
1535 let filter = combine_hwm_filter(edge_hwm, additional_filter);
1536
1537 let mut request = ScanRequest::all(&table_name).with_columns(actual_columns);
1538 if let Some(f) = filter {
1539 request = request.with_filter(f);
1540 }
1541
1542 scan_concat(backend, request).await
1543 }
1544
1545 pub async fn scan_main_vertex_table(
1550 &self,
1551 columns: &[&str],
1552 filter: Option<&FilterExpr>,
1553 ) -> Result<Option<arrow_array::RecordBatch>> {
1554 let backend = self.backend();
1555 let table_name = table_names::main_vertex_table_name();
1556
1557 if !backend.table_exists(table_name).await? {
1558 return Ok(None);
1559 }
1560
1561 let full_filter = combine_hwm_filter(self.version_high_water_mark(), filter);
1563
1564 let request = ScanRequest::all(table_name)
1565 .with_columns(columns.iter().map(|s| s.to_string()).collect());
1566 let request = match full_filter {
1567 Some(f) => request.with_filter(f),
1568 None => request,
1569 };
1570
1571 scan_concat(backend, request).await
1572 }
1573
1574 pub async fn scan_main_edge_table_stream(
1576 &self,
1577 filter: Option<&FilterExpr>,
1578 ) -> Result<
1579 Option<
1580 std::pin::Pin<Box<dyn futures::Stream<Item = Result<arrow_array::RecordBatch>> + Send>>,
1581 >,
1582 > {
1583 let backend = self.backend();
1584 let table_name = table_names::main_edge_table_name();
1585
1586 if !backend.table_exists(table_name).await? {
1587 return Ok(None);
1588 }
1589
1590 let mut request = ScanRequest::all(table_name);
1591 if let Some(f) = filter {
1592 request = request.with_filter(f.clone());
1593 }
1594
1595 let stream = backend.scan_stream(request).await?;
1596 Ok(Some(stream))
1597 }
1598
1599 pub async fn scan_vertex_table_stream(
1601 &self,
1602 label: &str,
1603 ) -> Result<
1604 Option<
1605 std::pin::Pin<Box<dyn futures::Stream<Item = Result<arrow_array::RecordBatch>> + Send>>,
1606 >,
1607 > {
1608 let backend = self.backend();
1609 let table_name = table_names::vertex_table_name(label);
1610
1611 if !backend.table_exists(&table_name).await? {
1612 return Ok(None);
1613 }
1614
1615 let stream = backend.scan_stream(ScanRequest::all(&table_name)).await?;
1616 Ok(Some(stream))
1617 }
1618
1619 pub async fn find_vertex_by_ext_id(&self, ext_id: &str) -> Result<Option<Vid>> {
1621 MainVertexDataset::find_by_ext_id(self.backend(), ext_id, self.version_high_water_mark())
1622 .await
1623 }
1624
1625 pub async fn get_vertex_ext_ids(&self) -> Result<std::collections::HashMap<Vid, String>> {
1637 use arrow_array::StringArray;
1638 let backend = self.backend.as_ref();
1639 let vtable = table_names::main_vertex_table_name();
1640 let mut out = std::collections::HashMap::new();
1641 if !backend.table_exists(vtable).await? {
1642 return Ok(out);
1643 }
1644 let request = ScanRequest::all(vtable)
1645 .with_filter(FilterExpr::not_deleted())
1646 .with_columns(vec!["_vid".to_string(), "ext_id".to_string()]);
1647 let batches = backend
1648 .scan(request)
1649 .await
1650 .map_err(|e| anyhow!("get_vertex_ext_ids: {}", e))?;
1651 for batch in batches {
1652 let vids = batch
1653 .column_by_name("_vid")
1654 .and_then(|c| c.as_any().downcast_ref::<UInt64Array>())
1655 .ok_or_else(|| anyhow!("get_vertex_ext_ids: missing/invalid _vid column"))?;
1656 let exts = batch
1657 .column_by_name("ext_id")
1658 .and_then(|c| c.as_any().downcast_ref::<StringArray>())
1659 .ok_or_else(|| anyhow!("get_vertex_ext_ids: missing/invalid ext_id column"))?;
1660 for i in 0..batch.num_rows() {
1661 if exts.is_null(i) {
1662 continue;
1663 }
1664 let ext = exts.value(i);
1665 if !ext.is_empty() {
1666 out.insert(Vid::from(vids.value(i)), ext.to_string());
1667 }
1668 }
1669 }
1670 Ok(out)
1671 }
1672
1673 pub async fn find_vertex_labels_by_vid(&self, vid: Vid) -> Result<Option<Vec<String>>> {
1675 MainVertexDataset::find_labels_by_vid(self.backend(), vid, self.version_high_water_mark())
1676 .await
1677 }
1678
1679 pub async fn find_edges_by_type_names(
1682 &self,
1683 type_names: &[&str],
1684 endpoint_filter: Option<(crate::storage::main_edge::EndpointSide, &[Vid])>,
1685 ) -> Result<Vec<(Eid, Vid, Vid, String, uni_common::Properties)>> {
1686 MainEdgeDataset::find_edges_by_type_names(self.backend(), type_names, endpoint_filter).await
1687 }
1688
1689 pub async fn scan_vertex_candidates(
1691 &self,
1692 label: &str,
1693 filter: Option<&FilterExpr>,
1694 ) -> Result<Vec<Vid>> {
1695 let backend = self.backend();
1696 let table_name = table_names::vertex_table_name(label);
1697
1698 if !backend.table_exists(&table_name).await? {
1699 return Ok(Vec::new());
1700 }
1701
1702 let full_filter = match filter {
1703 Some(f) => FilterExpr::all([FilterExpr::not_deleted(), f.clone()]),
1704 None => FilterExpr::not_deleted(),
1705 };
1706
1707 let request = ScanRequest::all(&table_name)
1708 .with_filter(full_filter)
1709 .with_columns(vec!["_vid".to_string()]);
1710
1711 let batches = backend.scan(request).await?;
1712
1713 let mut vids = Vec::new();
1714 for batch in batches {
1715 let vid_col = batch
1716 .column_by_name("_vid")
1717 .ok_or(anyhow!("Missing _vid"))?
1718 .as_any()
1719 .downcast_ref::<UInt64Array>()
1720 .ok_or(anyhow!("Invalid _vid"))?;
1721 for i in 0..batch.num_rows() {
1722 vids.push(Vid::from(vid_col.value(i)));
1723 }
1724 }
1725 Ok(vids)
1726 }
1727
1728 pub fn vertex_dataset(&self, label: &str) -> Result<VertexDataset> {
1734 let schema = self.schema_manager.schema();
1735 let label_meta = schema
1736 .labels
1737 .get(label)
1738 .ok_or_else(|| anyhow!("Label '{}' not found", label))?;
1739 Ok(VertexDataset::new(&self.base_uri, label, label_meta.id))
1740 }
1741
1742 #[cfg(feature = "lance-backend")]
1743 pub fn edge_dataset(
1744 &self,
1745 edge_type: &str,
1746 src_label: &str,
1747 dst_label: &str,
1748 ) -> Result<EdgeDataset> {
1749 let key = format!("edges_{edge_type}");
1750 match self.fork_branch_for(&key) {
1751 Some(branch) => Ok(EdgeDataset::new_branched(
1752 &self.base_uri,
1753 edge_type,
1754 src_label,
1755 dst_label,
1756 branch,
1757 )),
1758 None => Ok(EdgeDataset::new(
1759 &self.base_uri,
1760 edge_type,
1761 src_label,
1762 dst_label,
1763 )),
1764 }
1765 }
1766
1767 pub fn delta_dataset(&self, edge_type: &str, direction: &str) -> Result<DeltaDataset> {
1768 let key = format!("deltas_{edge_type}_{direction}");
1769 match self.fork_branch_for(&key) {
1770 Some(branch) => Ok(DeltaDataset::new_branched(
1771 &self.base_uri,
1772 edge_type,
1773 direction,
1774 branch,
1775 )),
1776 None => Ok(DeltaDataset::new(&self.base_uri, edge_type, direction)),
1777 }
1778 }
1779
1780 pub fn adjacency_dataset(
1781 &self,
1782 edge_type: &str,
1783 label: &str,
1784 direction: &str,
1785 ) -> Result<AdjacencyDataset> {
1786 let key = crate::backend::table_names::adjacency_table_name(edge_type, direction);
1793 match self.fork_branch_for(&key) {
1794 Some(branch) => Ok(AdjacencyDataset::new_branched(
1795 &self.base_uri,
1796 edge_type,
1797 label,
1798 direction,
1799 branch,
1800 )),
1801 None => Ok(AdjacencyDataset::new(
1802 &self.base_uri,
1803 edge_type,
1804 label,
1805 direction,
1806 )),
1807 }
1808 }
1809
1810 fn fork_branch_for(&self, dataset_name: &str) -> Option<String> {
1814 self.fork_scope
1815 .as_ref()
1816 .and_then(|s| s.branch_for(dataset_name))
1817 }
1818
1819 pub fn main_vertex_dataset(&self) -> MainVertexDataset {
1824 MainVertexDataset::new(&self.base_uri)
1825 }
1826
1827 pub fn main_edge_dataset(&self) -> MainEdgeDataset {
1832 MainEdgeDataset::new(&self.base_uri)
1833 }
1834
1835 #[cfg(feature = "lance-backend")]
1836 pub fn uid_index(&self, label: &str) -> Result<UidIndex> {
1837 Ok(UidIndex::new(&self.base_uri, label))
1838 }
1839
1840 #[cfg(feature = "lance-backend")]
1841 pub async fn inverted_index(&self, label: &str, property: &str) -> Result<InvertedIndex> {
1842 let schema = self.schema_manager.schema();
1843 let config = schema
1844 .indexes
1845 .iter()
1846 .find_map(|idx| match idx {
1847 IndexDefinition::Inverted(cfg)
1848 if cfg.label == label && cfg.property == property =>
1849 {
1850 Some(cfg.clone())
1851 }
1852 _ => None,
1853 })
1854 .ok_or_else(|| anyhow!("Inverted index not found for {}.{}", label, property))?;
1855
1856 InvertedIndex::new(&self.base_uri, config).await
1857 }
1858
1859 #[expect(clippy::too_many_arguments)]
1860 pub async fn vector_search(
1861 &self,
1862 label: &str,
1863 property: &str,
1864 query: &[f32],
1865 k: usize,
1866 filter: Option<&str>,
1867 opts: VectorQueryOpts,
1868 ctx: Option<&QueryContext>,
1869 ) -> Result<Vec<(Vid, f32)>> {
1870 use crate::backend::types::DistanceMetric as BackendMetric;
1871
1872 let schema = self.schema_manager.schema();
1874
1875 if let Some(meta) = schema
1880 .properties
1881 .get(label)
1882 .and_then(|props| props.get(property))
1883 && let uni_common::core::schema::DataType::Vector { dimensions } = &meta.r#type
1884 && query.len() != *dimensions
1885 {
1886 return Err(anyhow!(
1887 "vector dimension mismatch: query vector has {} dimensions but '{}.{}' is \
1888 declared VECTOR({})",
1889 query.len(),
1890 label,
1891 property,
1892 dimensions
1893 ));
1894 }
1895
1896 let metric = schema
1897 .vector_index_for_property(label, property)
1898 .map(|config| config.metric.clone())
1899 .unwrap_or(DistanceMetric::L2);
1900
1901 let backend = self.backend.as_ref();
1902 let name = table_names::vertex_table_name(label);
1903
1904 let mut results = Vec::new();
1905
1906 if backend.table_exists(&name).await? {
1908 let backend_metric = match &metric {
1909 DistanceMetric::L2 => BackendMetric::L2,
1910 DistanceMetric::Cosine => BackendMetric::Cosine,
1911 DistanceMetric::Dot => BackendMetric::Dot,
1912 DistanceMetric::L1 => BackendMetric::L2,
1917 _ => BackendMetric::L2,
1918 };
1919
1920 let mut filter_parts = vec![Self::build_active_filter(filter)];
1922 if ctx.is_some()
1923 && let Some(hwm) = self.version_high_water_mark()
1924 {
1925 filter_parts.push(FilterExpr::version_at_most(hwm));
1926 }
1927 let combined_filter = FilterExpr::all(filter_parts);
1928
1929 let fetch_k = if matches!(metric, DistanceMetric::L1) {
1933 backend.count_rows(&name, None).await.unwrap_or(k).max(k)
1934 } else {
1935 k
1936 };
1937
1938 let batches = backend
1939 .vector_search(
1940 &name,
1941 property,
1942 query,
1943 fetch_k,
1944 backend_metric,
1945 combined_filter,
1946 opts,
1947 )
1948 .await?;
1949
1950 results = extract_vid_and_vector_pairs(&batches, "_vid", property)?
1959 .into_iter()
1960 .filter(|(_, emb)| emb.len() == query.len())
1961 .map(|(vid, emb)| (vid, metric.compute_distance(&emb, query)))
1962 .collect();
1963 }
1964
1965 if let Some(qctx) = ctx {
1967 merge_l0_into_vector_results(&mut results, qctx, label, property, query, k, &metric);
1968 }
1969
1970 results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
1975 results.truncate(k);
1976
1977 Ok(results)
1978 }
1979
1980 #[expect(clippy::too_many_arguments)]
1991 pub async fn muvera_fde_candidates(
1992 &self,
1993 label: &str,
1994 fde_column: &str,
1995 fde_query: &[f32],
1996 k: usize,
1997 filter: Option<&str>,
1998 opts: VectorQueryOpts,
1999 ctx: Option<&QueryContext>,
2000 ) -> Result<Vec<(Vid, f32)>> {
2001 use crate::backend::types::{DistanceMetric as BackendMetric, FilterExpr};
2002
2003 let backend = self.backend.as_ref();
2004 let name = table_names::vertex_table_name(label);
2005 if !backend.table_exists(&name).await? {
2009 return Ok(Vec::new());
2010 }
2011
2012 let mut filter_parts = vec![Self::build_active_filter(filter)];
2013 if ctx.is_some()
2014 && let Some(hwm) = self.version_high_water_mark()
2015 {
2016 filter_parts.push(FilterExpr::version_at_most(hwm));
2017 }
2018 let combined_filter = FilterExpr::all(filter_parts);
2019
2020 let batches = backend
2021 .vector_search(
2022 &name,
2023 fde_column,
2024 fde_query,
2025 k,
2026 BackendMetric::Dot,
2027 combined_filter,
2028 opts,
2029 )
2030 .await?;
2031 extract_vid_score_pairs(&batches, "_vid", "_distance")
2032 }
2033
2034 #[expect(clippy::too_many_arguments)]
2055 pub async fn multivector_search(
2056 &self,
2057 label: &str,
2058 property: &str,
2059 query: &[Vec<f32>],
2060 k: usize,
2061 filter: Option<&str>,
2062 opts: VectorQueryOpts,
2063 ctx: Option<&QueryContext>,
2064 ) -> Result<Vec<(Vid, f32)>> {
2065 use crate::backend::types::{DistanceMetric as BackendMetric, FilterExpr};
2066
2067 let schema = self.schema_manager.schema();
2068 let metric = schema
2069 .vector_index_for_property(label, property)
2070 .map(|config| config.metric.clone())
2071 .unwrap_or(DistanceMetric::Cosine);
2072
2073 let backend = self.backend.as_ref();
2074 let name = table_names::vertex_table_name(label);
2075
2076 let branched = self
2087 .fork_scope
2088 .as_ref()
2089 .is_some_and(|s| s.branch_for(&name).is_some());
2090 if branched {
2091 if !backend.table_exists(&name).await? {
2094 return Ok(Vec::new());
2095 }
2096 let mut filter_parts = vec![Self::build_active_filter(filter)];
2097 if ctx.is_some()
2098 && let Some(hwm) = self.version_high_water_mark()
2099 {
2100 filter_parts.push(FilterExpr::version_at_most(hwm));
2101 }
2102 let request = ScanRequest::all(&name)
2103 .with_filter(FilterExpr::all(filter_parts))
2104 .with_columns(vec!["_vid".to_string()]);
2105 let batches = backend.scan(request).await?;
2106 let mut results = Vec::new();
2107 for batch in batches {
2108 let vid_col = batch
2109 .column_by_name("_vid")
2110 .ok_or(anyhow!("Missing _vid"))?
2111 .as_any()
2112 .downcast_ref::<UInt64Array>()
2113 .ok_or(anyhow!("Invalid _vid"))?;
2114 for i in 0..batch.num_rows() {
2115 results.push((Vid::from(vid_col.value(i)), 0.0_f32));
2116 }
2117 }
2118 return Ok(results);
2119 }
2120
2121 let mut results = Vec::new();
2122 if backend.table_exists(&name).await? {
2124 let backend_metric = match &metric {
2125 DistanceMetric::L2 => BackendMetric::L2,
2126 DistanceMetric::Cosine => BackendMetric::Cosine,
2127 DistanceMetric::Dot => BackendMetric::Dot,
2128 _ => BackendMetric::Cosine,
2129 };
2130
2131 let mut filter_parts = vec![Self::build_active_filter(filter)];
2132 if ctx.is_some()
2133 && let Some(hwm) = self.version_high_water_mark()
2134 {
2135 filter_parts.push(FilterExpr::version_at_most(hwm));
2136 }
2137 let combined_filter = FilterExpr::all(filter_parts);
2138
2139 let batches = backend
2140 .multivector_search(
2141 &name,
2142 property,
2143 query,
2144 k,
2145 backend_metric,
2146 combined_filter,
2147 opts,
2148 )
2149 .await?;
2150 results = extract_vid_score_pairs(&batches, "_vid", "_distance")?;
2151 }
2152
2153 Ok(results)
2154 }
2155
2156 pub async fn sparse_search(
2169 &self,
2170 label: &str,
2171 property: &str,
2172 query: &[(u32, f32)],
2173 k: usize,
2174 ) -> Result<Vec<(Vid, f32)>> {
2175 #[cfg(feature = "lance-backend")]
2176 {
2177 let name = table_names::vertex_table_name(label);
2178 let branched = self
2179 .fork_scope
2180 .as_ref()
2181 .is_some_and(|s| s.branch_for(&name).is_some());
2182 if branched {
2183 let backend = self.backend.as_ref();
2205 if !backend.table_exists(&name).await? {
2208 return Ok(Vec::new());
2209 }
2210 let request = ScanRequest::all(&name)
2211 .with_filter(Self::build_active_filter(None))
2212 .with_columns(vec!["_vid".to_string()]);
2213 let batches = backend.scan(request).await?;
2214 let mut results = Vec::new();
2215 for batch in batches {
2216 let vid_col = batch
2217 .column_by_name("_vid")
2218 .ok_or_else(|| anyhow!("Missing _vid"))?
2219 .as_any()
2220 .downcast_ref::<UInt64Array>()
2221 .ok_or_else(|| anyhow!("Invalid _vid"))?;
2222 for i in 0..batch.num_rows() {
2223 results.push((Vid::from(vid_col.value(i)), 0.0_f32));
2224 }
2225 }
2226 return Ok(results);
2227 }
2228 match self
2229 .index_manager()
2230 .sparse_vector_index(label, property)
2231 .await
2232 {
2233 Ok(idx) => idx.query_topk(query, k).await,
2234 Err(_) => Ok(Vec::new()),
2235 }
2236 }
2237 #[cfg(not(feature = "lance-backend"))]
2238 {
2239 let _ = (label, property, query, k);
2240 Ok(Vec::new())
2241 }
2242 }
2243
2244 pub async fn fts_search(
2260 &self,
2261 label: &str,
2262 property: &str,
2263 query: &str,
2264 k: usize,
2265 filter: Option<&str>,
2266 ctx: Option<&QueryContext>,
2267 ) -> Result<Vec<(Vid, f32)>> {
2268 let backend = self.backend.as_ref();
2269 let name = table_names::vertex_table_name(label);
2270
2271 let mut results = if backend.table_exists(&name).await? {
2272 let mut filter_parts = vec![Self::build_active_filter(filter)];
2274 if ctx.is_some()
2275 && let Some(hwm) = self.version_high_water_mark()
2276 {
2277 filter_parts.push(FilterExpr::version_at_most(hwm));
2278 }
2279 let combined_filter = FilterExpr::all(filter_parts);
2280
2281 let batches = backend
2282 .full_text_search(&name, property, query, k, combined_filter)
2283 .await?;
2284
2285 let mut fts_results = extract_vid_score_pairs(&batches, "_vid", "_score")?;
2286 fts_results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
2288 fts_results
2289 } else {
2290 Vec::new()
2291 };
2292
2293 if let Some(qctx) = ctx {
2295 merge_l0_into_fts_results(&mut results, qctx, label, property, query, k);
2296 }
2297
2298 Ok(results)
2299 }
2300
2301 #[cfg(feature = "lance-backend")]
2302 pub async fn get_vertex_by_uid(&self, uid: &UniId, label: &str) -> Result<Option<Vid>> {
2303 let index = self.uid_index(label)?;
2304 index.get_vid(uid).await
2305 }
2306
2307 #[cfg(feature = "lance-backend")]
2308 pub async fn insert_vertex_with_uid(&self, label: &str, vid: Vid, uid: UniId) -> Result<()> {
2309 let index = self.uid_index(label)?;
2310 index.write_mapping(&[(uid, vid)]).await
2311 }
2312
2313 pub async fn load_subgraph(
2314 &self,
2315 start_vids: &[Vid],
2316 edge_types: &[u32],
2317 max_hops: usize,
2318 direction: GraphDirection,
2319 l0: Option<&L0Buffer>,
2320 ) -> Result<WorkingGraph> {
2321 let mut graph = WorkingGraph::new();
2322 let schema = self.schema_manager.schema();
2323
2324 let label_map: HashMap<u16, String> = schema
2326 .labels
2327 .values()
2328 .map(|meta| {
2329 (
2330 meta.id,
2331 schema.label_name_by_id(meta.id).unwrap().to_owned(),
2332 )
2333 })
2334 .collect();
2335
2336 let edge_type_map: HashMap<u32, String> = schema
2337 .edge_types
2338 .values()
2339 .map(|meta| {
2340 (
2341 meta.id,
2342 schema.edge_type_name_by_id(meta.id).unwrap().to_owned(),
2343 )
2344 })
2345 .collect();
2346
2347 let target_edge_types: HashSet<u32> = edge_types.iter().copied().collect();
2348
2349 let mut frontier: Vec<Vid> = start_vids.to_vec();
2351 let mut visited: HashSet<Vid> = HashSet::new();
2352
2353 for &vid in start_vids {
2355 graph.add_vertex(vid);
2356 }
2357
2358 for _hop in 0..max_hops {
2359 let mut next_frontier = HashSet::new();
2360
2361 for &vid in &frontier {
2362 if visited.contains(&vid) {
2363 continue;
2364 }
2365 visited.insert(vid);
2366 graph.add_vertex(vid);
2367
2368 for &etype_id in &target_edge_types {
2370 let etype_name = edge_type_map
2371 .get(&etype_id)
2372 .ok_or_else(|| anyhow!("Unknown edge type ID: {}", etype_id))?;
2373
2374 let (dir_str, neighbor_is_dst) = match direction {
2378 GraphDirection::Outgoing => ("fwd", true),
2379 GraphDirection::Incoming => ("bwd", false),
2380 };
2381
2382 let mut edges: HashMap<Eid, EdgeState> = HashMap::new();
2383
2384 let _edge_ver = self
2389 .pinned_snapshot
2390 .as_ref()
2391 .and_then(|s| s.edges.get(etype_name).map(|es| es.lance_version));
2392
2393 let backend = self.backend();
2395 for current_src_label in label_map.values() {
2396 let adj_ds =
2397 match self.adjacency_dataset(etype_name, current_src_label, dir_str) {
2398 Ok(ds) => ds,
2399 Err(_) => continue,
2400 };
2401 if let Some((neighbors, eids)) =
2402 adj_ds.read_adjacency_backend(backend, vid).await?
2403 {
2404 for (n, eid) in neighbors.into_iter().zip(eids) {
2405 edges.insert(
2406 eid,
2407 EdgeState {
2408 neighbor: n,
2409 version: 0,
2410 deleted: false,
2411 },
2412 );
2413 }
2414 break; }
2416 }
2417
2418 let delta_ds = self.delta_dataset(etype_name, dir_str)?;
2420 let delta_entries = delta_ds
2421 .read_deltas(backend, vid, &schema, self.snapshot_version_hwm())
2422 .await?;
2423 Self::apply_delta_to_edges(&mut edges, delta_entries, neighbor_is_dst);
2424
2425 if let Some(l0) = l0 {
2427 Self::apply_l0_to_edges(&mut edges, l0, vid, etype_id, direction);
2428 }
2429
2430 Self::add_edges_to_graph(
2432 &mut graph,
2433 edges,
2434 vid,
2435 etype_id,
2436 neighbor_is_dst,
2437 &visited,
2438 &mut next_frontier,
2439 );
2440 }
2441 }
2442 frontier = next_frontier.into_iter().collect();
2443
2444 if frontier.is_empty() {
2446 break;
2447 }
2448 }
2449
2450 Ok(graph)
2451 }
2452
2453 fn apply_delta_to_edges(
2455 edges: &mut HashMap<Eid, EdgeState>,
2456 delta_entries: Vec<crate::storage::delta::L1Entry>,
2457 neighbor_is_dst: bool,
2458 ) {
2459 for entry in delta_entries {
2460 let neighbor = if neighbor_is_dst {
2461 entry.dst_vid
2462 } else {
2463 entry.src_vid
2464 };
2465 let current_ver = edges.get(&entry.eid).map(|s| s.version).unwrap_or(0);
2466
2467 if entry.version > current_ver {
2468 edges.insert(
2469 entry.eid,
2470 EdgeState {
2471 neighbor,
2472 version: entry.version,
2473 deleted: matches!(entry.op, Op::Delete),
2474 },
2475 );
2476 }
2477 }
2478 }
2479
2480 fn apply_l0_to_edges(
2482 edges: &mut HashMap<Eid, EdgeState>,
2483 l0: &L0Buffer,
2484 vid: Vid,
2485 etype_id: u32,
2486 direction: GraphDirection,
2487 ) {
2488 let l0_neighbors = l0.get_neighbors(vid, etype_id, direction);
2489 for (neighbor, eid, ver) in l0_neighbors {
2490 let current_ver = edges.get(&eid).map(|s| s.version).unwrap_or(0);
2491 if ver > current_ver {
2492 edges.insert(
2493 eid,
2494 EdgeState {
2495 neighbor,
2496 version: ver,
2497 deleted: false,
2498 },
2499 );
2500 }
2501 }
2502
2503 for (eid, state) in edges.iter_mut() {
2505 if l0.is_tombstoned(*eid) {
2506 state.deleted = true;
2507 }
2508 }
2509 }
2510
2511 fn add_edges_to_graph(
2513 graph: &mut WorkingGraph,
2514 edges: HashMap<Eid, EdgeState>,
2515 vid: Vid,
2516 etype_id: u32,
2517 neighbor_is_dst: bool,
2518 visited: &HashSet<Vid>,
2519 next_frontier: &mut HashSet<Vid>,
2520 ) {
2521 for (eid, state) in edges {
2522 if state.deleted {
2523 continue;
2524 }
2525 graph.add_vertex(state.neighbor);
2526
2527 if !visited.contains(&state.neighbor) {
2528 next_frontier.insert(state.neighbor);
2529 }
2530
2531 if neighbor_is_dst {
2532 graph.add_edge(vid, state.neighbor, eid, etype_id);
2533 } else {
2534 graph.add_edge(state.neighbor, vid, eid, etype_id);
2535 }
2536 }
2537 }
2538}
2539
2540fn extract_vid_score_pairs(
2542 batches: &[arrow_array::RecordBatch],
2543 vid_column: &str,
2544 score_column: &str,
2545) -> Result<Vec<(Vid, f32)>> {
2546 let mut results = Vec::new();
2547 for batch in batches {
2548 let vid_col = batch
2549 .column_by_name(vid_column)
2550 .ok_or_else(|| anyhow!("Missing {} column", vid_column))?
2551 .as_any()
2552 .downcast_ref::<UInt64Array>()
2553 .ok_or_else(|| anyhow!("Invalid {} column type", vid_column))?;
2554
2555 let score_col = batch
2556 .column_by_name(score_column)
2557 .ok_or_else(|| anyhow!("Missing {} column", score_column))?
2558 .as_any()
2559 .downcast_ref::<Float32Array>()
2560 .ok_or_else(|| anyhow!("Invalid {} column type", score_column))?;
2561
2562 for i in 0..batch.num_rows() {
2563 results.push((Vid::from(vid_col.value(i)), score_col.value(i)));
2564 }
2565 }
2566 Ok(results)
2567}
2568
2569fn extract_vid_and_vector_pairs(
2578 batches: &[arrow_array::RecordBatch],
2579 vid_column: &str,
2580 vector_column: &str,
2581) -> Result<Vec<(Vid, Vec<f32>)>> {
2582 use arrow_array::{Array, FixedSizeListArray};
2583 let mut out = Vec::new();
2584 for batch in batches {
2585 let vid_col = batch
2586 .column_by_name(vid_column)
2587 .ok_or_else(|| anyhow!("Missing {} column", vid_column))?
2588 .as_any()
2589 .downcast_ref::<UInt64Array>()
2590 .ok_or_else(|| anyhow!("Invalid {} column type", vid_column))?;
2591
2592 let vec_col = batch
2593 .column_by_name(vector_column)
2594 .ok_or_else(|| anyhow!("Missing vector column {}", vector_column))?
2595 .as_any()
2596 .downcast_ref::<FixedSizeListArray>()
2597 .ok_or_else(|| anyhow!("Vector column {} is not FixedSizeList", vector_column))?;
2598
2599 for i in 0..batch.num_rows() {
2600 if vec_col.is_null(i) {
2601 continue;
2602 }
2603 let element = vec_col.value(i);
2604 let floats = element
2605 .as_any()
2606 .downcast_ref::<Float32Array>()
2607 .ok_or_else(|| {
2608 anyhow!("Vector column {} inner type is not Float32", vector_column)
2609 })?;
2610 let emb: Vec<f32> = (0..floats.len()).map(|j| floats.value(j)).collect();
2611 out.push((Vid::from(vid_col.value(i)), emb));
2612 }
2613 }
2614 Ok(out)
2615}
2616
2617fn extract_embedding_from_props(
2631 props: &uni_common::Properties,
2632 property: &str,
2633) -> Option<Vec<f32>> {
2634 Vec::<f32>::try_from(props.get(property)?).ok()
2635}
2636
2637fn l0_buffers_in_precedence(ctx: &QueryContext) -> Vec<Arc<parking_lot::RwLock<L0Buffer>>> {
2640 let mut buffers: Vec<Arc<parking_lot::RwLock<L0Buffer>>> =
2641 ctx.pending_flush_l0s.iter().map(Arc::clone).collect();
2642 buffers.push(Arc::clone(&ctx.l0));
2643 if let Some(ref txn) = ctx.transaction_l0 {
2644 buffers.push(Arc::clone(txn));
2645 }
2646 buffers
2647}
2648
2649fn merge_l0_into_vector_results(
2659 results: &mut Vec<(Vid, f32)>,
2660 ctx: &QueryContext,
2661 label: &str,
2662 property: &str,
2663 query: &[f32],
2664 k: usize,
2665 metric: &DistanceMetric,
2666) {
2667 let buffers = l0_buffers_in_precedence(ctx);
2668
2669 let mut l0_candidates: HashMap<Vid, f32> = HashMap::new();
2671 let mut tombstoned: HashSet<Vid> = HashSet::new();
2673
2674 for buf_arc in &buffers {
2675 let buf = buf_arc.read();
2676
2677 for &vid in &buf.vertex_tombstones {
2679 tombstoned.insert(vid);
2680 }
2681
2682 for (&vid, labels) in &buf.vertex_labels {
2684 if !labels.iter().any(|l| l == label) {
2685 continue;
2686 }
2687 if let Some(props) = buf.vertex_properties.get(&vid)
2688 && let Some(emb) = extract_embedding_from_props(props, property)
2689 {
2690 if emb.len() != query.len() {
2691 continue; }
2699 let dist = metric.compute_distance(&emb, query);
2700 l0_candidates.insert(vid, dist);
2702 tombstoned.remove(&vid);
2704 }
2705 }
2706 }
2707
2708 if l0_candidates.is_empty() && tombstoned.is_empty() {
2710 return;
2711 }
2712
2713 results.retain(|(vid, _)| !tombstoned.contains(vid));
2715
2716 for (vid, dist) in &l0_candidates {
2718 if tombstoned.contains(vid) {
2724 continue;
2725 }
2726 if let Some(existing) = results.iter_mut().find(|(v, _)| v == vid) {
2727 existing.1 = *dist;
2728 } else {
2729 results.push((*vid, *dist));
2730 }
2731 }
2732
2733 results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
2735 results.truncate(k);
2736}
2737
2738pub fn collect_l0_label_candidates(ctx: &QueryContext, label: &str) -> (Vec<Vid>, HashSet<Vid>) {
2752 let buffers = l0_buffers_in_precedence(ctx);
2753
2754 let mut live: HashSet<Vid> = HashSet::new();
2755 let mut tombstoned: HashSet<Vid> = HashSet::new();
2756
2757 for buf_arc in &buffers {
2758 let buf = buf_arc.read();
2759
2760 for &vid in &buf.vertex_tombstones {
2762 tombstoned.insert(vid);
2763 live.remove(&vid);
2764 }
2765
2766 for (&vid, labels) in &buf.vertex_labels {
2769 if !labels.iter().any(|l| l == label) {
2770 continue;
2771 }
2772 if buf.vertex_properties.contains_key(&vid) {
2773 live.insert(vid);
2774 tombstoned.remove(&vid);
2775 }
2776 }
2777 }
2778
2779 (live.into_iter().collect(), tombstoned)
2780}
2781
2782fn compute_text_relevance(query: &str, text: &str) -> f32 {
2787 let query_tokens: HashSet<String> =
2788 query.split_whitespace().map(|t| t.to_lowercase()).collect();
2789 if query_tokens.is_empty() {
2790 return 0.0;
2791 }
2792 let text_tokens: HashSet<String> = text.split_whitespace().map(|t| t.to_lowercase()).collect();
2793 let hits = query_tokens
2794 .iter()
2795 .filter(|t| text_tokens.contains(t.as_str()))
2796 .count();
2797 hits as f32 / query_tokens.len() as f32
2798}
2799
2800fn extract_text_from_props<'a>(
2802 props: &'a uni_common::Properties,
2803 property: &str,
2804) -> Option<&'a str> {
2805 props.get(property)?.as_str()
2806}
2807
2808fn merge_l0_into_fts_results(
2818 results: &mut Vec<(Vid, f32)>,
2819 ctx: &QueryContext,
2820 label: &str,
2821 property: &str,
2822 query: &str,
2823 k: usize,
2824) {
2825 let buffers = l0_buffers_in_precedence(ctx);
2826
2827 let mut l0_candidates: HashMap<Vid, f32> = HashMap::new();
2829 let mut tombstoned: HashSet<Vid> = HashSet::new();
2831
2832 for buf_arc in &buffers {
2833 let buf = buf_arc.read();
2834
2835 for &vid in &buf.vertex_tombstones {
2837 tombstoned.insert(vid);
2838 }
2839
2840 for (&vid, labels) in &buf.vertex_labels {
2842 if !labels.iter().any(|l| l == label) {
2843 continue;
2844 }
2845 if let Some(props) = buf.vertex_properties.get(&vid)
2846 && let Some(text) = extract_text_from_props(props, property)
2847 {
2848 let score = compute_text_relevance(query, text);
2849 if score > 0.0 {
2850 l0_candidates.insert(vid, score);
2852 }
2853 tombstoned.remove(&vid);
2855 }
2856 }
2857 }
2858
2859 if l0_candidates.is_empty() && tombstoned.is_empty() {
2861 return;
2862 }
2863
2864 results.retain(|(vid, _)| !tombstoned.contains(vid));
2866
2867 for (vid, score) in &l0_candidates {
2869 if tombstoned.contains(vid) {
2873 continue;
2874 }
2875 if let Some(existing) = results.iter_mut().find(|(v, _)| v == vid) {
2876 existing.1 = *score;
2877 } else {
2878 results.push((*vid, *score));
2879 }
2880 }
2881
2882 results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
2884 results.truncate(k);
2885}